import argparse

import torch

torch.cuda.empty_cache() # clearing the occupied cuda memory

from torch.backends import cudnn

import torch.optim as optim

from torch.utils.data import DataLoader

import os

import numpy as np

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:256"

from dataset import LoadDataset

from model import InferenceNet, ECGnet

from loss import calculate_inference_loss, calculate_reconstruction_loss, calculate_ECG_reconstruction_loss, calculate_classify_loss

from utils import lossplot, lossplot_detailed, visualize_PC_with_label, ECG_visual_two, lossplot_classify, visualize_PC_with_twolabel

def train_ecg(args):

    DEVICE = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

    # DEVICE = torch.device('cpu')

    train_dataset = LoadDataset(path=args.partial_root, num_input=args.num_input, split='train')

    val_dataset = LoadDataset(path=args.partial_root, num_input=args.num_input, split='val')

    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True)

    val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True)

    cudnn.benchmark = True

    network = InferenceNet(in_ch=args.in_ch, out_ch=args.out_ch, num_input=args.num_input, z_dims=args.z_dims)

    if args.model is not None:

        print('Loaded trained model from {}.'.format(args.model))

        network.load_state_dict(torch.load(args.model))

    else:

        print('Begin training new model.')

    network.to(DEVICE)

    optimizer = optim.Adam(network.parameters(), lr=args.base_lr, weight_decay=args.weight_decay)

    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_decay_steps, gamma=args.lr_decay_rate)

    max_iter = int(len(train_dataset) / args.batch_size + 0.5)

    minimum_loss = 1e4

    best_epoch = 0

    lossfile_train = args.log_dir + "/training_loss.txt"

    lossfile_val = args.log_dir + "/val_loss.txt"

    lossfile_geometry_train = args.log_dir + "/training_calculate_inference_loss.txt"

    lossfile_geometry_val = args.log_dir + "/val_calculate_inference_loss.txt"

    lossfile_KL_train = args.log_dir + "/training_KL_loss.txt"

    lossfile_KL_val = args.log_dir + "/val_KL_loss.txt"

    lossfile_ecg_train = args.log_dir + "/training_ecg_loss.txt"

    lossfile_ecg_val = args.log_dir + "/val_ecg_loss.txt"

    for epoch in range(1, args.epochs + 1):

        if ((epoch % 25) == 0) and (epoch != 0):  

            lossplot_classify(lossfile_train, lossfile_val, lossfile_geometry_train, lossfile_geometry_val, lossfile_KL_train, lossfile_KL_val, lossfile_ecg_train, lossfile_ecg_val)

        f_train = open(lossfile_train, 'a')  # a: additional writing; w: overwrite writing

        f_val = open(lossfile_val, 'a')

        f_MI_train = open(lossfile_geometry_train, 'a')  # a: additional writing; w: overwrite writing

        f_MI_val = open(lossfile_geometry_val, 'a')

        f_KL_train = open(lossfile_KL_train, 'a')  # a: additional writing; w: overwrite writing

        f_KL_val = open(lossfile_KL_val, 'a')

        f_ecg_train = open(lossfile_ecg_train, 'a')  # a: additional writing; w: overwrite writing

        f_ecg_val = open(lossfile_ecg_val, 'a')

        # if ((epoch % 25) == 0) and (epoch != 0): 

        #     if  lamda_KL < 1:

        #         lamda_KL = 0.1*epoch*lamda_KL # 0.25

        #     else:

        #         lamda_KL = 0.1

        # training

        network.train()

        total_loss, iter_count = 0, 0

        for i, data in enumerate(train_dataloader, 1):

            partial_input, ECG_input, gt_MI, partial_input_coarse = data

            partial_input, ECG_input, gt_MI = partial_input.to(DEVICE), ECG_input.to(DEVICE), gt_MI.to(DEVICE)      

            partial_input_coarse = partial_input_coarse.to(DEVICE)      

            partial_input = partial_input.permute(0, 2, 1)

            optimizer.zero_grad()

            y_MI, y_ECG, mu, log_var = network(partial_input, ECG_input)

            loss_seg, KL_loss = calculate_classify_loss(y_MI, gt_MI, mu, log_var)

            loss_signal = calculate_ECG_reconstruction_loss(y_ECG, ECG_input)

            loss = loss_seg + args.lamda_KL*KL_loss

            check_grad = False

            if check_grad:

                print(loss_seg)

                print(loss_signal)

                print(KL_loss)

                print(loss.requires_grad)

                print(loss_seg.requires_grad)

                print(KL_loss.requires_grad)

                print(loss_signal.requires_grad)

            visual_check = False

            if visual_check:

                gd_ECG = ECG_input[0].cpu().detach().numpy()

                y_ECG = y_ECG[0].cpu().detach().numpy()

                ECG_visual_two(y_ECG, gd_ECG)

            loss.backward()

            optimizer.step()

            f_train.write(str(loss.item()))

            f_train.write('\n')

            f_MI_train.write(str(loss_seg.item()))

            f_MI_train.write('\n')

            f_KL_train.write(str(KL_loss.item()))

            f_KL_train.write('\n')

            f_ecg_train.write(str(loss_signal.item()))

            f_ecg_train.write('\n')

            iter_count += 1

            total_loss += loss.item()

            if i % 50 == 0:

                print("Training epoch {}/{}, iteration {}/{}: loss is {}".format(epoch, args.epochs, i, max_iter, loss.item()))

        scheduler.step()

        print("\033[96mTraining epoch {}/{}: avg loss = {}\033[0m".format(epoch, args.epochs, total_loss / iter_count))

        # evaluation

        network.eval()

        with torch.no_grad():

            total_loss, iter_count = 0, 0

            for i, data in enumerate(val_dataloader, 1):

                partial_input, ECG_input, gt_MI, partial_input_coarse = data

                partial_input, ECG_input, gt_MI = partial_input.to(DEVICE), ECG_input.to(DEVICE), gt_MI.to(DEVICE)  

                partial_input_coarse = partial_input_coarse.to(DEVICE)  

                partial_input = partial_input.permute(0, 2, 1)

                y_MI, y_ECG, mu, log_var = network(partial_input, ECG_input)

                loss_seg, KL_loss = calculate_classify_loss(y_MI, gt_MI, mu, log_var)

                loss_signal = calculate_ECG_reconstruction_loss(y_ECG, ECG_input)

                loss = loss_seg + args.lamda_KL*KL_loss

                total_loss += loss.item()

                iter_count += 1

                visual_check = False

                if visual_check:

                    gd_ECG = ECG_input[0].cpu().detach().numpy()

                    y_ECG = y_ECG[0].cpu().detach().numpy()

                    ECG_visual_two(y_ECG, gd_ECG)

                f_val.write(str(loss.item()))

                f_val.write('\n')

                f_MI_val.write(str(loss_seg.item()))

                f_MI_val.write('\n')

                f_KL_val.write(str(KL_loss.item()))

                f_KL_val.write('\n')

                f_ecg_val.write(str(loss_signal.item()))

                f_ecg_val.write('\n')

            mean_loss = total_loss / iter_count

            print("\033[35mValidation epoch {}/{}, loss is {}\033[0m".format(epoch, args.epochs, mean_loss))

            # records the best model and epoch

            if mean_loss < minimum_loss:

                best_epoch = epoch

                minimum_loss = mean_loss           

                strNetSaveName = 'net_model_classify.pkl'

                # strNetSaveName = 'net_with_%d.pkl' % epoch

                torch.save(network.state_dict(), args.log_dir + '/' + strNetSaveName)

        print("\033[4;37mBest model (lowest loss) in epoch {}\033[0m".format(best_epoch))

    lossplot(lossfile_train, lossfile_val)

def train(args):

    DEVICE = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

    # DEVICE = torch.device('cpu')

    train_dataset = LoadDataset(path=args.partial_root, num_input=args.num_input, split='train')

    val_dataset = LoadDataset(path=args.partial_root, num_input=args.num_input, split='val')

    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True)

    val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True)

    cudnn.benchmark = True

    network = InferenceNet(in_ch=args.in_ch, out_ch=args.out_ch, num_input=args.num_input, z_dims=args.z_dims)

    if args.model is not None:

        print('Loaded trained model from {}.'.format(args.model))

        network.load_state_dict(torch.load(args.model))

    else:

        print('Begin training new model.')

    network.to(DEVICE)

    optimizer = optim.Adam(network.parameters(), lr=args.base_lr, weight_decay=args.weight_decay)

    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_decay_steps, gamma=args.lr_decay_rate)

    max_iter = int(len(train_dataset) / args.batch_size + 0.5)

    minimum_loss = 1e4

    best_epoch = 0

    lossfile_train = args.log_dir + "/training_loss.txt"

    lossfile_val = args.log_dir + "/val_loss.txt"

    lossfile_geometry_train = args.log_dir + "/training_calculate_inference_loss.txt"

    lossfile_geometry_val = args.log_dir + "/val_calculate_inference_loss.txt"

    lossfile_compactness_train = args.log_dir + "/training_compactness_loss.txt"

    lossfile_compactness_val = args.log_dir + "/val_compactness_loss.txt"

    lossfile_KL_train = args.log_dir + "/training_KL_loss.txt"

    lossfile_KL_val = args.log_dir + "/val_KL_loss.txt"

    lossfile_PC_train = args.log_dir + "/training_PC_loss.txt"

    lossfile_PC_val = args.log_dir + "/val_PC_loss.txt"

    lossfile_ecg_train = args.log_dir + "/training_ecg_loss.txt"

    lossfile_ecg_val = args.log_dir + "/val_ecg_loss.txt"

    lossfile_RVp_train = args.log_dir + "/training_RVp_loss.txt"

    lossfile_RVp_val = args.log_dir + "/val_RVp_loss.txt"

    lossfile_size_train = args.log_dir + "/training_MIsize_loss.txt"

    lossfile_size_val = args.log_dir + "/val_MIsize_loss.txt"

    lamda_KL = args.lamda_KL

    for epoch in range(1, args.epochs + 1):

        if ((epoch % 25) == 0) and (epoch != 0):  

            lossplot_detailed(lossfile_train, lossfile_val, lossfile_geometry_train, lossfile_geometry_val, lossfile_KL_train, lossfile_KL_val, lossfile_compactness_train, lossfile_compactness_val, lossfile_PC_train, lossfile_PC_val, lossfile_ecg_train, lossfile_ecg_val, lossfile_RVp_train, lossfile_RVp_val, lossfile_size_train, lossfile_size_val)

        f_train = open(lossfile_train, 'a')  # a: additional writing; w: overwrite writing

        f_val = open(lossfile_val, 'a')

        f_MI_train = open(lossfile_geometry_train, 'a')  # a: additional writing; w: overwrite writing

        f_MI_val = open(lossfile_geometry_val, 'a')

        f_compactness_train = open(lossfile_compactness_train, 'a')  # a: additional writing; w: overwrite writing

        f_compactness_val = open(lossfile_compactness_val, 'a')

        f_KL_train = open(lossfile_KL_train, 'a')  # a: additional writing; w: overwrite writing

        f_KL_val = open(lossfile_KL_val, 'a')

        f_PC_train = open(lossfile_PC_train, 'a')  # a: additional writing; w: overwrite writing

        f_PC_val = open(lossfile_PC_val, 'a')

        f_ecg_train = open(lossfile_ecg_train, 'a')  # a: additional writing; w: overwrite writing

        f_ecg_val = open(lossfile_ecg_val, 'a')

        f_size_train = open(lossfile_size_train, 'a')  # a: additional writing; w: overwrite writing

        f_size_val = open(lossfile_size_val, 'a')

        f_RVp_train = open(lossfile_RVp_train, 'a')  # a: additional writing; w: overwrite writing

        f_RVp_val = open(lossfile_RVp_val, 'a')

        # if epoch != 0: 

        #     if  lamda_KL < 1:

        #         lamda_KL = 0.1*epoch*args.lamda_KL 

        #     else:

        #         lamda_KL = 0.1

        # training

        network.train()

        total_loss, iter_count = 0, 0

        for i, data in enumerate(train_dataloader, 1):

            partial_input, ECG_input, gt_MI, partial_input_coarse, MI_type = data

            partial_input, ECG_input, gt_MI = partial_input.to(DEVICE), ECG_input.to(DEVICE), gt_MI.to(DEVICE)      

            partial_input_coarse = partial_input_coarse.to(DEVICE)      

            partial_input = partial_input.permute(0, 2, 1)

            optimizer.zero_grad()

            y_MI, y_coarse, y_detail, y_ECG, mu, log_var = network(partial_input[:, 0:7, :], ECG_input)

            loss_seg, loss_compactness, loss_MI_RVpenalty, loss_MI_size, KL_loss = calculate_inference_loss(y_MI, gt_MI, mu, log_var, partial_input)

            loss_geo, loss_signal = calculate_reconstruction_loss(y_coarse, y_detail, partial_input_coarse, partial_input, y_ECG, ECG_input)

            loss = loss_seg + args.lamda_compact*loss_compactness + args.lamda_RVp*loss_MI_RVpenalty + args.lamda_MIsize*loss_MI_size + args.lamda_KL*KL_loss + args.lamda_recon*loss_geo # + args.lamda_recon*loss_signal # 

            check_grad = False

            if check_grad:

                print(loss.requires_grad)

                print(loss_seg.requires_grad)

                print(loss_compactness.requires_grad)

                print(loss_MI_RVpenalty.requires_grad)

                print(KL_loss.requires_grad)

                print(loss_MI_size.requires_grad)

                print(loss_geo.requires_grad)

                print(loss_signal.requires_grad)

            visual_check = False

            if visual_check:

                y_predict = y_MI[0].cpu().detach().numpy()

                y_gd = gt_MI[0].cpu().detach().numpy()

                x_input = partial_input[0].cpu().detach().numpy()

                y_predict_argmax = np.argmax(y_predict, axis=0)

                visualize_PC_with_twolabel(x_input[0:3, 0:args.num_input].transpose(), y_predict_argmax, y_gd, filename='RNmap_gd_pre.jpg')

            loss.backward()

            optimizer.step()

            f_train.write(str(loss.item()))

            f_train.write('\n')

            f_MI_train.write(str(loss_seg.item()))

            f_MI_train.write('\n')

            f_compactness_train.write(str(loss_compactness.item()))

            f_compactness_train.write('\n')

            f_KL_train.write(str(KL_loss.item()))

            f_KL_train.write('\n')

            f_PC_train.write(str(loss_geo.item()))

            f_PC_train.write('\n')

            f_ecg_train.write(str(loss_signal.item()))

            f_ecg_train.write('\n')

            f_size_train.write(str((loss_MI_size.item())))

            f_size_train.write('\n')

            f_RVp_train.write(str(loss_MI_RVpenalty.item()))

            f_RVp_train.write('\n')

            iter_count += 1

            total_loss += loss.item()

            if i % 50 == 0:

                print("Training epoch {}/{}, iteration {}/{}: loss is {}".format(epoch, args.epochs, i, max_iter, loss.item()))

        scheduler.step()

        print("\033[96mTraining epoch {}/{}: avg loss = {}\033[0m".format(epoch, args.epochs, total_loss / iter_count))

        # evaluation

        network.eval()

        with torch.no_grad():

            total_loss, iter_count = 0, 0

            for i, data in enumerate(val_dataloader, 1):

                partial_input, ECG_input, gt_MI, partial_input_coarse, MI_type = data

                partial_input, ECG_input, gt_MI = partial_input.to(DEVICE), ECG_input.to(DEVICE), gt_MI.to(DEVICE)  

                partial_input_coarse = partial_input_coarse.to(DEVICE)  

                partial_input = partial_input.permute(0, 2, 1)

                y_MI, y_coarse, y_detail, y_ECG, mu, log_var = network(partial_input[:, 0:7, :], ECG_input)

                loss_seg, loss_compactness, loss_MI_RVpenalty, loss_MI_size, KL_loss = calculate_inference_loss(y_MI, gt_MI, mu, log_var, partial_input)

                loss_geo, loss_signal = calculate_reconstruction_loss(y_coarse, y_detail, partial_input_coarse, partial_input, y_ECG, ECG_input)

                loss = loss_seg + args.lamda_compact*loss_compactness + args.lamda_RVp*loss_MI_RVpenalty + args.lamda_MIsize*loss_MI_size + args.lamda_KL*KL_loss + args.lamda_recon*loss_geo # + args.lamda_recon*loss_signal # 

                total_loss += loss.item()

                iter_count += 1

                if ((epoch % 25) == 0) and (epoch != 0) and (i == 1):  

                    y_predict = y_MI[0].cpu().detach().numpy()

                    y_gd = gt_MI[0].cpu().detach().numpy()

                    x_input = partial_input[0].cpu().detach().numpy()

                    y_predict_argmax = np.argmax(y_predict, axis=0)

                    visualize_PC_with_twolabel(x_input[0:3, 0:args.num_input].transpose(), y_predict_argmax, y_gd, filename='RNmap_gd_pre.jpg')

                f_val.write(str(loss.item()))

                f_val.write('\n')

                f_MI_val.write(str(loss_seg.item()))

                f_MI_val.write('\n')

                f_compactness_val.write(str(loss_compactness.item()))

                f_compactness_val.write('\n')

                f_KL_val.write(str(KL_loss.item()))

                f_KL_val.write('\n')

                f_PC_val.write(str(loss_geo.item()))

                f_PC_val.write('\n')

                f_ecg_val.write(str(loss_signal.item()))

                f_ecg_val.write('\n')

                f_size_val.write(str(loss_MI_size.item()))

                f_size_val.write('\n')

                f_RVp_val.write(str(loss_MI_RVpenalty.item()))

                f_RVp_val.write('\n')

            mean_loss = total_loss / iter_count

            print("\033[35mValidation epoch {}/{}, loss is {}\033[0m".format(epoch, args.epochs, mean_loss))

            # records the best model and epoch

            if mean_loss < minimum_loss:

                best_epoch = epoch

                minimum_loss = mean_loss           

                strNetSaveName = 'net_model.pkl'

                # strNetSaveName = 'net_with_%d.pkl' % epoch

                torch.save(network.state_dict(), args.log_dir + '/' + strNetSaveName)

        print("\033[4;37mBest model (lowest loss) in epoch {}\033[0m".format(best_epoch))

    lossplot(lossfile_train, lossfile_val)

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument('--partial_root', type=str, default='./Big_data_inference/meta_data/UKB_clinical_data/')

    parser.add_argument('--model', type=str, default=None) #'log/net_model.pkl'

    parser.add_argument('--in_ch', type=int, default=3+4) # coordinate dimension + label index

    parser.add_argument('--out_ch', type=int, default=3) # 3scar, BZ, normal/ 18 for ecg-based classification

    parser.add_argument('--z_dims', type=int, default=16)

    parser.add_argument('--num_input', type=int, default=1024*4)

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

    parser.add_argument('--lamda_recon', type=float, default=1) # 1

    parser.add_argument('--lamda_KL', type=float, default=1e-2) # 1e-2

    parser.add_argument('--lamda_MIsize', type=float, default=1) # 1

    parser.add_argument('--lamda_RVp', type=float, default=1) # 1 

    parser.add_argument('--lamda_compact', type=float, default=1) # 1

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

    parser.add_argument('--lr_decay_steps', type=int, default=50) 

    parser.add_argument('--lr_decay_rate', type=float, default=0.5) 

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

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

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

    parser.add_argument('--log_dir', type=str, default='log')

    args = parser.parse_args()

    train(args)