import argparse
from pathlib import Path
import utils
import torch
import numpy as np
import torch.backends.cudnn as cudnn
import pandas as pd
import time
import datetime
import os

# mixup
from timm.data import Mixup


# log about
# import mlflow

# for the challenge 2021
import ecg_dataset_2021
import engine_ecg_2021

# for the challenge 2020
import ecg_dataset_2020
import engine_ecg_2020

from engine_ecg_2021 import train_one_epoch, evaluate

# for the timm
from timm.models import create_model
import timm
from timm.scheduler.cosine_lr import CosineLRScheduler

import sys

from optimizer import NoamOpt

import models_mamba_ecg

def get_args_parser():
    parser = argparse.ArgumentParser('DeiT training and evaluation script', add_help=False)
    parser.add_argument('--batch-size', default=64, type=int)
    parser.add_argument('--epochs', default=300, type=int)

    # Model parameters
    parser.add_argument('--model', default='deit_base_patch16_224', type=str, metavar='MODEL', help='Name of model to train')

    # Parameters regarding the Drop
    parser.add_argument('--drop', type=float, default=0.0, metavar='PCT', help='Dropout rate (default: 0.)')
    parser.add_argument('--drop-path', type=float, default=0.1, metavar='PCT', help='Drop path rate (default: 0.1)')

    # setting the mode
    parser.add_argument('--train-mode', action='store_true')
    parser.add_argument('--no-train-mode', action='store_false', dest='train_mode')
    parser.set_defaults(train_mode=True)
    
    
    # * data augmentation method
    parser.add_argument('--mixup', type=float, default=0,
                        help='mixup alpha, mixup enabled if > 0. (default: 0)')
    parser.add_argument('--cutmix', type=float, default=0,
                        help='cutmix alpha, cutmix enabled if > 0. (default: 0)')
    parser.add_argument('--mixup_no_label', type=float, default=0,
                        help='this is the mixup but without interpolation of label. (default: 0)')
    parser.add_argument('--progressive_switch', type=bool, default=False,
                        help='this is the switch of progressive data augmentation method')
    
    # the output directory
    parser.add_argument('--output_dir', default='', help='path where to save, empty for no saving')

    # the device information
    parser.add_argument('--device', default='cuda', help='device to use for training / testing')

    parser.add_argument('--seed', default=0, type=int)
    parser.add_argument('--start_epoch', default=0, type=int, metavar='N', help='start epoch')


    parser.add_argument('--num_workers', default=10, type=int)
    
    parser.add_argument('--pin-mem', action='store_true', help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
    parser.add_argument('--no-pin-mem', action='store_false', dest='pin_mem', help='')
    parser.set_defaults(pin_mem=True)
    
    # special parameters for this strategy
    parser.add_argument('--depth', default=24, type=int)
    parser.add_argument('--block', default='VisionMamba', choices=['OriginalMamba', 'VisionMamba'], type=str, help='the selection of the block of the paradigm')
    parser.add_argument('--lead', default='12Lead', choices=['12Lead', 'RandomLead'], type=str, help='the number of Lead')
    parser.add_argument('--lrschedule', default='CosineAnnealing', choices=['Noam', 'CosineAnnealing'], type=str, help='the selection of learning rate strategy')

    # if use random token position
    parser.add_argument('--if_random_cls_token_position', action='store_true')
    parser.add_argument('--no_randargs.if_nan2numom_cls_token_position', action='store_false', dest='if_random_cls_token_position')
    parser.set_defaults(if_random_cls_token_position=False)    

    # if use random token rank
    parser.add_argument('--if_random_token_rank', action='store_true')
    parser.add_argument('--no_random_token_rank', action='store_false', dest='if_random_token_rank')
    parser.set_defaults(if_random_token_rank=False)

    # using for the journal 
    parser.add_argument('--fused_add_norm', type=bool, default=True, help='combines the element-wise addition (from residual connections) and normalization (e.g., RMSNorm)')
    parser.add_argument('--if_divide_out', type=bool, default=True, help='Should the result be divided by two when combining outputs from two directions?')
    parser.add_argument('--use_middle_cls_token', type=bool, default=True, help='Whether the class is inserted in the middle of sequence')

    # the switch of scenario
    parser.add_argument('--challenge_scenario', default='2021', choices=[2021, 2020], type=int, help='the scenario of Challenge')


    return parser


def collate(batch):
    # Left-zero padding
    ch = batch[0][0].shape[0]
    maxL = 8192
    X = np.zeros((len(batch), ch, maxL))
    
    for i in range(len(batch)):
        X[i, :, -batch[i][0].shape[-1]:] = batch[i][0]
    
    t = np.array([b[1] for b in batch])
    
    X = torch.from_numpy(X)
    t = torch.from_numpy(t)

    return X, t

def main(args, data_directory, model_directory, group_number):
    
    train_log_fp = open(args.output_dir + '/train_log_group_%d.txt' % group_number, 'a')
    print(args)
    train_log_fp.write("The is the configuration: {}\n".format(args))

    
    device = torch.device(args.device)
    
    print("This is the running device", device)
    train_log_fp.write("This is the running device:{}\n".format(device))

    seed = args.seed + utils.get_rank()
    torch.manual_seed(seed)
    np.random.seed(seed)

    print("This is the random's seed:", seed)
    train_log_fp.write("This is the random's seed:{}\n".format(seed))
    

    cudnn.benchmark = True

    
    dataset_file_all_address = "../../collection_of_all_datasets/"
    # below is building dataset(train and validation)

    print("This is building the training set:")
    ############ training area #########################
    the_training_address = data_directory + "/training_group" + data_directory[-1] + ".csv"
    df = pd.read_csv(the_training_address)
    print("Total the {} files will be fed into model for training".format(len(df['Name'])))
    train_log_fp.write("Total the {} files will be fed into model for training \n".format(len(df['Name'])))

    print("This is the first file of training:",df['Name'][0])


    training_header_files=[]
    
    for i in range(len(df['Name'])):
        each_header_file = dataset_file_all_address + df['Name'][i]
        training_header_files.append(each_header_file)
    
    collate_training = collate
    
    
    if args.mixup > 0:
        print("The mixup is using, and the percentage is:", args.mixup)
        train_log_fp.write("The mixup is using, and the percentage is:{}\n".format(args.mixup))
        # collate_training = collate_mixup
    
    if args.cutmix > 0:
        print("The cutmix is using, and the percentage is:", args.cutmix)
        train_log_fp.write("The cutmix is using, and the percentage is:{}\n".format(args.cutmix))
    if args.mixup_no_label > 0:
        print("The mixup is using without label's interpolation , and the percentage is:", args.mixup_no_label)
        train_log_fp.write("The mixup is using without label's interpolation , and the percentage is:{}\n".format(args.mixup_no_label))
    if args.progressive_switch:
        print("The progressive mixup is using.")
        train_log_fp.write("The progressive mixup is using.")

    if args.challenge_scenario == 2021:
        train_dataset = ecg_dataset_2021.dataset(training_header_files, Mixup = args.mixup, amount = len(df['Name']), cutMix=args.cutmix, Mixup_no_label_interpolate=args.mixup_no_label, progressive_switch=args.progressive_switch)
    elif args.challenge_scenario == 2020:
        train_dataset = ecg_dataset_2020.dataset(training_header_files, Mixup = args.mixup, amount = len(df['Name']), cutMix=args.cutmix, Mixup_no_label_interpolate=args.mixup_no_label, progressive_switch=args.progressive_switch)
    else:
        raise ValueError(f"No matching condition for value: {args.challenge_scenario}")


    if args.lead == "12Lead":
        lead_number = 12
        print("This 12 lead is using.")
        train_log_fp.write("This 12 lead is using.")
    elif args.lead == "RandomLead":
        lead_number = None
        print("This random lead is using.")
        train_log_fp.write("This random lead is using.")
    else:
        raise ValueError(f"No matching condition for value: {args.lead}")
    
    """
    we filter out the sample which the length is 8192 via the below code.
    just like the random shift windows.  
    """
    train_dataset.num_leads = lead_number
    train_dataset.sample = True
    ###################################################
    print("done")
    

    print("This is building the validation set:")
    ############ testing area #########################
    the_testing_address = data_directory + "/testing_group" + data_directory[-1]+".csv"
    df = pd.read_csv(the_testing_address)
    print("Total the {} files will be used as testing".format(len(df['Name'])))
    train_log_fp.write("Total the {} files will be used as testing \n".format(len(df['Name'])))

    print("This is the first file of testing:",df['Name'][0])

    testing_header_files=[]

    for i in range(len(df['Name'])):
        each_header_file = dataset_file_all_address + df['Name'][i]
        testing_header_files.append(each_header_file)
    

    if args.challenge_scenario == 2021:
        test_dataset = ecg_dataset_2021.dataset(testing_header_files, Mixup = 0)
    elif args.challenge_scenario == 2020:
        test_dataset = ecg_dataset_2020.dataset(testing_header_files, Mixup = 0)
    else:
        raise ValueError(f"No matching condition for value: {args.challenge_scenario}")


    test_dataset.num_leads = 12
    test_dataset.sample = True
    ###################################################
    print("done")

    sampler_train = torch.utils.data.RandomSampler(train_dataset)
    sampler_val = torch.utils.data.SequentialSampler(test_dataset)

    data_loader_train = torch.utils.data.DataLoader(
        train_dataset, sampler=sampler_train,
        batch_size=args.batch_size,
        collate_fn=collate_training,
        num_workers=args.num_workers,
        pin_memory=args.pin_mem,
        drop_last=True,
    )
    # Setting `drop_last=True` means that this incomplete batch will be dropped, 
    # ensuring that all batches fed to the model during training have **the same size.**
    
    data_loader_val = torch.utils.data.DataLoader(
        test_dataset, sampler=sampler_val,
        batch_size=int(1.5 * args.batch_size),
        collate_fn=collate,
        num_workers=args.num_workers,
        pin_memory=args.pin_mem,
        drop_last=False
    )

    if args.challenge_scenario == 2021:
        args.nb_classes = 26
    elif args.challenge_scenario == 2020:
        args.nb_classes = 27
    else:
        raise ValueError(f"No matching condition for value: {args.challenge_scenario}")


    print(f"Creating model: {args.model}")
    train_log_fp.write("Creating model:{}\n".format(args.model))
    model = create_model(
        args.model,
        pretrained=False,
        num_classes=args.nb_classes,
        drop_rate=args.drop,
        drop_path_rate=args.drop_path,
        drop_block_rate=None,
        block = args.block,
        depth = args.depth,
        fused_add_norm = args.fused_add_norm,
        use_middle_cls_token = args.use_middle_cls_token,
        img_size=8192
    )


    model.to(device)


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

    print('number of params:', n_parameters)
    train_log_fp.write("number of params:{}\n".format(n_parameters))
    
    optimizer = torch.optim.Adam(model.parameters(), lr=0.0006, betas=(0.9, 0.98), eps=1e-9)

    if args.lrschedule == "Noam":
        optimizer = NoamOpt(729, 1, 4000, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))
        print("The Noam is using as learning rate strategy.")
        train_log_fp.write("The Noam is using as learning rate strategy.")
        
    elif args.lrschedule == "CosineAnnealing":
        lr_scheduler = CosineLRScheduler(
        optimizer,
        t_initial=13,  # Number of epochs after warmup
        lr_min=1e-6,
        warmup_lr_init = 1e-5,
        warmup_t=5,
        cycle_limit=1,
        t_in_epochs=True,
        warmup_prefix=True  # Added to reach initial_lr
        )
        print("The cosing annealing is using as learning rate strategy.")
        train_log_fp.write("The cosing annealing is using as learning rate strategy.")
    else:
        raise ValueError(f"No matching condition for value: {args.lrschedule}")

    # below is for the cosine annealing schedule
    
    # the loss is set as normal BCE loss.
    criterion = torch.nn.BCEWithLogitsLoss()


    output_dir = Path(args.output_dir)

    print(f"Start training for {args.epochs} epochs")
    train_log_fp.write(f"Start training for {args.epochs} epochs\n")
    
    start_time = time.time()

    max_AUPRC = 0.0


    train_auprc_list = []
    testing_AUPRC_list = []
    loss_value_after_each_epcoh_list = []
    loss_value_testing_list = []

    testing_auroc_list = []
    testing_f1_list = []

    testing_subset_accuracy_list = []
    hamming_loss_list=[]
    challenge_score_list = []


    jump_count = 0

    thrs_list = []
    # loss_value_after_each_epcoh_testing_list = []
    # epoch is started from 0
    for epoch in range(args.start_epoch, args.epochs):
        
        train_log_fp.write("\n")
        train_log_fp.write("---------------------------------------------------------- \n")
        train_log_fp.write("---------------------------------------------------------- \n")
        train_log_fp.write("---------------------------------------------------------- \n")
        
        now = datetime.datetime.now()
        current_time = now.strftime("%H:%M:%S")
        print("Current Time =", current_time)
        train_log_fp.write(f"Current Time = {current_time} \n")
        train_log_fp.write(f"Current epoch: {epoch} \n")
        
        train_dataset.set_epoch(epoch)  # Update the current epoch in the dataset
        
        if args.lrschedule == "CosineAnnealing":
            lr_scheduler.step(epoch)


        ###### below is for the training ###########
        if args.challenge_scenario == 2021:
            train_auprc, loss_value_after_each_epcoh, thrs, scores_F1, scores_SubsetAccuracy, scores_HammingLoss = engine_ecg_2021.train_one_epoch(
                model, criterion, data_loader_train,
                optimizer, device, epoch,
                set_training_mode=args.train_mode,  # keep in eval mode for deit finetuning / train mode for training and deit III finetuning
                args=args,
            )
        elif args.challenge_scenario == 2020: 
            train_auprc, loss_value_after_each_epcoh, thrs, scores_F1, scores_SubsetAccuracy, scores_HammingLoss = engine_ecg_2020.train_one_epoch(
                model, criterion, data_loader_train,
                optimizer, device, epoch,
                set_training_mode=args.train_mode,  # keep in eval mode for deit finetuning / train mode for training and deit III finetuning
                args=args,
            )
        else:
            raise ValueError(f"No matching condition for value: {args.challenge_scenario}")
        

        train_auprc_list.append(round(train_auprc, 4))
        loss_value_after_each_epcoh_list.append(round(loss_value_after_each_epcoh, 4))
        print("**************************************************************")
        print("This is the list of Training of AUPRC:", train_auprc_list)
        train_log_fp.write(f"This is the list of Training of AUPRC: {train_auprc_list} \n")
        print("This is the list of Training of loss:", loss_value_after_each_epcoh_list)
        train_log_fp.write(f"This is the list of Training of loss: {loss_value_after_each_epcoh_list} \n")
        print("**************************************************************")


        ###### below is for the evaluation ###########

        if args.challenge_scenario == 2021:
            AUPRC, auroc, f1, hamming, subset_accuracy, challenge_score, loss_value_after_each_epoch_testing = engine_ecg_2021.evaluate(data_loader_val, model, thrs, scores_F1, scores_SubsetAccuracy, scores_HammingLoss, device)
        elif args.challenge_scenario == 2020: 
            AUPRC, auroc, f1, hamming, subset_accuracy, challenge_score, loss_value_after_each_epoch_testing = engine_ecg_2020.evaluate(data_loader_val, model, thrs, scores_F1, scores_SubsetAccuracy, scores_HammingLoss, device)
        else:
            raise ValueError(f"No matching condition for value: {args.challenge_scenario}")

        # print(f"Accuracy of the network on the {len(test_dataset)} test images: {AUPRC:.4f}")

        testing_AUPRC_list.append(round(AUPRC, 4))
        loss_value_testing_list.append(round(loss_value_after_each_epoch_testing, 4))
        testing_auroc_list.append(round(auroc, 4))
        testing_f1_list.append(round(f1, 4))
        testing_subset_accuracy_list.append((round(subset_accuracy, 4)))  
        hamming_loss_list.append((round(hamming, 4)))
        challenge_score_list.append((round(challenge_score, 4)))


        # loss_value_after_each_epcoh_testing_list.append(loss)
        print("**********************************************************")
        print("This is the list of Testing AUPRC:", testing_AUPRC_list)
        print("This is the list of Testing loss:", loss_value_testing_list)
        print("This is the list of Testing auroc:", testing_auroc_list)
        print("This is the list of Testing f1:", testing_f1_list)
        print("This is the list of subset accuracy:", testing_subset_accuracy_list)
        print("This is the list of hamming_loss:", hamming_loss_list)
        print("This is the list of challenge_score:", challenge_score_list)

        train_log_fp.write("---------------------------------------------------------- \n")
        train_log_fp.write(f"This is the list of Testing AUPRC: {testing_AUPRC_list} \n")
        train_log_fp.write(f"This is the list of Testing loss: {loss_value_testing_list} \n")
        train_log_fp.write(f"This is the list of Testing auroc: {testing_auroc_list} \n")
        train_log_fp.write(f"This is the list of Testing f1: {testing_f1_list} \n")
        train_log_fp.write(f"This is the list of Testing subset accuracy: {testing_subset_accuracy_list} \n")
        train_log_fp.write(f"This is the list of Testing hamming_loss: {hamming_loss_list} \n")
        train_log_fp.write(f"This is the list of challenge_score: {challenge_score_list} \n")
        print("**********************************************************")
        
        # here I have to put code regarding the auprc
        if max_AUPRC < AUPRC:
            max_AUPRC = AUPRC
            jump_count = 0
            if args.output_dir:
                checkpoint_paths = [output_dir / 'best_auprc_checkpoint.pth']
                for checkpoint_path in checkpoint_paths:
                    utils.save_on_master({
                        'model': model.state_dict(),
                        'optimizer': optimizer.optimizer.state_dict(),
                        'epoch': epoch,
                        'args': args,
                    }, checkpoint_path)
        else:
            jump_count = jump_count + 1
        

        print(f'This is the Max AUPRC: {max_AUPRC:.4f}')
        train_log_fp.write(f"This is the Max AUPRC: {max_AUPRC:.4f} \n")

        if jump_count > 4:
            print("This experimental will be finished at epoch:", epoch)
            train_log_fp.write(f"This experimental will be finished at epoch: {epoch} \n")
            break

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))
     
    train_log_fp.write('Training time {}\n'.format(total_time_str))
    train_log_fp.close()
    old_txt_file_name = args.output_dir + '/train_log_group_%d.txt' % group_number
    new_txt_file_name = args.output_dir + '/train_log_group_%d_MAX_AUPRC_%.4f_.txt' % (group_number, max_AUPRC)
    os.rename(old_txt_file_name, new_txt_file_name)


if __name__ == '__main__':
    parser = argparse.ArgumentParser('DeiT training and evaluation script, but this is for the multiple classification ECG.', parents=[get_args_parser()])
    args = parser.parse_args()
    
    if args.challenge_scenario == 2021: 
        data_directory = "./csv-file_2021_challenge/training_validation_testing/group"
    elif args.challenge_scenario == 2020:
        data_directory = "./csv-file_2020_challenge/training_validation_testing/group"
    else:
        raise ValueError(f"No matching condition for value: {args.challenge_scenario}")

    print("This is the scenario:", args.challenge_scenario)
    print("This is the scenario:", args.challenge_scenario)

    model_directory = "./model/model_group"

    for i in range(1,6):
        data_directory_x = data_directory + str(i)
        model_directory_x = model_directory +str(i)

        args.output_dir = f"./output/Scenario_{args.challenge_scenario}_{args.block}_depth_{args.depth}_{args.lead}_batchSize_{args.batch_size}_CutMix(random)_{args.cutmix}_MixUp_{args.mixup}_JTEHM_group{i}"

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

        print("This is the group", i)
        main(args, data_directory_x, model_directory_x, i)