#!/usr/bin/env python3

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

"""

@author: Zhi Huang

"""

import argparse, random

import torch

import torch.nn as nn

import torch.backends.cudnn as cudnn

import torch.nn.functional as F

import torch.optim as optim

from torch.utils.data import DataLoader

from torchvision import datasets, transforms

from torch.autograd import Variable

from collections import Counter

import pandas as pd

import matplotlib.pyplot as plt

import math

import random

from imblearn.over_sampling import RandomOverSampler

import pandas as pd

from lifelines.statistics import logrank_test

from lifelines.utils import concordance_index

import tables

import csv

import numpy as np

import json

from tqdm import tqdm

import gc

import copy

class SALMON(nn.Module):

    def __init__(self, input_dim, dropout_rate, length_of_data, label_dim):

        super(SALMON, self).__init__()

        self.length_of_data = length_of_data

        hidden1 = 8

        hidden2 = 4

        if input_dim == length_of_data['mRNAseq']: # mRNAseq

            self.encoder1 = nn.Sequential(nn.Linear(input_dim, hidden1),nn.Sigmoid())

            self.classifier = nn.Sequential(nn.Linear(hidden1, label_dim),nn.Sigmoid())

        if input_dim == length_of_data['miRNAseq']: # miRNAseq

            self.encoder2 = nn.Sequential(nn.Linear(input_dim, hidden2),nn.Sigmoid())

            self.classifier = nn.Sequential(nn.Linear(hidden2, label_dim),nn.Sigmoid())

        if input_dim == length_of_data['mRNAseq'] + length_of_data['miRNAseq']: # mRNAseq + miRNAseq

            self.encoder1 = nn.Sequential(nn.Linear(length_of_data['mRNAseq'], hidden1),nn.Sigmoid())

            self.encoder2 = nn.Sequential(nn.Linear(length_of_data['miRNAseq'], hidden2),nn.Sigmoid())

            self.classifier = nn.Sequential(nn.Linear(hidden1 + hidden2, label_dim),nn.Sigmoid())

        if input_dim == length_of_data['mRNAseq'] + length_of_data['miRNAseq'] + length_of_data['CNB'] + length_of_data['TMB']: # mRNAseq + miRNAseq + CNB + TMB

            hidden_cnv, hidden_tmb = length_of_data['CNB'], length_of_data['TMB']

            self.encoder1 = nn.Sequential(nn.Linear(length_of_data['mRNAseq'], hidden1),nn.Sigmoid())

            self.encoder2 = nn.Sequential(nn.Linear(length_of_data['miRNAseq'], hidden2),nn.Sigmoid())

            self.classifier = nn.Sequential(nn.Linear(hidden1 + hidden2 + hidden_cnv + hidden_tmb, label_dim),nn.Sigmoid())

        if input_dim == length_of_data['mRNAseq'] + length_of_data['miRNAseq'] + length_of_data['CNB'] + length_of_data['TMB'] + length_of_data['clinical']: # mRNAseq + miRNAseq + CNB + TMB + clinical

            hidden_cnv, hidden_tmb, hidden_clinical = length_of_data['CNB'], length_of_data['TMB'], length_of_data['clinical']

            self.encoder1 = nn.Sequential(nn.Linear(length_of_data['mRNAseq'], hidden1),nn.Sigmoid())

            self.encoder2 = nn.Sequential(nn.Linear(length_of_data['miRNAseq'], hidden2),nn.Sigmoid())

            self.classifier = nn.Sequential(nn.Linear(hidden1 + hidden2 + \

                                            hidden_cnv + hidden_tmb + hidden_clinical, label_dim),nn.Sigmoid())

        if input_dim == length_of_data['CNB'] + length_of_data['TMB'] + length_of_data['clinical']: # CNB + TMB + clinical

            hidden_cnv, hidden_tmb, hidden_clinical = length_of_data['CNB'], length_of_data['TMB'], length_of_data['clinical']

            self.classifier = nn.Sequential(nn.Linear(hidden_cnv + hidden_tmb + hidden_clinical, label_dim),nn.Sigmoid())

        if input_dim == length_of_data['mRNAseq'] + length_of_data['miRNAseq'] + length_of_data['clinical']: # mRNAseq + miRNAseq + clinical

            hidden_clinical = length_of_data['clinical']

            self.encoder1 = nn.Sequential(nn.Linear(length_of_data['mRNAseq'], hidden1),nn.Sigmoid())

            self.encoder2 = nn.Sequential(nn.Linear(length_of_data['miRNAseq'], hidden2),nn.Sigmoid())

            self.classifier = nn.Sequential(nn.Linear(hidden1 + hidden2 + \

                                            hidden_clinical, label_dim),nn.Sigmoid())

    def forward(self, x):

        input_dim = x.shape[1]

        x_d = None

        if input_dim == self.length_of_data['mRNAseq']: # mRNAseq

            code1 = self.encoder1(x)

            lbl_pred = self.classifier(code1) # predicted label

            code = code1

        if input_dim == self.length_of_data['miRNAseq']: # miRNAseq

            code2 = self.encoder2(x)

            lbl_pred = self.classifier(code2) # predicted label

            code = code2

        if input_dim == self.length_of_data['mRNAseq'] + self.length_of_data['miRNAseq']: # mRNAseq + miRNAseq

            code1 = self.encoder1(x[:,0:self.length_of_data['mRNAseq']])

            code2 = self.encoder2(x[:,self.length_of_data['mRNAseq']:])

            lbl_pred = self.classifier(torch.cat((code1, code2), 1)) # predicted label

            code = torch.cat((code1, code2), 1)

        if input_dim == self.length_of_data['mRNAseq'] + self.length_of_data['miRNAseq'] + self.length_of_data['CNB'] + self.length_of_data['TMB']: # mRNAseq + miRNAseq + CNB + TMB

            code1 = self.encoder1(x[:,0:self.length_of_data['mRNAseq']])

            code2 = self.encoder2(x[:,self.length_of_data['mRNAseq']: (self.length_of_data['mRNAseq'] + self.length_of_data['miRNAseq'])])

            lbl_pred = self.classifier(torch.cat((code1, code2, x[:,(self.length_of_data['mRNAseq'] + self.length_of_data['miRNAseq']):]), 1)) # predicted label

            code = torch.cat((code1, code2), 1)

        if input_dim == self.length_of_data['mRNAseq'] + self.length_of_data['miRNAseq'] + self.length_of_data['CNB'] + self.length_of_data['TMB'] + self.length_of_data['clinical']: # mRNAseq + miRNAseq + CNB + TMB + clinical

            code1 = self.encoder1(x[:,0:self.length_of_data['mRNAseq']])

            code2 = self.encoder2(x[:,self.length_of_data['mRNAseq']: (self.length_of_data['mRNAseq'] + self.length_of_data['miRNAseq'])])

            lbl_pred = self.classifier(torch.cat((code1, code2, x[:, (self.length_of_data['mRNAseq'] + self.length_of_data['miRNAseq']):]), 1)) # predicted label

            code = torch.cat((code1, code2), 1)

        if input_dim == self.length_of_data['CNB'] + self.length_of_data['TMB'] + self.length_of_data['clinical']: # CNB + TMB + clinical

            lbl_pred = self.classifier(x) # predicted label

            code = torch.FloatTensor([0])

        if input_dim == self.length_of_data['mRNAseq'] + self.length_of_data['miRNAseq'] + self.length_of_data['clinical']: # mRNAseq + miRNAseq + clinical

            code1 = self.encoder1(x[:,0:self.length_of_data['mRNAseq']])

            code2 = self.encoder2(x[:,self.length_of_data['mRNAseq']: (self.length_of_data['mRNAseq'] + self.length_of_data['miRNAseq'])])

            lbl_pred = self.classifier(torch.cat((code1, code2, x[:, (self.length_of_data['mRNAseq'] + self.length_of_data['miRNAseq']):]), 1)) # predicted label

            code = torch.cat((code1, code2), 1)

        return x_d, code, lbl_pred

def accuracy(output, labels):

    preds = output.max(1)[1].type_as(labels)

    correct = preds.eq(labels).double()

    correct = correct.sum()

    return correct / len(labels)

def accuracy_cox(hazards, labels):

    # This accuracy is based on estimated survival events against true survival events

    hazardsdata = hazards.cpu().numpy().reshape(-1)

    median = np.median(hazardsdata)

    hazards_dichotomize = np.zeros([len(hazardsdata)], dtype=int)

    hazards_dichotomize[hazardsdata > median] = 1

    labels = labels.data.cpu().numpy()

    correct = np.sum(hazards_dichotomize == labels)

    return correct / len(labels)

def cox_log_rank(hazards, labels, survtime_all):

    hazardsdata = hazards.cpu().numpy().reshape(-1)

    median = np.median(hazardsdata)

    hazards_dichotomize = np.zeros([len(hazardsdata)], dtype=int)

    hazards_dichotomize[hazardsdata > median] = 1

    survtime_all = survtime_all.data.cpu().numpy().reshape(-1)

    idx = hazards_dichotomize == 0

    labels = labels.data.cpu().numpy()

    T1 = survtime_all[idx]

    T2 = survtime_all[~idx]

    E1 = labels[idx]

    E2 = labels[~idx]

    results = logrank_test(T1, T2, event_observed_A=E1, event_observed_B=E2)

    pvalue_pred = results.p_value

    return(pvalue_pred)

def CIndex(hazards, labels, survtime_all):

    labels = labels.data.cpu().numpy()

    concord = 0.

    total = 0.

    N_test = labels.shape[0]

    labels = np.asarray(labels, dtype=bool)

    for i in range(N_test):

        if labels[i] == 1:

            for j in range(N_test):

                if survtime_all[j] > survtime_all[i]:

                    total = total + 1

                    if hazards[j] < hazards[i]: concord = concord + 1

                    elif hazards[j] < hazards[i]: concord = concord + 0.5

    return(concord/total)

def CIndex_lifeline(hazards, labels, survtime_all):

    labels = labels.data.cpu().numpy()

    hazards = hazards.cpu().numpy().reshape(-1)

    return(concordance_index(survtime_all, -hazards, labels))

def frobenius_norm_loss(a, b):

    loss = torch.sqrt(torch.sum(torch.abs(a-b)**2))

    return loss

def test(model, datasets, whichset, length_of_data, batch_size, cuda, verbose):

    x = datasets[whichset]['x']

    e = datasets[whichset]['e']

    t = datasets[whichset]['t']

    X = torch.FloatTensor(x)

    OS_event = torch.LongTensor(e)

    OS = torch.FloatTensor(t)

    dataloader = DataLoader(X, batch_size=batch_size, num_workers=1, pin_memory=True, shuffle=False)

    lblloader = DataLoader(OS_event, batch_size=batch_size, num_workers=1, pin_memory=True, shuffle=False)

    OSloader = DataLoader(OS, batch_size=batch_size, num_workers=1, pin_memory=True, shuffle=False)

    lbl_pred_all = None

    lbl_all = None

    survtime_all = None

    code_final = None

    loss_nn_sum = 0

    model.eval()

    iter = 0

    for data, lbl, survtime in zip(dataloader, lblloader, OSloader):

        graph = data

        graph = Variable(graph)

        lbl = Variable(lbl)

        if cuda:

            model = model.cuda()

            graph = graph.cuda()

            lbl = lbl.cuda()

        # ===================forward=====================

        output, code, lbl_pred = model(graph)

        if iter == 0:

            lbl_pred_all = lbl_pred

            lbl_all = lbl

            survtime_all = survtime

            code_final = code

        else:

            lbl_pred_all = torch.cat([lbl_pred_all, lbl_pred])

            lbl_all = torch.cat([lbl_all, lbl])

            survtime_all = torch.cat([survtime_all, survtime])

            code_final = torch.cat([code_final, code])

        current_batch_len = len(survtime)

        R_matrix_test = np.zeros([current_batch_len, current_batch_len], dtype=int)

        for i in range(current_batch_len):

            for j in range(current_batch_len):

                R_matrix_test[i,j] = survtime[j] >= survtime[i]

        test_R = torch.FloatTensor(R_matrix_test)

        test_R = Variable(test_R)

        if cuda:

            test_R = test_R.cuda()

        test_ystatus = lbl

        theta = lbl_pred.reshape(-1)

        exp_theta = torch.exp(theta)

        loss_nn = -torch.mean( (theta - torch.log(torch.sum( exp_theta*test_R ,dim=1))) * test_ystatus.float() )

        loss_nn_sum = loss_nn_sum + loss_nn.data.item()

        iter += 1

    code_final_4_original_data = code_final.data.cpu().numpy()

    acc_test = accuracy_cox(lbl_pred_all.data, lbl_all)

    pvalue_pred = cox_log_rank(lbl_pred_all.data, lbl_all, survtime_all)

    c_index = CIndex_lifeline(lbl_pred_all.data, lbl_all, survtime_all)

    if verbose > 0:

        print('\n[{:s}]\t\tloss (nn):{:.4f}'.format(whichset, loss_nn_sum),

                      'c_index: {:.4f}, p-value: {:.3e}'.format(c_index, pvalue_pred))

    return(code_final_4_original_data, loss_nn_sum, acc_test, \

           pvalue_pred, c_index, lbl_pred_all.data.cpu().numpy().reshape(-1), OS_event, survtime_all)

def init_weights(m):

    if type(m) == nn.Linear:

        m.weight.data.normal_(0, 0.5)

def train(datasets, num_epochs, batch_size, learning_rate, dropout_rate,

                        lambda_1, length_of_data, cuda, measure, verbose):

    x = datasets['train']['x']

    e = datasets['train']['e']

    t = datasets['train']['t']

    nodes_in = x.shape[1]

    X = torch.FloatTensor(x)

    OS_event = torch.LongTensor(e)

    OS = torch.FloatTensor(t)

    dataloader = DataLoader(X, batch_size=batch_size, num_workers=0, pin_memory=True, shuffle=False)

    lblloader = DataLoader(OS_event, batch_size=batch_size, num_workers=0, pin_memory=True, shuffle=False)

    OSloader = DataLoader(OS, batch_size=batch_size, num_workers=0, pin_memory=True, shuffle=False)

    cudnn.deterministic = True

    torch.cuda.manual_seed_all(666)

    torch.manual_seed(666)

    random.seed(666)

    model = SALMON(nodes_in, dropout_rate, length_of_data, label_dim = 1)

    if cuda:

        model.cuda()

    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=0)

    c_index_list = {}

    c_index_list['train'] = []

    c_index_list['test'] = []

    loss_nn_all = []

    pvalue_all = []

    c_index_all = []

    acc_train_all = []

    c_index_best = 0

    code_output = None

    for epoch in tqdm(range(num_epochs)):

        model.train()

        lbl_pred_all = None

        lbl_all = None

        survtime_all = None

        code_final = None

        loss_nn_sum = 0

        iter = 0

        gc.collect()

        for data, lbl, survtime in zip(dataloader, lblloader, OSloader):

            optimizer.zero_grad() # zero the gradient buffer

            graph = data

            if cuda:

                model = model.cuda()

                graph = graph.cuda()

                lbl = lbl.cuda()

            # ===================forward=====================

            output, code, lbl_pred = model(graph)

            if iter == 0:

                lbl_pred_all = lbl_pred

                survtime_all = survtime

                lbl_all = lbl

                code_final = code

            else:

                lbl_pred_all = torch.cat([lbl_pred_all, lbl_pred])

                lbl_all = torch.cat([lbl_all, lbl])

                survtime_all = torch.cat([survtime_all, survtime])

                code_final = torch.cat([code_final, code])

            # This calculation credit to Travers Ching https://github.com/traversc/cox-nnet

            # Cox-nnet: An artificial neural network method for prognosis prediction of high-throughput omics data

            current_batch_len = len(survtime)

            R_matrix_train = np.zeros([current_batch_len, current_batch_len], dtype=int)

            for i in range(current_batch_len):

                for j in range(current_batch_len):

                    R_matrix_train[i,j] = survtime[j] >= survtime[i]

            train_R = torch.FloatTensor(R_matrix_train)

            if cuda:

                train_R = train_R.cuda()

            train_ystatus = lbl

            theta = lbl_pred.reshape(-1)

            exp_theta = torch.exp(theta)

            loss_nn = -torch.mean( (theta - torch.log(torch.sum( exp_theta*train_R ,dim=1))) * train_ystatus.float() )

            l1_reg = None

            for W in model.parameters():

                if l1_reg is None:

                    l1_reg = torch.abs(W).sum()

                else:

                    l1_reg = l1_reg + torch.abs(W).sum() # torch.abs(W).sum() is equivalent to W.norm(1)

            loss = loss_nn + lambda_1 * l1_reg

            if verbose > 0:

                print("\nloss_nn: %.4f, L1: %.4f" % (loss_nn, lambda_1 * l1_reg))

            loss_nn_sum = loss_nn_sum + loss_nn.data.item()

            # ===================backward====================

            loss.backward()

            optimizer.step()

            iter += 1

            torch.cuda.empty_cache()

        code_final_4_original_data = code_final.data.cpu().numpy()

        if measure or epoch == (num_epochs - 1):

            acc_train = accuracy_cox(lbl_pred_all.data, lbl_all)

            pvalue_pred = cox_log_rank(lbl_pred_all.data, lbl_all, survtime_all)

            c_index = CIndex_lifeline(lbl_pred_all.data, lbl_all, survtime_all)

            c_index_list['train'].append(c_index)

            if c_index > c_index_best:

                c_index_best = c_index

                code_output = code_final_4_original_data

            if verbose > 0:

                print('\n[Training]\t loss (nn):{:.4f}'.format(loss_nn_sum),

                      'c_index: {:.4f}, p-value: {:.3e}'.format(c_index, pvalue_pred))

            pvalue_all.append(pvalue_pred)

            c_index_all.append(c_index)

            loss_nn_all.append(loss_nn_sum)

            acc_train_all.append(acc_train)

            whichset = 'test'

            code_validation, loss_nn_sum, acc_test, pvalue_pred, c_index_pred, lbl_pred_all, OS_event, OS = \

                test(model, datasets, whichset, length_of_data, batch_size, cuda, verbose)

            c_index_list['test'].append(c_index_pred)

    return(model, loss_nn_all, pvalue_all, c_index_all, c_index_list, acc_train_all, code_output)