import os
import time
import importlib
import collections
import pickle as cp
import glob
import numpy as np
import pandas as pd

from natsort import natsorted
from PIL import Image
import torch
import torchvision
import torch.nn.functional as F
import torchvision.models as models
import torchvision.transforms as transforms
from torch.utils.data import Subset
from torch.utils.data import Dataset, DataLoader

from sklearn.model_selection import StratifiedShuffleSplit
from sklearn.model_selection import train_test_split
from sklearn.model_selection import KFold

from networks import basenet, grayResNet2

DATASET_PATH = '/content/drive/My Drive/CD2 Project/'

TRAIN_DATASET_PATH = '/content/drive/My Drive/CD2 Project/data/classification data/aug&HGG+NL_train/'
TRAIN_TARGET_PATH = '/content/drive/My Drive/CD2 Project/data/classification data/train_augNL_classify_target.csv'
VAL_DATASET_PATH = '/content/drive/My Drive/CD2 Project/data/classification data/NL+HGG_val/'
VAL_TARGET_PATH = '/content/drive/My Drive/CD2 Project/data/classification data/val_nonaugNL_classify_target.csv'
TEST_DATASET_PATH = '/content/drive/My Drive/CD2 Project/data/classification data/NL+HGG_test/'
TEST_TARGET_PATH = '/content/drive/My Drive/CD2 Project/data/classification data/test_nonaugNL_classify_target.csv'

current_epoch = 0


def split_dataset(args, dataset, k):
    # load dataset
    X = list(range(len(dataset)))
    Y = dataset.targets

    # split to k-fold
    assert len(X) == len(Y)

    def _it_to_list(_it):
        return list(zip(*list(_it)))

    sss = StratifiedShuffleSplit(n_splits=k, random_state=args.seed, test_size=0.1)
    Dm_indexes, Da_indexes = _it_to_list(sss.split(X, Y))

    return Dm_indexes, Da_indexes




def concat_image_features(image, features, max_features=3):
    _, h, w = image.shape
    #print("\nfsize: ", features.size()) # (1, 240, 240)
   # features.size(0) = 64
    #print(features.size(0))
    #max_features = min(features.size(0), max_features)

    max_features = features.size(0)
    image_feature = image.clone()

    for i in range(max_features):
        # features torch.Size([64, 16, 16])

        feature = features[i:i+1]
        #torch.Size([1, 16, 16])

        _min, _max = torch.min(feature), torch.max(feature)
        feature = (feature - _min) / (_max - _min + 1e-6)
        # torch.Size([1, 16, 16])

        feature = torch.cat([feature]*1, 0)
        #feature = torch.cat([feature]*3, 0)
        # torch.Size([3, 16, 16]) -> [1, 16, 16]

        feature = feature.view(1, 1, feature.size(1), feature.size(2))
        #feature = feature.view(1, 3, feature.size(1), feature.size(2))
        # torch.Size([1, 3, 16, 16])-> [1, 1, 16, 16]

        feature = F.upsample(feature, size=(h,w), mode="bilinear")
        # torch.Size([1, 3, 32, 32])-> [1, 1, 32, 32]

        feature = feature.view(1, h, w) #(3, h, w) input of size 3072
        # torch.Size([3, 32, 32])->[1, 32, 32]

        #print("img_feature & feature size:\n", image_feature.size(),"\n", feature.size()) 
        # img_feature & feature size:
        # torch.Size([1, 32, 32]) -> [1, 32, 64]
        # torch.Size([3, 32, 32] ->[1, 32, 32]
        

        image_feature = torch.cat((image_feature, feature), 2) ### dim = 2
        #print("\nimg feature size: ", image_feature.size()) #[1, 240, 720]

    return image_feature

def get_model_name(args):
    from datetime import datetime, timedelta, timezone
    now = datetime.now(timezone.utc)
    tz = timezone(timedelta(hours=9))
    now = now.astimezone(tz)
    date_time = now.strftime("%B_%d_%H:%M:%S")
    model_name = '__'.join([date_time, args.network, str(args.seed)])
    return model_name



def dict_to_namedtuple(d):
    Args = collections.namedtuple('Args', sorted(d.keys()))

    for k,v in d.items():
        if type(v) is dict:
            d[k] = dict_to_namedtuple(v)

        elif type(v) is str:
            try:
                d[k] = eval(v)
            except:
                d[k] = v

    args = Args(**d)
    return args


def parse_args(kwargs):
    # combine with default args
    kwargs['dataset'] =  kwargs['dataset'] if 'dataset' in kwargs else 'BraTS'
    kwargs['network'] =  kwargs['network'] if 'network' in kwargs else 'resnet50'
    kwargs['optimizer'] = kwargs['optimizer'] if 'optimizer' in kwargs else 'adam'
    kwargs['learning_rate'] = kwargs['learning_rate'] if 'learning_rate' in kwargs else 0.01
    kwargs['seed'] =  kwargs['seed'] if 'seed' in kwargs else None
    kwargs['use_cuda'] =  kwargs['use_cuda'] if 'use_cuda' in kwargs else True
    kwargs['use_cuda'] =  kwargs['use_cuda'] and torch.cuda.is_available()
    kwargs['num_workers'] = kwargs['num_workers'] if 'num_workers' in kwargs else 4
    kwargs['print_step'] = kwargs['print_step'] if 'print_step' in kwargs else 50
    kwargs['val_step'] = kwargs['val_step'] if 'val_step' in kwargs else 50
    kwargs['scheduler'] = kwargs['scheduler'] if 'scheduler' in kwargs else 'exp'
    kwargs['batch_size'] = kwargs['batch_size'] if 'batch_size' in kwargs else 16
    kwargs['start_step'] = kwargs['start_step'] if 'start_step' in kwargs else 0
    kwargs['max_step'] = kwargs['max_step'] if 'max_step' in kwargs else 500
    kwargs['augment_path'] = kwargs['augment_path'] if 'augment_path' in kwargs else None

    # to named tuple
    args = dict_to_namedtuple(kwargs)
    return args, kwargs


def select_model(args):
    # grayResNet2
    resnet_dict = {'resnet18':grayResNet2.resnet18(), 'resnet34':grayResNet2.resnet34(), 
    'resnet50':grayResNet2.resnet50(), 'resnet101':grayResNet2.resnet101(), 'resnet152':grayResNet2.resnet152()}

    if args.network in resnet_dict:
        model = resnet_dict[args.network]
    # else: # 3 channels
    # backbone = models.__dict__[args.network]()
    # model = basenet.BaseNet(backbone, args)

    #print(model) # print model architecture
    return model


def select_optimizer(args, model):
    if args.optimizer == 'sgd':
        optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=0.0001)
    elif args.optimizer == 'rms':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=args.learning_rate)
    elif args.optimizer == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
    else:
        raise Exception('Unknown Optimizer')
    return optimizer


def select_scheduler(args, optimizer):
    if not args.scheduler or args.scheduler == 'None':
        return None
    elif args.scheduler =='clr':
        return torch.optim.lr_scheduler.CyclicLR(optimizer, 0.01, 0.015, mode='triangular2', step_size_up=250000, cycle_momentum=False)
    elif args.scheduler =='exp':
        return torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9999283, last_epoch=-1)
    else:
        raise Exception('Unknown Scheduler')

def get_train_transform(args, model, transform, log_dir=None):
    if args.dataset == 'cifar10':
        transform = transforms.Compose([
            transforms.Pad(4),
            transforms.RandomCrop(32),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor()
        ])

    else:
        transform = transforms.Compose([
            transforms.Resize([240, 240]),
            transforms.ToTensor() 
    ])

    return transform

def get_valid_transform(args, model):
    if args.dataset == 'cifar10':
        val_transform = transforms.Compose([
            transforms.Resize(32),
            transforms.ToTensor()
        ])

    else:
        val_transform = transforms.Compose([
            transforms.Resize([240, 240]),
            transforms.ToTensor() 
    ])

    return val_transform

class CustomDataset(Dataset): 
  def __init__(self, data_path, csv_path, transform):
    self.path = data_path
    self.imgs = natsorted(os.listdir(data_path))
    self.len = len(self.imgs)
    self.transform = transform

    df = pd.read_csv(csv_path)
    targets_list = []

    for fname in self.imgs:
        row = df.loc[df['filename'] == fname]
        targets_list.append(row.iloc[0, 1])

    self.targets = targets_list

  def __len__(self):
      return self.len

  def __getitem__(self, idx):
    img_loc = os.path.join(self.path, self.imgs[idx])
    targets = self.targets[idx]
    image = Image.open(img_loc)
    image = self.transform(image)
    #print("\n idx, img, targets: ", idx, img_loc, targets)
    return image, targets



def get_dataset(args, transform, split='train'):
    assert split in ['train', 'val', 'test']

    if args.dataset == 'cifar10':
        train = split in ['train', 'val', 'trainval']
        dataset = torchvision.datasets.CIFAR10(DATASET_PATH,
                                               train=train,
                                               transform=transform,
                                               download=True)

        if split in ['train', 'val']:
            split_path = os.path.join(DATASET_PATH,
                    'cifar-10-batches-py', 'train_val_index.cp')

            if not os.path.exists(split_path):
                [train_index], [val_index] = split_dataset(args, dataset, k=1)
                split_index = {'train':train_index, 'val':val_index}
                cp.dump(split_index, open(split_path, 'wb'))

            split_index = cp.load(open(split_path, 'rb'))
            dataset = Subset(dataset, split_index[split])

    else:                                            
        if split in ['train']:
            dataset = CustomDataset(TRAIN_DATASET_PATH, TRAIN_TARGET_PATH, transform)
        elif split in ['val']: 
            dataset = CustomDataset(VAL_DATASET_PATH, VAL_TARGET_PATH, transform)
        else : #  test
            dataset = CustomDataset(TEST_DATASET_PATH, TEST_TARGET_PATH, transform)


    return dataset


def get_dataloader(args, dataset, shuffle=False, pin_memory=True):
    data_loader = torch.utils.data.DataLoader(dataset,
                                              batch_size=args.batch_size,
                                              shuffle=shuffle,
                                              num_workers=args.num_workers,
                                              pin_memory=pin_memory)
    return data_loader



def get_inf_dataloader(args, dataset):
    global current_epoch
    data_loader = iter(get_dataloader(args, dataset, shuffle=True))

    while True:
        try:
            batch = next(data_loader)

        except StopIteration:
            current_epoch += 1
            data_loader = iter(get_dataloader(args, dataset, shuffle=True))
            batch = next(data_loader)

        yield batch




def train_step(args, model, optimizer, scheduler, criterion, batch, step, device=None):
    model.train()
    images, target = batch

    if device:
        images = images.to(device)
        target = target.to(device)

    elif args.use_cuda:
        images = images.cuda(non_blocking=True)
        target = target.cuda(non_blocking=True)

    # compute output
    start_t = time.time()
    output= model(images)
    forward_t = time.time() - start_t
    loss = criterion(output, target)

    # measure accuracy and record loss
    acc1 = accuracy(output, target, topk=(1, ))[0]
    acc1 /= images.size(0)


    # compute gradient and do SGD step
    optimizer.zero_grad()
    start_t = time.time()
    loss.backward()
    backward_t = time.time() - start_t
    optimizer.step()
    if scheduler: scheduler.step()

    # if writer and step % args.print_step == 0:
    #     n_imgs = min(images.size(0), 10)
    #     tag = 'train/' + str(step)
    #     for j in range(n_imgs):
    #         writer.add_image(tag,
    #                 concat_image_features(images[j], first[j]), global_step=step)

    return acc1, loss, forward_t, backward_t


#_acc1= accuracy(output, target, topk=(1,))
def accuracy(output, target, topk=(1,)):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))
        
        # print("\noutout: ", output.size()) #(32, 1000)
        # print("\npred: ", pred.size()) #(5, 32)
        # print("\ncorrect: ", correct.size()) #(5, 32)

        res = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
            res.append(correct_k)
        return res

def validate(args, model, criterion, valid_loader, step, device=None):
    # switch to evaluate mode
    model.eval()

    acc1 = 0
    samples = 0
    infer_t = 0

    with torch.no_grad():
        for i, (images, target) in enumerate(valid_loader):

            start_t = time.time()
            if device:
                images = images.to(device)
                target = target.to(device)

            elif args.use_cuda is not None:
                images = images.cuda(non_blocking=True)
                target = target.cuda(non_blocking=True)

            # compute output
            output  = model(images)
            loss = criterion(output, target)
            infer_t += time.time() - start_t

            # measure accuracy and record loss
            _acc1 = accuracy(output, target, topk=(1, ))[0]
            acc1 += _acc1
            samples += images.size(0)

    #print("\nsamples: ", samples) 4640
    acc1 /= samples

    # if writer:
    #     n_imgs = min(images.size(0), 10)
    #     for j in range(n_imgs):
    #         writer.add_image('valid/input_image',
    #                 concat_image_features(images[j], first[j]), global_step=step)

    return acc1, loss, infer_t