import argparse

import datetime

import json

import random

from re import A

import time

from pathlib import Path

import os

from torch.utils.data.dataset import Subset

from datasets.MyData import MyDataset

import numpy as np

import torch

from torch.utils.data import DataLoader, DistributedSampler,random_split

from sklearn.model_selection import KFold

import util.misc as utils

from models.engine import evaluate, train_one_epoch_SAM,test

from models import build_model

from models.sam import SAM

from torch.utils.tensorboard import SummaryWriter

import pandas as pd

def get_args_parser():

    parser = argparse.ArgumentParser('Set transformer detector', add_help=False)

    parser.add_argument('--lr', default=1e-4, type=float)

    parser.add_argument('--batch_size', default=8, type=int)

    parser.add_argument('--weight_decay', default=1e-4, type=float)

    parser.add_argument('--epochs', default=300, type=int)

    parser.add_argument('--lr_drop', default=200, type=int)

    parser.add_argument('--clip_max_norm', default=0.1, type=float,

                        help='gradient clipping max norm')

    parser.add_argument('--position_embedding', default='3Dlearned', type=str,

                        help="Type of positional embedding to use on top of the image features")

    # * Transformer

    parser.add_argument('--enc_layers', default=6, type=int,

                        help="Number of encoding layers in the transformer")

    parser.add_argument('--dec_layers', default=6, type=int,

                        help="Number of decoding layers in the transformer")

    parser.add_argument('--dim_feedforward', default=2048, type=int,

                        help="Intermediate size of the feedforward layers in the transformer blocks")

    parser.add_argument('--hidden_dim', default=256, type=int,

                        help="Size of the embeddings (dimension of the transformer)")

    parser.add_argument('--dropout', default=0.1, type=float,

                        help="Dropout applied in the transformer")

    parser.add_argument('--nheads', default=8, type=int,

                        help="Number of attention heads inside the transformer's attentions")

    parser.add_argument('--num_queries', default=100, type=int,

                        help="Number of query slots")

    parser.add_argument('--pre_norm', action='store_true',default=False)

    parser.add_argument('--pretrained_path', default='', type=str, help="path of pretrained model")

    # dataset parameters

    parser.add_argument('--dataset_file', default='coco')

    parser.add_argument('--output_dir', default='./',

                        help='path where to save, empty for no saving')

    parser.add_argument('--device', default='cuda',

                        help='device to use for training / testing')

    parser.add_argument('--seed', default=42, type=int)

    parser.add_argument('--resume', default='', help='resume from checkpoint')

    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',

                        help='start epoch')

    parser.add_argument('--eval', action='store_true')

    parser.add_argument('--num_workers', default=2, type=int)

    parser.add_argument('--kfoldNum', default=5, type=int)

    parser.add_argument('--dataDir',type=str,help='path of the data')

    parser.add_argument('--externalDataDir',type=str,help='path of the external test data')

    # distributed training parameters

    parser.add_argument('--world_size', default=1, type=int,

                        help='number of distributed processes')

    parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training')

    # cluster parameters

    parser.add_argument('--group_Q', action='store_true', default=False)

    parser.add_argument('--group_K', action='store_true', default=False)

    parser.add_argument('--cuda-devices', default=None)

    parser.add_argument('--max_num_cluster', default=64, type=int)

    parser.add_argument('--sequence_len', default=15000, type=int)

    # gridsearch parameters

    parser.add_argument('--withPosEmbedding', action='store_true', default=False)

    parser.add_argument('--seq_pool', action='store_true', default=False,help='use attention pooling layer for aggregating patch risk score') 

    parser.add_argument('--withLN', action='store_true', default=False)

    parser.add_argument('--withEmbeddingPreNorm', action='store_true', default=False,help='Pre-normalize the patch representation before feeding them into ViT')

    parser.add_argument('--input_pool', action='store_true', default=False)

    parser.add_argument('--mixUp', action='store_true', default=False)

    parser.add_argument('--SAM', action='store_true', default=False)

    return parser

def main(args):

    allDataset = MyDataset(root_dir=args.dataDir,sequence_len=args.sequence_len,max_num_cluster=args.max_num_cluster,status = 'test',input_pool=args.input_pool)

    kfoldSplits=KFold(n_splits=args.kfoldNum,shuffle=True,random_state=args.seed)  

    splitIdx = kfoldSplits.split(np.arange(len(allDataset)))

    CIndexTest = []

    CIndexExternalTest = []

    IBSTest = []

    IBSExternalTest = []

    correlationCoeffTest = []

    IPCWCIndexTest = []

    IPCWCIndexExternalTest = []

    for fold, (train_idx,test_idx) in enumerate(splitIdx):

        utils.init_distributed_mode(args)

        print(args)

        device = torch.device(args.device)

        # fix the seed for reproducibility

        seed = args.seed + utils.get_rank()

        torch.manual_seed(seed)

        np.random.seed(seed)

        random.seed(seed)

        model, criterion = build_model(args)

        model.to(device)

        model_without_ddp = model

        if args.distributed:

            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])

            model_without_ddp = model.module

        n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)

        print('number of params:', n_parameters)

        if args.pretrained_path!='':

            pretrainedModel = torch.load(args.pretrained_path)

            model_without_ddp.load_state_dict(pretrainedModel['model'], strict=False)

        if args.SAM:

            base_optimizer = torch.optim.Adam # define an optimizer for the "sharpness-aware" update

            optimizer_SAM = SAM(model_without_ddp.parameters(), base_optimizer,lr=args.lr,weight_decay=args.weight_decay)

            optimizer = optimizer_SAM

            lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop) 

        else:

            optimizer = torch.optim.AdamW(model_without_ddp.parameters(), lr=args.lr,

                                          weight_decay=args.weight_decay)

            lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)

        best_loss = 1e12

        tb_writer = SummaryWriter()

        dataset_train,dataset_val = random_split(Subset(allDataset,train_idx),[int(len(train_idx)*0.8),len(train_idx)-int(len(train_idx)*0.8)],generator=torch.Generator().manual_seed(args.seed))

        dataset_test = Subset(allDataset,test_idx)

        dataset_train_all = Subset(allDataset,train_idx)

        if args.distributed:

            sampler_train = DistributedSampler(dataset_train)

            sampler_val = DistributedSampler(dataset_val, shuffle=False)

            sampler_test = DistributedSampler(dataset_test, shuffle=False)

        else:

            sampler_train = torch.utils.data.RandomSampler(dataset_train)

            sampler_val = torch.utils.data.SequentialSampler(dataset_val)

            sampler_test = torch.utils.data.SequentialSampler(dataset_test)

        batch_sampler_train = torch.utils.data.BatchSampler(

            sampler_train, args.batch_size, drop_last=False)

        data_loader_train = DataLoader(dataset_train, batch_sampler=batch_sampler_train,

                                    collate_fn=None, num_workers=args.num_workers)

        data_loader_val = DataLoader(dataset_val, args.batch_size, sampler=sampler_val,

                                    drop_last=False,num_workers=args.num_workers)

        data_loader_test = DataLoader(dataset_test, args.batch_size, sampler=sampler_test,

                                    drop_last=False, num_workers=args.num_workers)

        output_dir = Path(args.output_dir)

        if args.kfoldNum>1:

            output_dir = output_dir/f'fold{fold}'

            output_dir.mkdir(parents=True, exist_ok=True)

        if args.resume:

            checkpoint = torch.load(args.resume, map_location='cpu')

            model_without_ddp.load_state_dict(checkpoint['model'])

            if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:

                optimizer.load_state_dict(checkpoint['optimizer'])

                lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])

                args.start_epoch = checkpoint['epoch'] + 1      

        if args.eval:

            val_stats = evaluate(model, criterion, data_loader_train, device, args.output_dir,'Validation')

        print(f"fold {fold}/{args.kfoldNum} Start training")

        start_time = time.time()

        # train the model 

        for epoch in range(args.start_epoch, args.epochs):

            if args.distributed:

                sampler_train.set_epoch(epoch)

            train_stats = train_one_epoch_SAM(

                model, criterion, data_loader_train, optimizer, device, epoch,fold,tb_writer,

                args.clip_max_norm,mixUp=args.mixUp,SAM=args.SAM)    

            lr_scheduler.step()

            train_eval_stats = evaluate(

                model, criterion, data_loader_train, device, args.output_dir,'Train')

            val_stats = evaluate(

                model, criterion, data_loader_val, device, args.output_dir, 'Validation')

            test_stats = evaluate(

                model, criterion, data_loader_test, device, args.output_dir, 'Test')

            is_best = val_stats['loss'] < best_loss

            best_loss = min(val_stats['loss'], best_loss)            

            if args.output_dir:

                checkpoint_paths = [output_dir / 'checkpoint.pth']

                # extra checkpoint before LR drop and every 100 epochs

                if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 100 == 0:

                    checkpoint_paths.append(output_dir / f'checkpoint{epoch:04}.pth')

                if is_best:

                    checkpoint_paths.append(output_dir / f'model_best.pth.tar')

                for checkpoint_path in checkpoint_paths:

                    utils.save_on_master({

                        'model': model_without_ddp.state_dict(),

                        'optimizer': optimizer.state_dict(),

                        'lr_scheduler': lr_scheduler.state_dict(),

                        'epoch': epoch,

                        'args': args,

                    }, checkpoint_path)

            log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},

                        **{f'train_eval{k}': v for k, v in train_eval_stats.items()},

                        **{f'val_{k}': v for k, v in val_stats.items()},

                        **{f'test_{k}': v for k, v in test_stats.items()},

                        'epoch': epoch,

                        'n_parameters': n_parameters}

            temploss = train_eval_stats

            tb_writer.add_scalar('train/loss'+'fold{}'.format(fold), temploss['loss'], epoch)

            tb_writer.add_scalar('train/CIndex'+'fold{}'.format(fold), temploss['CIndex'], epoch)                        

            temploss = val_stats

            tb_writer.add_scalar('val/loss'+'fold{}'.format(fold), temploss['loss'], epoch)

            tb_writer.add_scalar('val/CIndex'+'fold{}'.format(fold), temploss['CIndex'], epoch) 

            temploss = test_stats

            tb_writer.add_scalar('test/loss'+'fold{}'.format(fold), temploss['loss'], epoch)

            tb_writer.add_scalar('test/CIndex'+'fold{}'.format(fold), temploss['CIndex'], epoch)

            if args.output_dir and utils.is_main_process():

                with (output_dir / f"trainingLog_fold{fold}.txt").open("a") as f:

                    f.write(json.dumps(log_stats) + "\n")  

        # evaluate the best model on internal and external datasets

        bestModelPath =  output_dir / f'model_best.pth.tar' 

        bestCheckpoint = torch.load(bestModelPath)

        bestModel, testcriterion = build_model(args)

        bestModel.to(device)

        bestModel.load_state_dict(bestCheckpoint['model'])

        data_loader_train_all = DataLoader(dataset_train_all, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True)

        dataset_external_test = MyDataset(root_dir=externalDataDir,sequence_len=args.sequence_len,max_num_cluster=args.max_num_cluster,status='externalTest',input_pool=args.input_pool)

        dataset_external_test.status = 'externalTest'

        data_loader_external_test = DataLoader(dataset_external_test, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True)

        externalOutputDir = output_dir / 'externalTest'

        Path(externalOutputDir).mkdir(parents=True, exist_ok=True)

        internalTestBestModelStatus = test(bestModel, testcriterion, data_loader_test, data_loader_train_all, device, output_dir,fold,coxBiomarkerRisk)

        externalTestBestModelStatus = test(bestModel, testcriterion, data_loader_external_test, data_loader_train_all, device, externalOutputDir,fold)

        log_stats = {

            **{f'testBestModel_{k}': v for k, v in internalTestBestModelStatus.items()},

            **{f'ExternalTestBestModel_{k}': v for k, v in externalTestBestModelStatus.items()},

            'fold': fold

            }

        CIndexTest.append(internalTestBestModelStatus['CIndex'])

        CIndexExternalTest.append(externalTestBestModelStatus['CIndex'])

        IPCWCIndexTest.append(internalTestBestModelStatus['IPCWCIndex'])

        IPCWCIndexExternalTest.append(externalTestBestModelStatus['IPCWCIndex'])

        IBSTest.append(internalTestBestModelStatus['IBSTest'])

        IBSExternalTest.append(externalTestBestModelStatus['IBSTest'])

        if args.output_dir and utils.is_main_process():

            with (output_dir / f"testingLog_fold{fold}.txt").open("a") as f:

                f.write(json.dumps(log_stats) + "\n")  

        total_time = time.time() - start_time

        total_time_str = str(datetime.timedelta(seconds=int(total_time)))

        print('Training time {}'.format(total_time_str))

    output_dir = Path(args.output_dir)

    AverageBestModelStatus = {'AverageCIndexTest':np.mean(CIndexTest), 'AverageCIndexExternalTest':np.mean(CIndexExternalTest),\

        'AverageIPCWCIndexTest':np.mean(IPCWCIndexTest), 'AverageIPCWCIndexExternalTest':np.mean(IPCWCIndexExternalTest),\

        'AverageIBSTest':np.mean(IBSTest),'AverageIBSExternalTest':np.mean(IBSExternalTest),\

        'stdCIndexTest':np.std(CIndexTest), 'stdCIndexExternalTest':np.std(CIndexExternalTest),\

        'stdIPCWCIndexTest':np.std(IPCWCIndexTest), 'stdIPCWCIndexExternalTest':np.std(IPCWCIndexExternalTest),\

        'stdIBSTest':np.std(IBSTest),'stdIBSExternalTest':np.std(IBSExternalTest),'AverageCorrelationCoeffTest':np.mean(correlationCoeffTest)}

    log_stats = {

        **{f'{k}': v for k, v in AverageBestModelStatus.items()}

        }    

    if args.output_dir and utils.is_main_process():

        with (output_dir / "logAverage.txt").open("a") as f:

            f.write(json.dumps(log_stats) + "\n")      

    print(AverageBestModelStatus)

if __name__ == '__main__':

    now = datetime.datetime.now()

    dt_string = now.strftime("%d%m%Y_%H%M%S")    

    parser = argparse.ArgumentParser('DETR training and evaluation script', parents=[get_args_parser()])

    args = parser.parse_args()

    if args.cuda_devices is not None:

        os.environ["CUDA_VISIBLE_DEVICES"]=args.cuda_devices

    if args.output_dir:

        Path(args.output_dir).mkdir(parents=True, exist_ok=True)

    hyperparameters = vars(args)

    hyperparameter_stas = {

        **{f'{k}': v for k, v in hyperparameters.items()}

        }

    if args.output_dir and utils.is_main_process():

        with (Path(args.output_dir) / "logHyperparameters.txt").open("a") as f:

            f.write(json.dumps(hyperparameter_stas) + "\n")      

    main(args)