import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torchvision import datasets, models
from torch.utils.data import DataLoader
# from torchsummary import summary

from util import make_transforms
import os
import time
import copy
import argparse
from tensorboardX import SummaryWriter

cwd = os.getcwd()

# data 폴더에 원하는 데이터를 형식에 맞추어 저장
DATA_PATH = os.path.join(cwd, 'data')
class_names = os.listdir(DATA_PATH)
base_dir = DATA_PATH

## Create the data loader
def load_dataset(format, batch_size):
    data_path = os.path.join(cwd, format)
    dataset = datasets.ImageFolder(
        root=data_path,
        transform=chosen_transforms[format]
    )
    data_loader = DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=4,
        shuffle=True
    )
    return data_loader, len(dataset), dataset.classes

## Start Training

def train_model(model, criterion, optimizer, scheduler, num_epochs=10):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                optimizer.step()
                scheduler.step()
                model.train()  # Set model to training mode
            else:
                model.eval()  # Set model to evaluate mode

            current_loss = 0.0
            current_corrects = 0

            # Here's where the training happens
            print('Iterating through data...')

            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # We need to zero the gradients, don't forget it
                optimizer.zero_grad()

                # Time to carry out the forward training poss
                # We only need to log the loss stats if we are in training phase
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # We want variables to hold the loss statistics
                current_loss += loss.item() * inputs.size(0)
                current_corrects += torch.sum(preds == labels.data)
                summary.add_scalar('loss/current_loss', current_loss, epoch)

            epoch_loss = current_loss / dataset_sizes[phase]
            epoch_acc = current_corrects.double() / dataset_sizes[phase]

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))

            # Make a copy of the model if the accuracy on the validation set has improved
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
        summary.add_scalar('loss/loss', epoch_loss, epoch)
        summary.add_scalar('acc/acc', epoch_acc, epoch)

        print()

    time_since = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_since // 60, time_since % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # Now we'll load in the best model weights and return it
    model.load_state_dict(best_model_wts)
    return model

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--use-cuda', action='store_true', default=True,
                        help='Use NVIDIA GPU acceleration')
    parser.add_argument('--epoch', type=int, default=6,
                        help='Number of epoch')
    parser.add_argument('--save', type=str, default='model',
                        help='Name of saved model')
    args = parser.parse_args()
    args.use_cuda = args.use_cuda and torch.cuda.is_available()
    if args.use_cuda:
        print("Using GPU for acceleration")
    else:
        print("Using CPU for computation")

    return args


if __name__ == '__main__':


    args = get_args()
    summary = SummaryWriter()
    chosen_transforms = make_transforms()


    ## Set code to run on device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    train_loader, train_size, class_names = load_dataset('data/train', 8)
    print("Train Data Set size is: ", train_size)
    print("Class Names are: ", class_names)
    inputs, classes = next(iter(train_loader))
#    print(inputs.shape, classes.shape)

    ## Set up Pre-Trained Model

    CLASS_NUM = len(class_names)
    EPOCHS = args.epoch
    SAVED_NAME = args.save + '.pt'
    ## Load pretrained model
    resnet50 = models.resnet50(pretrained=False)

    # Freeze model parameters
    # for param in resnet50.parameters():
    #     param.requires_grad = False

    ## Change the final layer of the resnet model
    # Change the final layer of ResNet50 Model for Transfer Learning
    fc_inputs = resnet50.fc.in_features

    # resnet50.fc = nn.Linear(fc_inputs, CLASS_NUM)
    resnet50.fc = nn.Sequential(
        nn.Linear(fc_inputs, 128),
        nn.ReLU(),
        nn.Dropout(0.4),
        nn.Linear(128, CLASS_NUM)
    )


    # Convert model to be used on GPU
    resnet50 = resnet50.to(device)

    # summary(resnet50, (3, 227, 227))

    # Define Optimizer and Loss Function
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(resnet50.parameters())
    # optimizer = optim.SGD(resnet50.fc.parameters(), lr=0.001, momentum=0.9)
    # Decay LR by a factor of 0.1 every 3 epochs
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

    dataloaders = {}
    dataset_sizes = {}
    idx_to_class = {}
    batch_size = 4
    dataloaders['train'], dataset_sizes['train'], class_names = load_dataset('data/train', batch_size)
    dataloaders['val'], dataset_sizes['val'], _ = load_dataset('data/val', batch_size)
    for idx, label in enumerate(class_names):
        idx_to_class[idx] = label

    base_model = train_model(resnet50, criterion, optimizer, exp_lr_scheduler, num_epochs=EPOCHS)
    saved_model = os.path.join(cwd, SAVED_NAME)
    torch.save(base_model, saved_model)
    # torch.save(base_model.state_dict(), cwd + '/model_state_dict.pt')
    print("Save Complete to ", saved_model)