# Base / Native

import math

import os

import pickle

import re

import warnings

warnings.filterwarnings('ignore')

# Numerical / Array

import lifelines

from lifelines.utils import concordance_index

from lifelines import CoxPHFitter

from lifelines.datasets import load_regression_dataset

from lifelines.utils import k_fold_cross_validation

from lifelines.statistics import logrank_test

from imblearn.over_sampling import RandomOverSampler

import matplotlib as mpl

import matplotlib.pyplot as plt

import matplotlib.font_manager as font_manager

import numpy as np

import pandas as pd

from PIL import Image

import pylab

import scipy

import seaborn as sns

from sklearn import preprocessing

from sklearn.model_selection import train_test_split, KFold

from sklearn.metrics import average_precision_score, auc, f1_score, roc_curve, roc_auc_score

from sklearn.preprocessing import LabelBinarizer

from scipy import interp

mpl.rcParams['axes.linewidth'] = 3 #set the value globally

# Torch

import torch

import torch.nn as nn

from torch.nn import init, Parameter

from torch.utils.data._utils.collate import *

from torch.utils.data.dataloader import default_collate

import torch_geometric

from torch_geometric.data import Batch

################

# Regularization

################

def regularize_weights(model, reg_type=None):

    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)

    return l1_reg

def regularize_path_weights(model, reg_type=None):

    l1_reg = None

    for W in model.module.classifier.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)

    for W in model.module.linear.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)

    return l1_reg

def regularize_MM_weights(model, reg_type=None):

    l1_reg = None

    if model.module.__hasattr__('omic_net'):

        for W in model.module.omic_net.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)

    if model.module.__hasattr__('linear_h_path'):

        for W in model.module.linear_h_path.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)

    if model.module.__hasattr__('linear_h_omic'):

        for W in model.module.linear_h_omic.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)

    if model.module.__hasattr__('linear_h_grph'):

        for W in model.module.linear_h_grph.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)

    if model.module.__hasattr__('linear_z_path'):

        for W in model.module.linear_z_path.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)

    if model.module.__hasattr__('linear_z_omic'):

        for W in model.module.linear_z_omic.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)

    if model.module.__hasattr__('linear_z_grph'):

        for W in model.module.linear_z_grph.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)

    if model.module.__hasattr__('linear_o_path'):

        for W in model.module.linear_o_path.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)

    if model.module.__hasattr__('linear_o_omic'):

        for W in model.module.linear_o_omic.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)

    if model.module.__hasattr__('linear_o_grph'):

        for W in model.module.linear_o_grph.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)

    if model.module.__hasattr__('encoder1'):

        for W in model.module.encoder1.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)

    if model.module.__hasattr__('encoder2'):

        for W in model.module.encoder2.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)

    if model.module.__hasattr__('classifier'):

        for W in model.module.classifier.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)

    return l1_reg

def regularize_MM_omic(model, reg_type=None):

    l1_reg = None

    if model.module.__hasattr__('omic_net'):

        for W in model.module.omic_net.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)

    return l1_reg

################

# Network Initialization

################

def init_weights(net, init_type='orthogonal', init_gain=0.02):

    """Initialize network weights.

    Parameters:

        net (network)   -- network to be initialized

        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal

        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.

    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might

    work better for some applications. Feel free to try yourself.

    """

    def init_func(m):  # define the initialization function

        classname = m.__class__.__name__

        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):

            if init_type == 'normal':

                init.normal_(m.weight.data, 0.0, init_gain)

            elif init_type == 'xavier':

                init.xavier_normal_(m.weight.data, gain=init_gain)

            elif init_type == 'kaiming':

                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')

            elif init_type == 'orthogonal':

                init.orthogonal_(m.weight.data, gain=init_gain)

            else:

                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)

            if hasattr(m, 'bias') and m.bias is not None:

                init.constant_(m.bias.data, 0.0)

        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.

            init.normal_(m.weight.data, 1.0, init_gain)

            init.constant_(m.bias.data, 0.0)

    print('initialize network with %s' % init_type)

    net.apply(init_func)  # apply the initialization function <init_func>

def init_max_weights(module):

    for m in module.modules():

        if type(m) == nn.Linear:

            stdv = 1. / math.sqrt(m.weight.size(1))

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

            m.bias.data.zero_()

def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):

    """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights

    Parameters:

        net (network)      -- the network to be initialized

        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal

        gain (float)       -- scaling factor for normal, xavier and orthogonal.

        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2

    Return an initialized network.

    """

    if len(gpu_ids) > 0:

        assert(torch.cuda.is_available())

        net.to(gpu_ids[0])

        net = torch.nn.DataParallel(net, gpu_ids)           # multi-GPUs

    if init_type != 'max' and init_type != 'none':

        print("Init Type:", init_type)

        init_weights(net, init_type, init_gain=init_gain)

    elif init_type == 'none':

        print("Init Type: Not initializing networks.")

    elif init_type == 'max':

        print("Init Type: Self-Normalizing Weights")

    return net

################

# Freeze / Unfreeze

################

def unfreeze_unimodal(opt, model, epoch):

    if opt.mode == 'graphomic':

        if epoch == 5:

            dfs_unfreeze(model.module.omic_net)

            print("Unfreezing Omic")

        if epoch == 5:

            dfs_unfreeze(model.module.grph_net)

            print("Unfreezing Graph")

    elif opt.mode == 'pathomic':

        if epoch == 5:

            dfs_unfreeze(model.module.omic_net)

            print("Unfreezing Omic")

    elif opt.mode == 'pathgraph':

        if epoch == 5:

            dfs_unfreeze(model.module.grph_net)

            print("Unfreezing Graph")

    elif opt.mode == "pathgraphomic":

        if epoch == 5:

            dfs_unfreeze(model.module.omic_net)

            print("Unfreezing Omic")

        if epoch == 5:

            dfs_unfreeze(model.module.grph_net)

            print("Unfreezing Graph")

    elif opt.mode == "omicomic":

        if epoch == 5:

            dfs_unfreeze(model.module.omic_net)

            print("Unfreezing Omic")

    elif opt.mode == "graphgraph":

        if epoch == 5:

            dfs_unfreeze(model.module.grph_net)

            print("Unfreezing Graph")

def dfs_freeze(model):

    for name, child in model.named_children():

        for param in child.parameters():

            param.requires_grad = False

        dfs_freeze(child)

def dfs_unfreeze(model):

    for name, child in model.named_children():

        for param in child.parameters():

            param.requires_grad = True

        dfs_unfreeze(child)

def print_if_frozen(module):

    for idx, child in enumerate(module.children()):

        for param in child.parameters():

            if param.requires_grad == True:

                print("Learnable!!! %d:" % idx, child)

            else:

                print("Still Frozen %d:" % idx, child)

def unfreeze_vgg_features(model, epoch):

    epoch_schedule = {30:45}

    unfreeze_index = epoch_schedule[epoch]

    for idx, child in enumerate(model.features.children()):

        if idx > unfreeze_index:

            print("Unfreezing %d:" %idx, child)

            for param in child.parameters(): 

                param.requires_grad = True

        else:

            print("Still Frozen %d:" %idx, child)

            continue

################

# Collate Utils

################

def mixed_collate(batch):

    elem = batch[0]

    elem_type = type(elem)    

    transposed = zip(*batch)

    return [Batch.from_data_list(samples, []) if type(samples[0]) is torch_geometric.data.data.Data else default_collate(samples) for samples in transposed]

################

# Survival Utils

################

def CoxLoss(survtime, censor, hazard_pred, device):

    # 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_mat = 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_mat[i,j] = survtime[j] >= survtime[i]

    R_mat = torch.FloatTensor(R_mat).to(device)

    theta = hazard_pred.reshape(-1)

    exp_theta = torch.exp(theta)

    loss_cox = -torch.mean((theta - torch.log(torch.sum(exp_theta*R_mat, dim=1))) * censor)

    return loss_cox

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(hazardsdata, labels):

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

    median = np.median(hazardsdata)

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

    hazards_dichotomize[hazardsdata > median] = 1

    correct = np.sum(hazards_dichotomize == labels)

    return correct / len(labels)

def cox_log_rank(hazardsdata, labels, survtime_all):

    median = np.median(hazardsdata)

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

    hazards_dichotomize[hazardsdata > median] = 1

    idx = hazards_dichotomize == 0

    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):

    concord = 0.

    total = 0.

    N_test = labels.shape[0]

    for i in range(N_test):

        if labels[i] == 1:

            for j in range(N_test):

                if survtime_all[j] > survtime_all[i]:

                    total += 1

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

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

    return(concord/total)

def CIndex_lifeline(hazards, labels, survtime_all):

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

################

# Data Utils

################

def addHistomolecularSubtype(data):

    """

    Molecular Subtype: IDHwt == 0, IDHmut-non-codel == 1, IDHmut-codel == 2

    Histology Subtype: astrocytoma == 0, oligoastrocytoma == 1, oligodendroglioma == 2, glioblastoma == 3

    """

    subtyped_data = data.copy()

    subtyped_data.insert(loc=0, column='Histomolecular subtype', value=np.ones(len(data)))

    idhwt_ATC = np.logical_and(data['Molecular subtype'] == 0, np.logical_or(data['Histology'] == 0, data['Histology'] == 3))

    subtyped_data.loc[idhwt_ATC, 'Histomolecular subtype'] = 'idhwt_ATC'

    idhmut_ATC = np.logical_and(data['Molecular subtype'] == 1, np.logical_or(data['Histology'] == 0, data['Histology'] == 3))

    subtyped_data.loc[idhmut_ATC, 'Histomolecular subtype'] = 'idhmut_ATC'

    ODG = np.logical_and(data['Molecular subtype'] == 2, data['Histology'] == 2)

    subtyped_data.loc[ODG, 'Histomolecular subtype'] = 'ODG'

    return subtyped_data

def changeHistomolecularSubtype(data):

    """

    Molecular Subtype: IDHwt == 0, IDHmut-non-codel == 1, IDHmut-codel == 2

    Histology Subtype: astrocytoma == 0, oligoastrocytoma == 1, oligodendroglioma == 2, glioblastoma == 3

    """

    data = data.drop(['Histomolecular subtype'], axis=1)

    subtyped_data = data.copy()

    subtyped_data.insert(loc=0, column='Histomolecular subtype', value=np.ones(len(data)))

    idhwt_ATC = np.logical_and(data['Molecular subtype'] == 0, np.logical_or(data['Histology'] == 0, data['Histology'] == 3))

    subtyped_data.loc[idhwt_ATC, 'Histomolecular subtype'] = 'idhwt_ATC'

    idhmut_ATC = np.logical_and(data['Molecular subtype'] == 1, np.logical_or(data['Histology'] == 0, data['Histology'] == 3))

    subtyped_data.loc[idhmut_ATC, 'Histomolecular subtype'] = 'idhmut_ATC'

    ODG = np.logical_and(data['Molecular subtype'] == 2, data['Histology'] == 2)

    subtyped_data.loc[ODG, 'Histomolecular subtype'] = 'ODG'

    return subtyped_data

def getCleanAllDataset(dataroot='./data/TCGA_GBMLGG/', ignore_missing_moltype=False, ignore_missing_histype=False, use_rnaseq=False):

    ### 1. Joining all_datasets.csv with grade data. Looks at columns with misisng samples

    metadata = ['Histology', 'Grade', 'Molecular subtype', 'TCGA ID', 'censored', 'Survival months']

    all_dataset = pd.read_csv(os.path.join(dataroot, 'all_dataset.csv')).drop('indexes', axis=1)

    all_dataset.index = all_dataset['TCGA ID']

    all_grade = pd.read_csv(os.path.join(dataroot, 'grade_data.csv'))

    all_grade['Histology'] = all_grade['Histology'].str.replace('astrocytoma (glioblastoma)', 'glioblastoma', regex=False)

    all_grade.index = all_grade['TCGA ID']

    assert pd.Series(all_dataset.index).equals(pd.Series(sorted(all_grade.index)))

    all_dataset = all_dataset.join(all_grade[['Histology', 'Grade', 'Molecular subtype']], how='inner')

    cols = all_dataset.columns.tolist()

    cols = cols[-3:] + cols[:-3]

    all_dataset = all_dataset[cols]

    if use_rnaseq:

        gbm = pd.read_csv(os.path.join(dataroot, 'mRNA_Expression_z-Scores_RNA_Seq_RSEM.txt'), sep='\t', skiprows=1, index_col=0)

        lgg = pd.read_csv(os.path.join(dataroot, 'mRNA_Expression_Zscores_RSEM.txt'), sep='\t', skiprows=1, index_col=0)

        gbm = gbm[gbm.columns[~gbm.isnull().all()]]

        lgg = lgg[lgg.columns[~lgg.isnull().all()]]

        glioma_RNAseq = gbm.join(lgg, how='inner').T

        glioma_RNAseq = glioma_RNAseq.dropna(axis=1)

        glioma_RNAseq.columns = [gene+'_rnaseq' for gene in glioma_RNAseq.columns]

        glioma_RNAseq.index = [patname[:12] for patname in glioma_RNAseq.index]

        glioma_RNAseq = glioma_RNAseq.iloc[~glioma_RNAseq.index.duplicated()]

        glioma_RNAseq.index.name = 'TCGA ID'

        all_dataset = all_dataset.join(glioma_RNAseq, how='inner')

    pat_missing_moltype = all_dataset[all_dataset['Molecular subtype'].isna()].index

    pat_missing_idh = all_dataset[all_dataset['idh mutation'].isna()].index

    pat_missing_1p19q = all_dataset[all_dataset['codeletion'].isna()].index

    print("# Missing Molecular Subtype:", len(pat_missing_moltype))

    print("# Missing IDH Mutation:", len(pat_missing_idh))

    print("# Missing 1p19q Codeletion:", len(pat_missing_1p19q))

    assert pat_missing_moltype.equals(pat_missing_idh)

    assert pat_missing_moltype.equals(pat_missing_1p19q)

    pat_missing_grade =  all_dataset[all_dataset['Grade'].isna()].index

    pat_missing_histype = all_dataset[all_dataset['Histology'].isna()].index

    print("# Missing Histological Subtype:", len(pat_missing_histype))

    print("# Missing Grade:", len(pat_missing_grade))

    assert pat_missing_histype.equals(pat_missing_grade)

    ### 2. Impute Missing Genomic Data: Removes patients with missing molecular subtype / idh mutation / 1p19q. Else imputes with median value of each column. Fills missing Molecular subtype with "Missing"

    if ignore_missing_moltype: 

        all_dataset = all_dataset[all_dataset['Molecular subtype'].isna() == False]

    for col in all_dataset.drop(metadata, axis=1).columns:

        all_dataset['Molecular subtype'] = all_dataset['Molecular subtype'].fillna('Missing')

        all_dataset[col] = all_dataset[col].fillna(all_dataset[col].median())

    ### 3. Impute Missing Histological Data: Removes patients with missing histological subtype / grade. Else imputes with "missing" / grade -1

    if ignore_missing_histype: 

        all_dataset = all_dataset[all_dataset['Histology'].isna() == False]

    else:

        all_dataset['Grade'] = all_dataset['Grade'].fillna(1)

        all_dataset['Histology'] = all_dataset['Histology'].fillna('Missing')

    all_dataset['Grade'] = all_dataset['Grade'] - 2

    ### 4. Adds Histomolecular subtype

    ms2int = {'Missing':-1, 'IDHwt':0, 'IDHmut-non-codel':1, 'IDHmut-codel':2}

    all_dataset[['Molecular subtype']] = all_dataset[['Molecular subtype']].applymap(lambda s: ms2int.get(s) if s in ms2int else s)

    hs2int = {'Missing':-1, 'astrocytoma':0, 'oligoastrocytoma':1, 'oligodendroglioma':2, 'glioblastoma':3}

    all_dataset[['Histology']] = all_dataset[['Histology']].applymap(lambda s: hs2int.get(s) if s in hs2int else s)

    all_dataset = addHistomolecularSubtype(all_dataset)

    metadata.extend(['Histomolecular subtype'])

    all_dataset['censored'] = 1 - all_dataset['censored']

    return metadata, all_dataset

################

# Analysis Utils

################

def count_parameters(model):

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

def hazard2grade(hazard, p):

    if hazard < p[0]:

        return 0

    elif hazard < p[1]:

        return 1

    return 2

def p(n):

    def percentile_(x):

        return np.percentile(x, n)

    percentile_.__name__ = 'p%s' % n

    return percentile_

def natural_sort(l): 

    convert = lambda text: int(text) if text.isdigit() else text.lower() 

    alphanum_key = lambda key: [ convert(c) for c in re.split('([0-9]+)', key) ] 

    return sorted(l, key = alphanum_key)

def CI_pm(data, confidence=0.95):

    a = 1.0 * np.array(data)

    n = len(a)

    m, se = np.mean(a), scipy.stats.sem(a)

    h = se * scipy.stats.t.ppf((1 + confidence) / 2., n-1)

    return str("{0:.4f} ± ".format(m) + "{0:.3f}".format(h))

def CI_interval(data, confidence=0.95):

    a = 1.0 * np.array(data)

    n = len(a)

    m, se = np.mean(a), scipy.stats.sem(a)

    h = se * scipy.stats.t.ppf((1 + confidence) / 2., n-1)

    return str("{0:.3f}, ".format(m-h) + "{0:.3f}".format(m+h))

def poolSurvTestPD(ckpt_name='./checkpoints/TCGA_GBMLGG/surv_15_rnaseq/', model='pathgraphomic_fusion', split='test', zscore=False, agg_type='Hazard_mean'):

    all_dataset_regstrd_pooled = []    

    ignore_missing_moltype = 1 if 'omic' in model else 0

    ignore_missing_histype = 1 if 'grad' in ckpt_name else 0

    use_patch, roi_dir, use_vgg_features = ('_patch_', 'all_st_patches_512', 1) if ((('path' in model) or ('graph' in model)) and ('cox' not in model)) else ('_', 'all_st', 0)

    use_rnaseq = '_rnaseq' if ('rnaseq' in ckpt_name and 'path' != model and 'pathpath' not in model and 'graph' != model and 'graphgraph' not in model) else ''

    for k in range(1,16):

        pred = pickle.load(open(ckpt_name+'/%s/%s_%d%spred_%s.pkl' % (model, model, k, use_patch, split), 'rb'))    

        if 'cox' not in model:

            surv_all = pd.DataFrame(np.stack(np.delete(np.array(pred), 3))).T

            surv_all.columns = ['Hazard', 'Survival months', 'censored', 'Grade']

            data_cv = pickle.load(open('./data/TCGA_GBMLGG/splits/gbmlgg15cv_%s_%d_%d_%d%s.pkl' % (roi_dir, ignore_missing_moltype, ignore_missing_histype, use_vgg_features, use_rnaseq), 'rb'))

            data_cv_splits = data_cv['cv_splits']

            data_cv_split_k = data_cv_splits[k]

            assert np.all(data_cv_split_k[split]['t'] == pred[1]) # Data is correctly registered

            all_dataset = data_cv['data_pd'].drop('TCGA ID', axis=1)

            all_dataset_regstrd = all_dataset.loc[data_cv_split_k[split]['x_patname']] # Subset of "all_datasets" (metadata) that is registered with "pred" (predictions)

            assert np.all(np.array(all_dataset_regstrd['Survival months']) == pred[1])

            assert np.all(np.array(all_dataset_regstrd['censored']) == pred[2])

            assert np.all(np.array(all_dataset_regstrd['Grade']) == pred[4])

            all_dataset_regstrd.insert(loc=0, column='Hazard', value = np.array(surv_all['Hazard']))

            all_dataset_regstrd.index.name = 'TCGA ID'

            hazard_agg = all_dataset_regstrd.groupby('TCGA ID').agg({'Hazard': ['mean', 'median', max, p(0.25), p(0.75)]})

            hazard_agg.columns = ["_".join(x) for x in hazard_agg.columns.ravel()]

            hazard_agg = hazard_agg[[agg_type]]

            hazard_agg.columns = ['Hazard']

            pred = hazard_agg.join(all_dataset, how='inner')

        if zscore: pred['Hazard'] = scipy.stats.zscore(np.array(pred['Hazard']))

        all_dataset_regstrd_pooled.append(pred)

    all_dataset_regstrd_pooled = pd.concat(all_dataset_regstrd_pooled)

    all_dataset_regstrd_pooled = changeHistomolecularSubtype(all_dataset_regstrd_pooled)

    return all_dataset_regstrd_pooled

def getAggHazardCV(ckpt_name='./checkpoints/TCGA_GBMLGG/surv_15_rnaseq/', model='pathgraphomic_fusion', split='test', agg_type='Hazard_mean'):

    result = []

    ignore_missing_moltype = 1 if 'omic' in model else 0

    ignore_missing_histype = 1 if 'grad' in ckpt_name else 0

    use_patch, roi_dir, use_vgg_features = ('_patch_', 'all_st_patches_512', 1) if (('path' in model) or ('graph' in model)) else ('_', 'all_st', 0)

    use_rnaseq = '_rnaseq' if ('rnaseq' in ckpt_name and 'path' != model and 'pathpath' not in model and 'graph' != model and 'graphgraph' not in model) else ''

    for k in range(1,16):

        pred = pickle.load(open(ckpt_name+'/%s/%s_%d%spred_%s.pkl' % (model, model, k, use_patch, split), 'rb'))    

        surv_all = pd.DataFrame(np.stack(np.delete(np.array(pred), 3))).T

        surv_all.columns = ['Hazard', 'Survival months', 'censored', 'Grade']

        data_cv = pickle.load(open('./data/TCGA_GBMLGG/splits/gbmlgg15cv_%s_%d_%d_%d%s.pkl' % (roi_dir, ignore_missing_moltype, ignore_missing_histype, use_vgg_features, use_rnaseq), 'rb'))

        data_cv_splits = data_cv['cv_splits']

        data_cv_split_k = data_cv_splits[k]

        assert np.all(data_cv_split_k[split]['t'] == pred[1]) # Data is correctly registered

        all_dataset = data_cv['data_pd'].drop('TCGA ID', axis=1)

        all_dataset_regstrd = all_dataset.loc[data_cv_split_k[split]['x_patname']] # Subset of "all_datasets" (metadata) that is registered with "pred" (predictions)

        assert np.all(np.array(all_dataset_regstrd['Survival months']) == pred[1])

        assert np.all(np.array(all_dataset_regstrd['censored']) == pred[2])

        assert np.all(np.array(all_dataset_regstrd['Grade']) == pred[4])

        all_dataset_regstrd.insert(loc=0, column='Hazard', value = np.array(surv_all['Hazard']))

        all_dataset_regstrd.index.name = 'TCGA ID'

        hazard_agg = all_dataset_regstrd.groupby('TCGA ID').agg({'Hazard': ['mean', max, p(0.75)]})

        hazard_agg.columns = ["_".join(x) for x in hazard_agg.columns.ravel()]

        hazard_agg = hazard_agg[[agg_type]]

        hazard_agg.columns = ['Hazard']

        all_dataset_hazard = hazard_agg.join(all_dataset, how='inner')

        cin = CIndex_lifeline(all_dataset_hazard['Hazard'], all_dataset_hazard['censored'], all_dataset_hazard['Survival months'])

        result.append(cin)

    return result

def calcGradMetrics(ckpt_name='./checkpoints/grad_15/', model='pathgraphomic_fusion', split='test', avg='micro'):

    auc_all = []

    ap_all = []

    f1_all = []

    f1_gradeIV_all = []

    ignore_missing_moltype = 1 if 'omic' in model else 0

    ignore_missing_histype = 1 if 'grad' in ckpt_name else 0

    use_patch, roi_dir, use_vgg_features = ('_patch_', 'all_st_patches_512', 1) if (('path' in model) or ('graph' in model)) else ('_', 'all_st', 0)

    for k in range(1,16):

        pred = pickle.load(open(ckpt_name+'/%s/%s_%d%spred_%s.pkl' % (model, model, k, use_patch, split), 'rb'))    

        grade_pred, grade = np.array(pred[3]), np.array(pred[4])

        enc = LabelBinarizer()

        enc.fit(grade)

        grade_oh = enc.transform(grade)

        rocauc = roc_auc_score(grade_oh, grade_pred, avg)

        ap = average_precision_score(grade_oh, grade_pred, average=avg)

        f1 = f1_score(grade_pred.argmax(axis=1), grade, average=avg)

        f1_gradeIV = f1_score(grade_pred.argmax(axis=1), grade, average=None)[2]

        auc_all.append(rocauc)

        ap_all.append(ap)

        f1_all.append(f1)

        f1_gradeIV_all.append(f1_gradeIV)

    return np.array([CI_pm(auc_all), CI_pm(ap_all), CI_pm(f1_all), CI_pm(f1_gradeIV_all)])

################

# Plot Utils

################

def makeKaplanMeierPlot(ckpt_name='./checkpoints/surv_15_rnaseq/', model='omic', split='test', zscore=False, agg_type='Hazard_mean'):

    def hazard2KMCurve(data, subtype):

        p = np.percentile(data['Hazard'], [33, 66])

        if p[0] == p[1]: p[0] = 2.99997

        data.insert(0, 'grade_pred', [hazard2grade(hazard, p) for hazard in data['Hazard']])

        kmf_pred = lifelines.KaplanMeierFitter()

        kmf_gt = lifelines.KaplanMeierFitter()

        def get_name(model):

            mode2name = {'pathgraphomic':'Pathomic F.', 'pathomic':'Pathomic F.', 'graphomic':'Pathomic F.', 'path':'Histology CNN', 'graph':'Histology GCN', 'omic':'Genomic SNN'}

            for mode in mode2name.keys():

                if mode in model: return mode2name[mode]

            return 'N/A'

        fig = plt.figure(figsize=(10, 10), dpi=600)

        ax = plt.subplot()

        censor_style = {'ms': 20, 'marker': '+'}

        temp = data[data['Grade']==0]

        kmf_gt.fit(temp['Survival months']/365, temp['censored'], label="Grade II")

        kmf_gt.plot(ax=ax, show_censors=True, ci_show=False, c='g', linewidth=3, ls='--', markerfacecolor='black', censor_styles=censor_style)

        temp = data[data['grade_pred']==0]

        kmf_pred.fit(temp['Survival months']/365, temp['censored'], label="%s (Low)" % get_name(model))

        kmf_pred.plot(ax=ax, show_censors=True, ci_show=False, c='g', linewidth=4, ls='-', markerfacecolor='black', censor_styles=censor_style)

        temp = data[data['Grade']==1]

        kmf_gt.fit(temp['Survival months']/365, temp['censored'], label="Grade III")

        kmf_gt.plot(ax=ax, show_censors=True, ci_show=False, c='b', linewidth=3, ls='--', censor_styles=censor_style)

        temp = data[data['grade_pred']==1]

        kmf_pred.fit(temp['Survival months']/365, temp['censored'], label="%s (Mid)" % get_name(model))

        kmf_pred.plot(ax=ax, show_censors=True, ci_show=False, c='b', linewidth=4, ls='-', censor_styles=censor_style)

        if subtype != 'ODG':    

            temp = data[data['Grade']==2]

            kmf_gt.fit(temp['Survival months']/365, temp['censored'], label="Grade IV")

            kmf_gt.plot(ax=ax, show_censors=True, ci_show=False, c='r', linewidth=3, ls='--', censor_styles=censor_style)

            temp = data[data['grade_pred']==2]

            kmf_pred.fit(temp['Survival months']/365, temp['censored'], label="%s (High)" % get_name(model))

            kmf_pred.plot(ax=ax, show_censors=True, ci_show=False, c='r', linewidth=4, ls='-', censor_styles=censor_style)

        ax.set_xlabel('')

        ax.set_ylim(0, 1)

        ax.set_yticks(np.arange(0, 1.001, 0.5))

        ax.tick_params(axis='both', which='major', labelsize=40)    

        plt.legend(fontsize=32, prop=font_manager.FontProperties(family='Arial', style='normal', size=32))

        if subtype != 'idhwt_ATC': ax.get_legend().remove()

        return fig

    data = poolSurvTestPD(ckpt_name, model, split, zscore, agg_type)

    for subtype in ['idhwt_ATC', 'idhmut_ATC', 'ODG']:

        fig = hazard2KMCurve(data[data['Histomolecular subtype'] == subtype], subtype)

        fig.savefig(ckpt_name+'/%s_KM_%s.png' % (model, subtype))

    fig = hazard2KMCurve(data, 'all')

    fig.savefig(ckpt_name+'/%s_KM_%s.png' % (model, 'all'))

def makeHazardSwarmPlot(ckpt_name='./checkpoints/surv_15_rnaseq/', model='path', split='test', zscore=True, agg_type='Hazard_mean'):

    mpl.rcParams['font.family'] = "arial"

    data = poolSurvTestPD(ckpt_name=ckpt_name, model=model, split=split, zscore=zscore, agg_type=agg_type)

    data = data[data['Grade'] != -1]

    data = data[data['Histomolecular subtype'] != -1]

    data['Grade'] = data['Grade'].astype(int).astype(str)

    data['Grade'] = data['Grade'].str.replace('0', 'Grade II', regex=False)

    data['Grade'] = data['Grade'].str.replace('1', 'Grade III', regex=False)

    data['Grade'] = data['Grade'].str.replace('2', 'Grade IV', regex=False)

    data['Histomolecular subtype'] = data['Histomolecular subtype'].str.replace('idhwt_ATC', 'IDH-wt \n astryocytoma', regex=False)

    data['Histomolecular subtype'] = data['Histomolecular subtype'].str.replace('idhmut_ATC', 'IDH-mut \n astrocytoma', regex=False)

    data['Histomolecular subtype'] = data['Histomolecular subtype'].str.replace('ODG', 'Oligodendroglioma', regex=False)

    fig, ax = plt.subplots(dpi=600)

    ax.set_ylim([-2, 2.5]) # plt.ylim(-2, 2)

    ax.spines['right'].set_visible(False)

    ax.spines['top'].set_visible(False)

    ax.set_yticks(np.arange(-2, 2.001, 1))

    sns.swarmplot(x = 'Histomolecular subtype', y='Hazard', data=data, hue='Grade',

                  palette={"Grade II":"#AFD275" , "Grade III":"#7395AE", "Grade IV":"#E7717D"}, 

                  size = 4, alpha = 0.9, ax=ax)

    ax.set_xlabel('') # ax.set_xlabel('Histomolecular subtype', size=16)

    ax.set_ylabel('') # ax.set_ylabel('Hazard (Z-Score)', size=16)

    ax.tick_params(axis='y', which='both', labelsize=20)

    ax.tick_params(axis='x', which='both', labelsize=15)

    ax.tick_params(axis='x', which='both', labelbottom='off') # doesn't work??

    ax.legend(prop={'size': 8})

    fig.savefig(ckpt_name+'/%s_HSP.png' % (model))

def makeHazardBoxPlot(ckpt_name='./checkpoints/surv_15_rnaseq/', model='omic', split='test', zscore=True, agg_type='Hazard_mean'):

    mpl.rcParams['font.family'] = "arial"

    data = poolSurvTestPD(ckpt_name, model, split, zscore, 'Hazard_mean')

    data['Grade'] = data['Grade'].astype(int).astype(str)

    data['Grade'] = data['Grade'].str.replace('0', 'II', regex=False)

    data['Grade'] = data['Grade'].str.replace('1', 'III', regex=False)

    data['Grade'] = data['Grade'].str.replace('2', 'IV', regex=False)

    fig, axes = plt.subplots(nrows=1, ncols=3, gridspec_kw={'width_ratios': [3, 3, 2]}, dpi=600)

    plt.subplots_adjust(wspace=0, hspace=0)

    plt.ylim(-2, 2)

    plt.yticks(np.arange(-2, 2.001, 1))

    #color_dict = {0: '#CF9498', 1: '#8CC7C8', 2: '#AAA0C6'}

    #color_dict = {0: '#F76C6C', 1: '#A8D0E6', 2: '#F8E9A1'}

    color_dict = ['#F76C6C', '#A8D0E6', '#F8E9A1']

    subtypes = ['idhwt_ATC', 'idhmut_ATC', 'ODG']

    for i in range(len(subtypes)):

        axes[i].spines["top"].set_visible(False)

        axes[i].spines["right"].set_visible(False)

        axes[i].xaxis.grid(False)

        axes[i].yaxis.grid(False)

        if i > 0: 

            axes[i].get_yaxis().set_visible(False)

            axes[i].spines["left"].set_visible(False)

        order = ["II","III","IV"] if subtypes[i] != 'ODG' else ["II", "III"]

        axes[i].xaxis.label.set_visible(False)

        axes[i].yaxis.label.set_visible(False)

        axes[i].tick_params(axis='y', which='both', labelsize=20)

        axes[i].tick_params(axis='x', which='both', labelsize=15)

        datapoints = data[data['Histomolecular subtype'] == subtypes[i]]

        sns.boxplot(y='Hazard', x="Grade", data=datapoints, ax = axes[i], color=color_dict[i], order=order)

        sns.stripplot(y='Hazard', x='Grade', data=datapoints, alpha=0.2, jitter=0.2, color='k', ax = axes[i], order=order)

        axes[i].set_ylim(-2.5, 2.5)

        axes[i].set_yticks(np.arange(-2.0, 2.1, 1))

    #axes[2].legend(prop={'size': 10})

    fig.savefig(ckpt_name+'/%s_HBP.png' % (model))

def makeAUROCPlot(ckpt_name='./checkpoints/grad_15/', model_list=['path', 'omic', 'pathgraphomic_fusion'], split='test', avg='micro', use_zoom=False):

    mpl.rcParams['font.family'] = "arial"

    colors = {'path':'dodgerblue', 'graph':'orange', 'omic':'green', 'pathgraphomic_fusion':'crimson'}

    names = {'path':'Histology CNN', 'graph':'Histology GCN', 'omic':'Genomic SNN', 'pathgraphomic_fusion':'Pathomic F.'}

    zoom_params = {0:([0.2, 0.4], [0.8, 1.0]), 

                   1:([0.25, 0.45], [0.75, 0.95]),

                   2:([0.0, 0.2], [0.8, 1.0]),

                   'micro':([0.15, 0.35], [0.8, 1.0])}

    mean_fpr = np.linspace(0, 1, 100)

    classes = [0, 1, 2, avg]

    ### 1. Looping over classes

    for i in classes:

        print("Class: " + str(i))

        fi = pylab.figure(figsize=(10,10), dpi=600, linewidth=0.2)

        axi = plt.subplot()

        ### 2. Looping over models

        for m, model in enumerate(model_list):

            ignore_missing_moltype = 1 if 'omic' in model else 0

            ignore_missing_histype = 1 if 'grad' in ckpt_name else 0

            use_patch, roi_dir, use_vgg_features = ('_patch_', 'all_st_patches_512', 1) if (('path' in model) or ('graph' in model)) else ('_', 'all_st', 0)

            ###. 3. Looping over all splits

            tprs, pres, aucrocs, rocaucs, = [], [], [], []

            for k in range(1,16):

                pred = pickle.load(open(ckpt_name+'/%s/%s_%d%spred_%s.pkl' % (model, model, k, use_patch, split), 'rb'))    

                grade_pred, grade = np.array(pred[3]), np.array(pred[4])

                enc = LabelBinarizer()

                enc.fit(grade)

                grade_oh = enc.transform(grade)

                if i != avg:

                    pres.append(average_precision_score(grade_oh[:, i], grade_pred[:, i])) # from https://scikit-learn.org/stable/auto_examples/model_selection/plot_precision_recall.html

                    fpr, tpr, thresh = roc_curve(grade_oh[:,i], grade_pred[:,i], drop_intermediate=False)

                    aucrocs.append(auc(fpr, tpr)) # https://scikit-learn.org/stable/auto_examples/model_selection/plot_roc.html

                    rocaucs.append(roc_auc_score(grade_oh[:,i], grade_pred[:,i])) # https://scikit-learn.org/stable/modules/generated/sklearn.metrics.roc_auc_score.html#sklearn.metrics.roc_auc_score

                    tprs.append(interp(mean_fpr, fpr, tpr))

                    tprs[-1][0] = 0.0

                else:

                    # A "micro-average": quantifying score on all classes jointly

                    pres.append(average_precision_score(grade_oh, grade_pred, average=avg))

                    fpr, tpr, thresh = roc_curve(grade_oh.ravel(), grade_pred.ravel())

                    aucrocs.append(auc(fpr, tpr))

                    rocaucs.append(roc_auc_score(grade_oh, grade_pred, avg))

                    tprs.append(interp(mean_fpr, fpr, tpr))

                    tprs[-1][0] = 0.0

            mean_tpr = np.mean(tprs, axis=0)

            mean_tpr[-1] = 1.0

            #mean_auc = auc(mean_fpr, mean_tpr)

            mean_auc = np.mean(aucrocs)

            std_auc = np.std(aucrocs)

            print('\t'+'%s - AUC: %0.3f ± %0.3f' % (model, mean_auc, std_auc))

            if use_zoom:

                alpha, lw = (0.8, 6) if model =='pathgraphomic_fusion' else (0.5, 6)

                plt.plot(mean_fpr, mean_tpr, color=colors[model],

                     label=r'%s (AUC = %0.3f $\pm$ %0.3f)' % (names[model], mean_auc, std_auc), lw=lw, alpha=alpha)

                std_tpr = np.std(tprs, axis=0)

                tprs_upper = np.minimum(mean_tpr + std_tpr, 1)

                tprs_lower = np.maximum(mean_tpr - std_tpr, 0)

                plt.fill_between(mean_fpr, tprs_lower, tprs_upper, color=colors[model], alpha=0.1)

                plt.xlim([zoom_params[i][0][0]-0.005, zoom_params[i][0][1]+0.005])

                plt.ylim([zoom_params[i][1][0]-0.005, zoom_params[i][1][1]+0.005])

                axi.set_xticks(np.arange(zoom_params[i][0][0], zoom_params[i][0][1]+0.001, 0.05))

                axi.set_yticks(np.arange(zoom_params[i][1][0], zoom_params[i][1][1]+0.001, 0.05))

                axi.tick_params(axis='both', which='major', labelsize=26)

            else:

                alpha, lw = (0.8, 4) if model =='pathgraphomic_fusion' else (0.5, 3)

                plt.plot(mean_fpr, mean_tpr, color=colors[model],

                     label=r'%s (AUC = %0.3f $\pm$ %0.3f)' % (names[model], mean_auc, std_auc), lw=lw, alpha=alpha)

                std_tpr = np.std(tprs, axis=0)

                tprs_upper = np.minimum(mean_tpr + std_tpr, 1)

                tprs_lower = np.maximum(mean_tpr - std_tpr, 0)

                plt.fill_between(mean_fpr, tprs_lower, tprs_upper, color=colors[model], alpha=0.1)

                plt.xlim([-0.05, 1.05])

                plt.ylim([-0.05, 1.05])

                axi.set_xticks(np.arange(0, 1.001, 0.2))

                axi.set_yticks(np.arange(0, 1.001, 0.2))

                axi.legend(loc="lower right", prop={'size': 20})

                axi.tick_params(axis='both', which='major', labelsize=30)

                #plt.plot([0, 1], [0, 1], linestyle='--', lw=2, color='navy', alpha=.8)

    figures = [manager.canvas.figure

               for manager in mpl._pylab_helpers.Gcf.get_all_fig_managers()]

    zoom = '_zoom' if use_zoom else ''

    for i, fig in enumerate(figures):

        fig.savefig(ckpt_name+'/AUC_%s%s.png' % (classes[i], zoom))