import csv
import copy
import time
from tqdm import tqdm
import torch
import numpy as np
import os
from datetime import datetime
import pathlib
import matplotlib.pyplot as plt


def load_checkpoint(bpath):

    checkpoint_folder = os.path.join(bpath, 'checkpoint')
    checkpoint_filename = os.path.join(
        checkpoint_folder, 'checkpoint.pth.tar')

    bestweights_filename = os.path.join(
        checkpoint_folder, 'best_weights_checkpoint.pth.tar')

    file = pathlib.Path(checkpoint_filename)

    if not file.exists():
        return None, None, None, None, None, None

    file = pathlib.Path(bestweights_filename)

    best_weight = None
    if file.exists():
        best_weight = torch.load(bestweights_filename)
        best_weight = best_weight['state_dict']

    checkpoint = torch.load(checkpoint_filename)

    return checkpoint['epoch'], checkpoint['state_dict'], best_weight, checkpoint['optimizer'], checkpoint['best_loss'], checkpoint['best_pred']


def save_checkpoint(bpath, state, is_best=False):

    checkpoint_folder = os.path.join(bpath, 'checkpoint')

    if is_best:
        best_pred = state['best_pred']
        with open(os.path.join(checkpoint_folder, 'best_pred.txt'), 'w') as f:
            f.write(str(best_pred))

        best_pred = state['best_loss']
        with open(os.path.join(checkpoint_folder, 'best_loss.txt'), 'w') as f:
            f.write(str(best_pred))

        torch.save(state, os.path.join(checkpoint_folder,
                                       'best_weights_checkpoint.pth.tar'))

    torch.save(state, os.path.join(checkpoint_folder,
                                   'checkpoint.pth.tar'))


def train_model(model, criterion, dataloaders, optimizer, scheduler, metrics, bpath, num_epochs=3):

    start_epoch, state_dict, bweights, optm, bloss, bpred = load_checkpoint(
        bpath)

    if start_epoch is not None:
        print("")
        print("NEW CHECKPOINT FOUND! LAST EPOCH ", start_epoch)
        print("")
        model.load_state_dict(state_dict)
        start_epoch += 1

        best_model_wts = copy.deepcopy(bweights)
        best_loss = float(bloss)

        best_Train_dice = 1e-5
        best_Valid_dice = bpred
    else:
        start_epoch = 1
        best_model_wts = copy.deepcopy(model.state_dict())
        best_loss = 1e10

        best_Train_dice = 1e-5
        best_Valid_dice = 1e-5

    since = time.time()

    # Use gpu if available
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)
    # Initialize the log file for training and testing loss and metrics
    fieldnames = ['epoch', 'Train_loss', 'Valid_loss'] + \
        [f'Train_{m}' for m in metrics.keys()] + \
        [f'Valid_{m}' for m in metrics.keys()]
    with open(os.path.join(bpath, 'log.csv'), 'w', newline='') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()

    for epoch in range(start_epoch, num_epochs+1):
        print('Epoch {}/{}'.format(epoch, num_epochs))
        print('-' * 10)
        # Each epoch has a training and validation phase
        # Initialize batch summary
        batchsummary = {a: [0] for a in fieldnames}

        for phase in ['Train', 'Valid']:
            if phase == 'Train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            # Iterate over data.

            for sample in tqdm(iter(dataloaders[phase])):

                inputs = sample['image'].to(device)
                masks = sample['mask'].to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # track history if only in train
                with torch.set_grad_enabled(phase == 'Train'):
                    outputs = model(inputs)
                    # loss = criterion(outputs['out'], masks)
                    loss = criterion(outputs, masks)

                    # y_pred = outputs['out'].data.cpu().numpy().squeeze(1)
                    y_pred = outputs.data.cpu().numpy().squeeze(1)
                    y_true = masks.data.cpu().numpy().squeeze(1)

                    for name, metric in metrics.items():
                        if name == 'dice' or name == 'dice_target':
                            # Use a classification threshold of 0.5
                            val_metric = metric(y_pred > 0.5, y_true > 0)

                            if val_metric is not None:
                                batchsummary[f'{phase}_{name}'].append(
                                    val_metric)

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

            batchsummary['epoch'] = epoch
            epoch_loss = loss
            batchsummary[f'{phase}_loss'] = epoch_loss.item()
            print('{} Loss: {:.4f}'.format(phase, loss))

        print('New LR: ', scheduler.get_last_lr())
        scheduler.step()

        for field in fieldnames[3:]:
            batchsummary[field] = np.mean(batchsummary[field])

        print(batchsummary)

        epoch_valid_dice = np.mean(batchsummary['Valid_dice_tumor'])
        is_best = False
        with open(os.path.join(bpath, 'log.csv'), 'a', newline='') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            writer.writerow(batchsummary)

            SAVE_BESTLOSS_WEIGTH = False
            if SAVE_BESTLOSS_WEIGTH:
                # deep copy the model
                if phase == 'Valid' and loss < best_loss:
                    print('\nnew best loss: {:.4f} in epoch {}\n'.format(
                        loss, epoch))
                    best_loss = loss
                    best_model_wts = copy.deepcopy(model.state_dict())
                    now = datetime.now()
                    str_datetime = now.strftime("%Y%m%d_%H_%M_%S")

                    best_Train_dice = np.mean(batchsummary['Train_dice'])
                    best_Valid_dice = np.mean(batchsummary['Valid_dice'])

                    torch.save(model, os.path.join(
                        bpath, 'weights_partial_epch{}_{}.pt'.format(epoch, str_datetime)))
            else:
                # deep copy the model
                if phase == 'Valid' and epoch_valid_dice > best_Valid_dice:
                    is_best = True
                    print('\nNew valid dice: {:.4f} in epoch {}\n'.format(
                        epoch_valid_dice, epoch))
                    best_loss = loss.item()
                    best_model_wts = copy.deepcopy(model.state_dict())
                    now = datetime.now()
                    str_datetime = now.strftime("%Y%m%d_%H_%M_%S")

                    best_Train_dice = np.mean(batchsummary['Train_dice'])
                    best_Valid_dice = epoch_valid_dice

                    torch.save(model, os.path.join(
                        bpath, 'weights_partial_diceval_epch{}_{}.pt'.format(epoch, str_datetime)))

                    # torch.save(model, os.path.join(
                    #     bpath, 'model_weights_partial.pt'))

            save_checkpoint(bpath, {
                'epoch': epoch,
                'state_dict': model.state_dict(),
                'optimizer': optimizer.state_dict(),
                'best_pred': best_Valid_dice,
                'best_loss': best_loss
            }, is_best=is_best)

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Lowest by valid dice Loss: {:4f}'.format(best_loss))
    print('Max valid Dice: {:4f}'.format(best_Valid_dice))

    # load best model weights
    model.load_state_dict(best_model_wts)

    return best_Train_dice, best_Valid_dice