criterion = nn.CrossEntropyLoss()
writer = SummaryWriter('runs/caer')

train_losses, val_losses = [], []
train_accs, val_accs = [], []

def train(args, dev, model, optimizer, dataloaders):
    train_loader, val_loader, test_loader = dataloaders
    for epoch in range(args.epochs):
        print(epoch)
        count = 0
        train_loss, val_loss = 0,0
        train_acc, val_acc = 0,0

        for (images,labels) in train_loader:
            #if (count == 6):
             #   break
            images = images.to(dev)
            labels = labels.to(dev)
            out_pred = model(images)
            labels = labels.type(torch.LongTensor).to(dev)

            loss = criterion(out_pred, labels)

            print(loss.item())

            train_loss += loss.item()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            count += 1

            total = labels.size(0)
            _, predicted = torch.max(out_pred.data, 1)
            correct = (predicted == labels).sum().item()
            train_acc += float(correct) / float(total)

        else:
            with torch.no_grad():
                model.eval()
                for (images,labels) in val_loader:
                    #labels = to_one_hot(labels)

                    images = images.to(dev)
                    labels = labels.to(dev)

                    val_out_pred = model(images)
                    labels = labels.type(torch.LongTensor).to(dev)

                    loss = criterion(val_out_pred, labels)
                    print(loss.item())
                    val_loss += loss.item()

                    total = labels.size(0)
                    _, predicted = torch.max(val_out_pred.data, 1)
                    correct = (predicted == labels).sum().item()
                    val_acc += float(correct) / float(total)

            model.train()

        print("Epoch: {}/{}.. ".format(epoch, args.epochs),
              "Training Loss: {:.3f}.. ".format(train_loss/len(train_loader)),
              "Validation Loss: {:.3f}.. ".format(val_loss/len(val_loader)),
              "Training Accuracy: {: .3f}.. ".format(train_acc/len(train_loader)*100),
              "Validation Accuracy: {:.3f}".format(val_acc/len(val_loader)*100))

        train_losses.append(train_loss/len(train_loader))
        val_losses.append(val_loss/len(val_loader))
        train_accs.append(train_acc/len(train_loader)*100)
        val_accs.append(val_acc/len(val_loader)*100)
        writer.add_scalar('Train/Loss', train_loss/len(train_loader), epoch)
        writer.add_scalar('Val/Loss', val_loss/len(val_loader), epoch)
        writer.flush()

    checkpoint = {'model': model,
              'state_dict': model.state_dict(),
              'optimizer' : optimizer.state_dict()}
    torch.save(checkpoint, 'caer_model.pth')
    evaluate(model, test_loader, args, dev)

def evaluate(model, test_loader, args, dev):
    for i, data in enumerate(test_loader):
        images, labels = data
        images, labels = images.to(dev), labels.to(dev)

        with torch.no_grad():
            out_pred = model(images)
            _, predicted = torch.max(out_pred.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print('Test Accuracy of the model on test images: {} %'.format(100 * correct / total))
    return