from .models import CNNClassifier, save_model
from .utils import ConfusionMatrix, load_data, LABEL_NAMES
import torch
import torchvision
from torchvision import transforms
from torchvision.transforms import functional as F
import torch.utils.tensorboard as tb
import torch.nn as nn
import torch.optim as optim

TRAIN_PATH = "data/train"
VALID_PATH = "data/valid"
EARLY_STOP = 20

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

transformations = transforms.Compose([
    transforms.RandomRotation(degrees=10),
    transforms.ColorJitter(0.3, 0.4, 0.3, 0.3),
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),

    transforms.ToTensor(),
    #transforms.Normalize([0.5, 0.5, 0.5],[0.5,0.5,0.5])
    transforms.Normalize([0.425, 0.425, 0.425],[0.25,0.25,0.25])
])

def train(args):
    from os import path
    model = CNNClassifier().to(device)
    loss_function = nn.CrossEntropyLoss()
    optimizer = None
    if args.optimizer == "SGD":
        optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.9)
    else:
        optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
    data_loader = load_data(TRAIN_PATH, batch_size=args.batch_size, transformations=transformations)
    validation_accuracy = 0
    for epoch in range(0, args.epochs):
        model.train()
        running_loss = 0.0
        for index, (inputs, labels) in enumerate(data_loader):
            inputs = inputs.to(device)
            labels = labels.to(device)
            optimizer.zero_grad()

            outputs = model(inputs)
            loss = loss_function(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        model.eval()
        confusion = ConfusionMatrix(6)
        for img, label in load_data(VALID_PATH):
            confusion.add(model(img.to(device)).argmax(1).cpu(), label)
        print("global accuracy: ", confusion.global_accuracy, " loss at epoch ", epoch, running_loss)
        if validation_accuracy > confusion.global_accuracy and epoch > args.early_stop:
            exit()
        else:
            validation_accuracy = confusion.global_accuracy
        save_model(model)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()

    parser.add_argument('--log_dir')
    parser.add_argument('--epochs', type=int, default=10)
    parser.add_argument('--batch_size', type=int, default=200)
    parser.add_argument('--optimizer', type=str, default="adam")
    parser.add_argument('--learning_rate', type=float, default=0.001)
    parser.add_argument('--early_stop', type=int, default=5000)
    # Put custom arguments here

    args = parser.parse_args()
    train(args)
