#!/usr/bin/env python3

# -*- coding: utf-8 -*-

"""

@author: Zhi Huang

"""

import sys, os

sys.path.append("/home/zhihuan/Documents/SALMON/model")

import SALMON

import pandas as pd

import argparse

import torch

import torch.nn as nn

import torch.nn.functional as F

import torch.optim as optim

from torch.utils.data import DataLoader

from torchvision import transforms

from torch.autograd import Variable

from collections import Counter

import pandas as pd

import math

import random

from imblearn.over_sampling import RandomOverSampler

from lifelines.statistics import logrank_test

import json

import tables

import logging

import csv

import numpy as np

import optunity

import pickle

import time

from sklearn.model_selection import KFold

from sklearn import preprocessing

import matplotlib

matplotlib.use('Agg')

import matplotlib.pyplot as plt

plt.ioff()

def parse_args():

    parser = argparse.ArgumentParser()

    parser.add_argument('--dataset_dir', default='/home/zhihuan/Documents/SALMON/data/LUAD/multiomics_preprocessing_results/', help="datasets")

    parser.add_argument('--num_epochs', type=int, default=100, help="Number of epochs to train for. Default: 100")

    parser.add_argument('--measure_while_training', action='store_true', default=False, help='disables measure while training (make program faster)')

    parser.add_argument('--batch_size', type=int, default=64, help="Number of batches to train/test for. Default: 256")

    parser.add_argument('--dataset', type=int, default=7)

    parser.add_argument('--nocuda', action='store_true', default=False, help='disables CUDA training')

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

    parser.add_argument('--results_dir', default='/home/zhihuan/Documents/SALMON/experiments/Results/LUAD', help="results dir")

    return parser.parse_args()

if __name__=='__main__':

    torch.cuda.empty_cache()

    args = parse_args()

    # model file

    num_epochs = args.num_epochs

    batch_size = args.batch_size

    learning_rate_range = 10**np.arange(-4,-1,0.3)

    cuda = True

    verbose = 0

    measure_while_training = True

    dropout_rate = 0

    lambda_1 = 1e-6 # L1

    # 5-fold data

    tempdata = {}

    tempdata['clinical'] = pd.read_csv(args.dataset_dir + 'clinical.csv', index_col = 0).reset_index(drop = True)

    tempdata['mRNAseq_eigengene'] = pd.read_csv(args.dataset_dir + 'mRNAseq_eigengene_matrix.csv', index_col = 0).reset_index(drop = True)

    tempdata['miRNAseq_eigengene'] = pd.read_csv(args.dataset_dir + 'miRNAseq_eigengene_matrix.csv', index_col = 0).reset_index(drop = True)

    tempdata['TMB'] = pd.read_csv(args.dataset_dir + 'TMB.csv', index_col = 0).reset_index(drop = True)

    tempdata['CNB'] = pd.read_csv(args.dataset_dir + 'CNB.csv', index_col = 0).reset_index(drop = True)

    tempdata['CNB']['log2_LENGTH_KB'] = np.log2(tempdata['CNB']['LENGTH_KB'].values + 1)

    print('0:MALE\t\t1:FEMALE\n0:Alive\t\t1:Dead')

    tempdata['clinical']['gender'] = (tempdata['clinical']['gender'].values == 'MALE').astype(int)

    tempdata['clinical']['vital_status'] = (tempdata['clinical']['vital_status'].values == 'Dead').astype(int)

    data = {}

    data['x'] = pd.concat((tempdata['mRNAseq_eigengene'], tempdata['miRNAseq_eigengene'], tempdata['CNB']['log2_LENGTH_KB'], tempdata['TMB']['All_TMB'], tempdata['clinical'][['gender','age_at_initial_pathologic_diagnosis']]), axis = 1).values.astype(np.double)

    all_column_names = ['mRNAseq_' + str(i+1) for i in range(tempdata['mRNAseq_eigengene'].shape[1])] + \

                            ['miRNAseq_' + str(i+1) for i in range(tempdata['miRNAseq_eigengene'].shape[1])] + \

                            ['CNB', 'TMB', 'GENDER', 'AGE']

    print('perform min-max scaler on all input features')

    scaler = preprocessing.MinMaxScaler()

    scaler.fit(data['x'])

    data['x'] = scaler.transform(data['x'])

    data['e'] = tempdata['clinical']['vital_status'].values.astype(np.int32)

    data['t'] = tempdata['clinical']['survival_days'].values.astype(np.double)

    if args.dataset == 1:

        dataset_subset = "1_RNAseq"

        data['column_names'] = ['mRNAseq_' + str(i+1) for i in range(tempdata['mRNAseq_eigengene'].shape[1])]

    elif args.dataset == 2:

        dataset_subset = "2_miRNAseq"

        data['column_names'] = ['miRNAseq_' + str(i+1) for i in range(tempdata['miRNAseq_eigengene'].shape[1])]

    elif args.dataset == 3:

        dataset_subset = "3_RNAseq+miRNAseq"

        data['column_names'] = ['mRNAseq_' + str(i+1) for i in range(tempdata['mRNAseq_eigengene'].shape[1])] + \

                                ['miRNAseq_' + str(i+1) for i in range(tempdata['miRNAseq_eigengene'].shape[1])]

    elif args.dataset == 4:

        dataset_subset = "4_RNAseq+miRNAseq+cnb+tmb"

        data['column_names'] = ['mRNAseq_' + str(i+1) for i in range(tempdata['mRNAseq_eigengene'].shape[1])] + \

                                ['miRNAseq_' + str(i+1) for i in range(tempdata['miRNAseq_eigengene'].shape[1])] + \

                                ['CNB', 'TMB']

    elif args.dataset == 5:

        dataset_subset = "5_RNAseq+miRNAseq+clinical"

        data['column_names'] = ['mRNAseq_' + str(i+1) for i in range(tempdata['mRNAseq_eigengene'].shape[1])] + \

                                ['miRNAseq_' + str(i+1) for i in range(tempdata['miRNAseq_eigengene'].shape[1])] + \

                                ['GENDER', 'AGE']

    elif args.dataset == 6:

        dataset_subset = "6_cnb+tmb+clinical"

        data['column_names'] = ['CNB', 'TMB', 'GENDER', 'AGE']

    elif args.dataset == 7:

        dataset_subset = "7_RNAseq+miRNAseq+cnb+tmb+clinical"

        data['column_names'] = ['mRNAseq_' + str(i+1) for i in range(tempdata['mRNAseq_eigengene'].shape[1])] + \

                                ['miRNAseq_' + str(i+1) for i in range(tempdata['miRNAseq_eigengene'].shape[1])] + \

                                ['CNB', 'TMB', 'GENDER', 'AGE']

    print('subsetting data...')

    data['x'] = data['x'][:, [i for i, c in enumerate(all_column_names) if c in data['column_names']]]

    kf = KFold(n_splits=5, shuffle=True, random_state=666)

    datasets_5folds = {}

    for ix, (train_index, test_index) in enumerate(kf.split(data['x']), start = 1):

        datasets_5folds[ix] = {}

        datasets_5folds[ix]['train'] = {}

        datasets_5folds[ix]['train']['x'] = data['x'][train_index, :]

        datasets_5folds[ix]['train']['e'] = data['e'][train_index]

        datasets_5folds[ix]['train']['t'] = data['t'][train_index]

        datasets_5folds[ix]['test'] = {}

        datasets_5folds[ix]['test']['x'] = data['x'][train_index, :]

        datasets_5folds[ix]['test']['e'] = data['e'][train_index]

        datasets_5folds[ix]['test']['t'] = data['t'][train_index]

    for i in range(1, len(datasets_5folds) + 1):

        print("5 fold CV -- %d/5" % i)

        # dataset

        TIMESTRING  = time.strftime("%Y%m%d-%H.%M.%S", time.localtime())

        results_dir_dataset = args.results_dir + '/' + dataset_subset + '/run_' + TIMESTRING + '_fold_' + str(i)

        if not os.path.exists(results_dir_dataset):

            os.makedirs(results_dir_dataset)

        logging.basicConfig(filename=results_dir_dataset+'/mainlog.log',level=logging.DEBUG)

    #    print("Arguments:",args)

    #    logging.info("Arguments: %s" % args)

        datasets = datasets_5folds[i]

        length_of_data = {}

        length_of_data['mRNAseq'] = tempdata['mRNAseq_eigengene'].shape[1]

        length_of_data['miRNAseq'] = tempdata['miRNAseq_eigengene'].shape[1]

        length_of_data['CNB'] = 1

        length_of_data['TMB'] = 1

        length_of_data['clinical'] = 2

    # =============================================================================

    # # Finding optimal learning rate w.r.t. concordance index

    # =============================================================================

        ci_list = []

        for j, lr in enumerate(learning_rate_range):

            print("[%d/%d] current lr: %.4E" %((j+1), len(learning_rate_range), lr))

            logging.info("[%d/%d] current lr: %.4E" %((j+1), len(learning_rate_range), lr))

            model, loss_nn_all, pvalue_all, c_index_all, c_index_list, acc_train_all, code_output = \

                 SALMON.train(datasets, num_epochs, batch_size, lr, dropout_rate,\

                                         lambda_1, length_of_data, cuda, measure_while_training, verbose)

            epochs_list = range(num_epochs)

            plt.figure(figsize=(8,4))

            plt.plot(epochs_list, c_index_list['train'], "b--",linewidth=1)

            plt.plot(epochs_list, c_index_list['test'], "g-",linewidth=1)

            plt.legend(['train', 'test'])

            plt.xlabel("epochs")

            plt.ylabel("Concordance index")

            plt.savefig(results_dir_dataset + "/convergence_%02d_lr=%.2E.png" % (j, lr),dpi=300)

            plt.close()

            code_test, loss_nn_sum, acc_test, pvalue_pred, c_index_pred, lbl_pred_all, OS_event_test, OS_test = \

                SALMON.test(model, datasets, 'test', length_of_data, batch_size, cuda, verbose)

            ci_list.append(c_index_pred)

            print("current concordance index: ", c_index_pred,"\n")

            logging.info("current concordance index: %.10f\n" % c_index_pred)

        optimal_lr = learning_rate_range[np.argmax(ci_list)]

        print("Optimal learning rate: %.4E, optimal c-index: %.10f" % (optimal_lr, np.max(ci_list)))

        logging.info("Optimal learning rate: %.4E, optimal c-index: %.10f" % (optimal_lr, np.max(ci_list)))

    # =============================================================================

    # # Training 

    # =============================================================================

        model, loss_nn_all, pvalue_all, c_index_all, c_index_list, acc_train_all, code_output = \

                 SALMON.train(datasets, num_epochs, batch_size, optimal_lr, dropout_rate,\

                                         lambda_1, length_of_data, cuda, measure_while_training, verbose)

        code_train, loss_nn_sum, acc_train, pvalue_pred, c_index_pred, lbl_pred_all_train, OS_event_train, OS_train = \

            SALMON.test(model, datasets, 'train', length_of_data, batch_size, cuda, verbose)

        print("[Final] Apply model to training set: c-index: %.10f, p-value: %.10e" % (c_index_pred, pvalue_pred))

        logging.info("[Final] Apply model to training set: c-index: %.10f, p-value: %.10e" % (c_index_pred, pvalue_pred))

        code_test, loss_nn_sum, acc_test, pvalue_pred, c_index_pred, lbl_pred_all_test, OS_event_test, OS_test = \

            SALMON.test(model, datasets, 'test', length_of_data, batch_size, cuda, verbose)

        print("[Final] Apply model to testing set: c-index: %.10f, p-value: %.10e" % (c_index_pred, pvalue_pred))

        logging.info("[Final] Apply model to testing set: c-index: %.10f, p-value: %.10e" % (c_index_pred, pvalue_pred))

        with open(results_dir_dataset + '/model.pickle', 'wb') as handle:

            pickle.dump(model, handle, protocol=pickle.HIGHEST_PROTOCOL)

        with open(results_dir_dataset + '/c_index_list_by_epochs.pickle', 'wb') as handle:

            pickle.dump(c_index_list, handle, protocol=pickle.HIGHEST_PROTOCOL)

        with open(results_dir_dataset + '/hazard_ratios_lbl_pred_all_train.pickle', 'wb') as handle:

            pickle.dump(lbl_pred_all_train, handle, protocol=pickle.HIGHEST_PROTOCOL)

        with open(results_dir_dataset + '/OS_event_train.pickle', 'wb') as handle:

            pickle.dump(OS_event_train, handle, protocol=pickle.HIGHEST_PROTOCOL)

        with open(results_dir_dataset + '/OS_train.pickle', 'wb') as handle:

            pickle.dump(OS_train, handle, protocol=pickle.HIGHEST_PROTOCOL)

        with open(results_dir_dataset + '/code_train.pickle', 'wb') as handle:

            pickle.dump(code_train, handle, protocol=pickle.HIGHEST_PROTOCOL)

        with open(results_dir_dataset + '/hazard_ratios_lbl_pred_all_test.pickle', 'wb') as handle:

            pickle.dump(lbl_pred_all_test, handle, protocol=pickle.HIGHEST_PROTOCOL)

        with open(results_dir_dataset + '/OS_event_test.pickle', 'wb') as handle:

            pickle.dump(OS_event_test, handle, protocol=pickle.HIGHEST_PROTOCOL)

        with open(results_dir_dataset + '/OS_test.pickle', 'wb') as handle:

            pickle.dump(OS_test, handle, protocol=pickle.HIGHEST_PROTOCOL)

        with open(results_dir_dataset + '/code_test.pickle', 'wb') as handle:

            pickle.dump(code_test, handle, protocol=pickle.HIGHEST_PROTOCOL)

        epochs_list = range(num_epochs)

        plt.figure(figsize=(8,4))

        plt.plot(epochs_list, c_index_list['train'], "b--",linewidth=1)

        plt.plot(epochs_list, c_index_list['test'], "g-",linewidth=1)

        plt.legend(['train', 'test'])

        plt.xlabel("epochs")

        plt.ylabel("Concordance index")

        plt.savefig(results_dir_dataset + "/convergence.png",dpi=300)

        plt.close()