"""

A collection of functions used to analyze radiomic features. This code appends data from RT, primary tumor

recurrence, and survival times. The example code is at the end of this file.

"""

import os

import sys

file = '/home/matt/Documents/SegSarcoma'

sys.path.append(file)

import pptx

import pickle

import json

from datetime import datetime, timedelta

import datetime as dt

from time import time

from subprocess import call

from lifelines import KaplanMeierFitter

import numpy as np

import pandas as pd

from matplotlib import pyplot as plt

from matplotlib import rcParams

import seaborn as sns

from glob2 import glob

from scipy import stats

from Crawler.crawler_radiomics import load_study_data

from Radiomics.radiomic_functions import *

from PIL import Image, ImageDraw, ImageFont

# Set up plotting properties

sns.set()

sns.set_style('whitegrid')

sns.set_context("paper")

PATH_TO_MRMR = '/home/matt/Downloads/mrmr_c_src/mrmr'

def concate_contrasts_small(df):

    """

    Takes a dataframe containing radiomcis data for all animals and contrasta.

    Returns a dataframe with contrasts of the same animal concatenated together.

    Args:

        df:

    Returns:

    """

    # Turn multiple contrast images into a single vector of data

    keys = df.keys()

    cat_cols = [i + '_T1' for i in keys] + \

               [i + '_T1C' for i in keys] + \

               [i + '_T2' for i in keys]

    key = [[i + '_T1' for i in keys],

           [i + '_T1C' for i in keys],

           [i + '_T2' for i in keys]]

    df_cat = pd.DataFrame(columns=cat_cols, index=range(1))

    ind = 0

    for i in range(3):

        tmp = df.iloc[i]

        # Append data into larger array

        df_cat[key[i]] = [n for n in tmp.get_values().ravel()]

    return df_cat

def strip_excluded(rfiles, exclude):

    """

    Remove animals which either did not have two scans or died shortly after surgery.

    Args:

        rfiles (list): list of all animals

        exclude (list): list of animals to exclude

    Returns:

        (list): filtered list

    """

    out_rfiles = rfiles

    for file in rfiles:

        path, _ = os.path.split(file)

        path, _ = os.path.split(path)

        _, animal = os.path.split(path)

        if animal in exclude:

            out_rfiles.remove(file)

    return out_rfiles

def keep_animals(rfiles, keep):

    """

    Only keep specific animals.

    Args:

        rfiles (list): list of all animals

        keep (list): list of animals to keep

    Returns:

        (list): filtered list

    """

    out_rfiles = []

    for file in rfiles:

        path, _ = os.path.split(file)

        path, _ = os.path.split(path)

        _, animal = os.path.split(path)

        if animal in keep:

            out_rfiles.append(file)

    return out_rfiles

def select_region(rfiles, region):

    """

    Select features based on the region they were calculated in.

    Args:

        rfiles (list): list of all animals

        region (str): 'tumor', 'bed', or 'edge'

    Returns:

    """

    out_rfiles = rfiles

    for _ in range(2):

        for file in rfiles:

            if region == 'tumor':

                if '_bed' in file:

                    out_rfiles.remove(file)

                elif 'edge' in file:

                    out_rfiles.remove(file)

            elif region == 'bed':

                if '_bed' not in file:

                    out_rfiles.remove(file)

            elif region == 'edge':

                if '_edge' not in file:

                    out_rfiles.remove(file)

        rfiles = out_rfiles

    return out_rfiles

def load_radiomics(radiomics_paths, exclude, region, keep, group):

    """

    Load radimic data (output of PyRadiomics)

    Args:

        radiomics_paths (str): path to radiomic files

        exclude (str): amimals to exclude

        region (str): which region to select

        keep (str): animals to keep

        group (str): pre/post RT and control/treatment

    Returns:

        pandas dataframe: radiomic data

    """

    # Load radiomics file data

    with open(radiomics_paths, 'r') as f:

        rfiles = f.readlines()

    # Remove line endings '\n'

    rfiles = [i.strip() for i in rfiles]

    # Sort radiomics files

    rfiles = strip_excluded(rfiles, exclude)

    rfiles = select_region(rfiles, region)

    # Only load specific animals

    rfiles = keep_animals(rfiles, keep)

    # Load data and combine contrasts

    df = load_and_concate(rfiles)

    df['Group'] = group

    # df['Region'] = region

    df.index = range(df.shape[0])

    return df

def load_study_logs(summary_file, log_file):

    """

    Load study logs which contain relevant dates

    Args:

        summary_file:

        log_file:

    Returns:

    """

    # Load summary file

    summary = load_study_data(summary_file)

    # Convert animal IDs to match the log file

    summary['Kirsch lab iD'] = ['K' + str(int(i)) for i in summary['Kirsch lab iD']]

    if not log_file == '':

        # Load log file

        with open(log_file) as f:

            log = json.load(f)

    else:

        log = 0

    # Amputation date

    treat_date = summary['Date of Antibody treatment   (0.2mg/mouse Ip)  1st']

    # Compute time until today

    summary['SinceAmputation'] = (datetime.now() - treat_date)

    return summary, log

def capitalize_method(key):

    """

    Capitalizes an list of strings

    Args:

        key (str): a feature name to capitalize

    Returns:

        (str): capitalized name

    """

    if str in key:

        key = key.replace(str, str.upper())

    return key

def load_and_concate(rfiles):

    """

    Loads and concatenates radiomic features from different MR contrasts

    Args:

        rfiles (list): list of radiomic files (one for each scan)

    Returns:

        pandas dataframe: radiomic features

    """

    animal_id = []

    for i, file in enumerate(rfiles):

        df = pd.DataFrame.from_csv(file)

        if i == 0:

            cat_df = concate_contrasts_small(df)

        else:

            df = concate_contrasts_small(df)

            cat_df = pd.concat([cat_df, df])

        path, _ = os.path.split(file)

        path, _ = os.path.split(path)

        _, animal = os.path.split(path)

        if animal not in animal_id:

            animal_id.append(animal)

        else:

            print(animal)

    # Filter out non-feature names

    features = cat_df.keys()

    feature_names = list(

        sorted(filter(lambda k: k.startswith("original_"), features)))

    cat_df = cat_df[feature_names]

     # Add animal name

    cat_df['ID'] = animal_id

    # Update df index

    cat_df.index = range(cat_df.shape[0])

    return cat_df

def load_recurrence_log(recurrence_file, threshold):

    """

    Load information about recurrence

    Args:

        recurrence_file (str): file containing primary tumor recurrence information.

        threshold (int): a limit on how long a mouse must survive post-surgury to be included

    Returns:

        pandas dataframe: animal and recurrence information

    """

    # Read in log file

    dat = pd.read_excel(recurrence_file)

    # Add 'K' to each animal name

    dat['Animal'] = ['K' + str(n) for n in dat['Animal']]

    # Filter out animals without amputation

    dat = dat[dat['AmputationDate'].notnull()]

    # Fill in death dates with today

    dat['DeathDate'][dat['DeathDate'].isnull()] = datetime.now()

    # Find lifespan post amputation

    dat['AmputationDate'] = [i.date() for i in dat['AmputationDate']]

    dat['DeathDate'] = [i.date() for i in dat['DeathDate']]

    dat['Recurrence'] = [i.date() for i in dat['Recurrence']]

    lifespan = dat['DeathDate'] - dat['AmputationDate']

    dat['since_amp'] = [int(life.days) for life in lifespan]

    # Filter out short lifespans without recurrence

    thresh = timedelta(threshold)

    recurr = ~pd.isna(dat['Recurrence'])

    inds = np.logical_or(lifespan > thresh, recurr)

    dat = dat[inds]

    # dat = dat[lifespan > thresh]

    # Get recurrence

    dat['bool_recur'] = dat['Recurrence'].notnull().tolist()

    # Reindex

    dat = dat.reset_index(drop=True)

    return dat

def sort_study_data(rec_df, sum_df, exclude):

    # Create a new dictionary to sort recurrence data

    df = {'animalID': [], 'group': [], 'recurrence': [], 'rec_days': []}

    # Populate df with data from both logs

    n = 0

    for i, animal in enumerate(sum_df['Kirsch lab iD']):

        if any(animal == rec_df['Animal']) and animal not in exclude:

            tmp = rec_df[rec_df['Animal'].str.match(animal)]

            df['animalID'].append(animal)

            df['group'].append(sum_df['Group'][i])

            df['recurrence'].append(tmp['bool_recur'].tolist()[0])

            df['rec_days'].append(int(tmp['since_amp']))

    # Make a Pandas df

    df = pd.DataFrame.from_dict(df)

    # Print general stats

    nrec = sum(df['recurrence'])

    norec = sum(df['recurrence']==False)

    nrecpd1 = sum([df['group'][i] == 'PD1' and df['recurrence'][i] for i in range(len(df))])

    nreccon = sum([df['group'][i] == 'Control' and df['recurrence'][i] for i in range(len(df))])

    norecpd1 = sum([df['group'][i] == 'PD1' and not df['recurrence'][i] for i in range(len(df))])

    noreccon = sum([df['group'][i] == 'Control' and not df['recurrence'][i] for i in range(len(df))])

    print('Number of recurrent tumors:     %d' % nrec)

    print('\tControl:    \t%d' % nreccon)

    print('\tPD1:        \t%d' % noreccon)

    print('Number of non-recurrent tumors: %d' % norec)

    print('\tControl:    \t%d' % noreccon)

    print('\tPD1:        \t%d' % norecpd1)

    return df

def append_rec(df_rec, df_tumor):

    """

    Append radiomics, RT, survival, and recurrence data together

    Args:

        df_rec (pandas dataframe): recurrence

        df_tumor (pandas dataframe): radiomics

    Returns:

        appended pandas dataframe

    """

    # Get the list of animals

    animals = df_rec['animalID'].tolist()

    keys = df_rec.keys().tolist()

    keys.remove('animalID')

    keys.remove('group')

    # Extend the output dataframe

    ltumor = len(df_tumor['ID'])

    df_out = df_tumor.copy()

    for key in keys:

        df_out[key] = pd.Series(np.zeros(ltumor), index=df_out.index)

    for animal in animals:

        rec = df_rec[df_rec['animalID'] == animal]

        for key in keys:

            df_out.loc[df_out['ID'] == animal, key] = rec[key].tolist()

    return df_out

def examine_recurrence(df_tumor, spath, group='Pre'):

    """

    Run classifiers to predict recurrence

    Args:

        df_tumor (pandas dataframe): radiomic data

        spath (str): output path

        group (str): pre/post RT

    Returns:

    """

    # Ensure spath exists

    if not os.path.exists(spath):

        os.makedirs(spath)

    # Separate control and PD1

    df_cnt = df_tumor[df_tumor['Group'] == group + 'Cnt']

    df_pd1 = df_tumor[df_tumor['Group'] == group + 'PD1']

    # Separate by larger pre/post

    df_tumor = df_tumor.loc[(df_tumor['Group'] == group + 'Cnt') | (df_tumor['Group'] == group + 'PD1'), :]

    # Get numeric data

    keys = df_tumor.keys().tolist()

    rm_keys = ['ID', 'Group', 'rec_days']

    [keys.remove(i) for i in rm_keys]

    lab_keys = 'recurrence'

    data_keys = keys

    data_keys.remove(lab_keys)

    # Scale measurements

    rad = df_tumor[data_keys].copy()

    # rad = scale_features(rad)

    label = df_tumor[lab_keys]

    ids = df_tumor['ID']

    # bar_plot(rad_data)

    # tmp = rad.loc[label, 'original_firstorder_10Percentile_T1']

    # rad.loc[label, 'original_firstorder_10Percentile_T1'] = tmp.copy()

    # Plot metrics

    plot_radiomics(rad, label, ids, spath)

    # SVM classification

    svm_classifier_loc(scale_features(rad), label, spath)

    # tspath = os.path.join(spath, 'cnt')

    # svm_classifier_loc(scale_features(df_cnt[data_keys]), df_cnt[lab_keys], tspath)

    # tspath = os.path.join(spath, 'pd1')

    # svm_classifier_loc(scale_features(df_pd1[data_keys]), df_pd1[lab_keys], tspath)

    # NN classification

    nn_classifier_loc(scale_features(rad), label, spath)

    # tspath = os.path.join(spath, 'cnt')

    # nn_classifier_loc(scale_features(df_cnt[data_keys]), df_cnt[lab_keys], tspath)

    # tspath = os.path.join(spath, 'pd1')

    # nn_classifier_loc(scale_features(df_pd1[data_keys]), df_pd1[lab_keys], tspath)

    # sname = os.path.join(spath, 'full.csv')

    # pd.concat((scale_features(rad), label), axis=1).to_csv(sname, index=False)

    # sname = os.path.join(spath, 'cnt.csv')

    # pd.concat((scale_features(df_cnt[data_keys]), df_cnt[lab_keys]), axis=1).to_csv(sname, index=False)

    # sname = os.path.join(spath, 'pd1.csv')

    # pd.concat((scale_features(df_pd1[data_keys]), df_pd1[lab_keys]), axis=1).to_csv(sname, index=False)

def svm_classifier_loc(rad, label, spath):

    """

    Attempts at SVM classification using the full dataset and PCA-based reduced datasets.

    Args:

        rad (pandas dataframe): radiomic features

        label (pandas dataframe): classification label

        spath (str): output directory

    Returns:

    """

    sns.set(font_scale=1.1)

    # Ensure spath exists

    if not os.path.exists(spath):

        os.makedirs(spath)

    old_stdout = sys.stdout

    log_file_path = os.path.join(spath, 'SVM_results.txt')

    log_file = open(log_file_path, 'w')

    sys.stdout = log_file

    print('Training samples: %d, test samples: %d' % (int(rad.shape[0]*0.75), int(rad.shape[0]*0.25)))

    print('Number of features %d' % rad.shape[1])

    # Get index

    inds_1 = label == True

    inds_2 = label == False

    # Create label vector

    label = label.get_values().astype(np.float64)

    rad = rad.get_values().astype(np.float64)

    # Plot euclidean distances - multidimensional scaling

    fig = plot_euclidean_distances(rad, inds_1, inds_2)

    fig.savefig(os.path.join(spath, 'SVM_EuclDist_SVM.png'), dpi=300)

    plt.close(fig)

    # SVM without reduction

    print('Trying SVM classifier without dimensionality reduction:')

    print('-' * 50)

    # fig, clf = svm_classifier(rad.get_values(), inds_1, inds_2)

    fig = svm_cross(rad, label)

    fig.savefig(os.path.join(spath, 'SVM_NoPCA_roc.svg'))

    fig.savefig(os.path.join(spath, 'SVM_NoPCA_roc.png'), dpi=200)

    plt.close(fig)

    with open(os.path.join(spath, 'NoReduction_SVM.model'), 'wb') as f:

        pickle.dump(clf, f)

    # PCA analysis

    # a look at 2 components and SVM

    print('\n')

    print('Trying SVM classifier with 2 principal components:')

    print('-' * 50)

    rad_pca = pca_reduction(rad, npcomps=2)

    fig = svm_cross(rad_pca, label)

    fig.savefig(os.path.join(spath, 'SVM_PCA.svg'))

    plt.close(fig)

    fig = svm_cross(rad_pca, label)

    fig.savefig(os.path.join(spath, 'SVM_PCA2_roc.svg'))

    fig.savefig(os.path.join(spath, 'SVM_PCA2_roc.png'), dpi=300)

    plt.close(fig)

    with open(os.path.join(spath, 'PCA2_SVM.model'), 'wb') as f:

        pickle.dump(clf, f)

    # a look at variance

    print('\n')

    npcomps, fig = pca_analyis(rad)

    print('Trying SVM classifier with %d principal components found from data variance:' % npcomps)

    print('-' * 50)

    rad_pca = pca_reduction(rad, npcomps)

    fig.savefig(os.path.join(spath, 'Cumulative_var.svg'))

    # fig.savefig(os.path.join(spath, 'Cumulative_var.png'), dpi=300)

    plt.close(fig)

    fig = svm_cross(rad_pca, label)

    fig = svm_cross(rad_pca, label)

    # fig.savefig(os.path.join(spath, 'SVM_PCA_roc_var.svg'))

    # fig.savefig(os.path.join(spath, 'SVM_PCA_roc_var.png'), dpi=300)

    plt.close(fig)

    with open(os.path.join(spath, 'SVM_PCA95.model'), 'wb') as f:

        pickle.dump(clf, f)

    # Close log file

    sys.stdout = old_stdout

    log_file.close()

def nn_classifier_loc(rad, label, spath):

    """

    Attempts at NN classification using the full dataset and PCA-based reduced datasets.

    Args:

        rad (pandas dataframe): radiomic features

        label (pandas dataframe): classification label

        spath (str): output directory

    Returns:

    """

    sns.set(font_scale=1.1)

    # Ensure spath exists

    if not os.path.exists(spath):

        os.makedirs(spath)

    old_stdout = sys.stdout

    log_file_path = os.path.join(spath, 'NN_results.txt')

    log_file = open(log_file_path, 'w')

    sys.stdout = log_file

    print('Training samples: %d, test samples: %d' % (int(rad.shape[0]*0.75), int(rad.shape[0]*0.25)))

    print('Number of features %d' % rad.shape[1])

    # Get index

    inds_1 = label == True

    inds_2 = label == False

    # Create label vector

    label = label.get_values().astype(np.float64)

    rad = rad.get_values().astype(np.float64)

    # SVM without reduction

    print('Trying NN classifier without dimensionality reduction:')

    print('-' * 50)

    # fig, clf = svm_classifier(rad.get_values(), inds_1, inds_2)

    fig = neural_network_cross(rad, label)

    fig.savefig(os.path.join(spath, 'NN_NoPCA_roc.svg'))

    fig.savefig(os.path.join(spath, 'NN_NoPCA_roc.png'), dpi=200)

    plt.close(fig)

    # PCA analysis

    # a look at 2 components and SVM

    print('\n')

    print('Trying NN classifier with 2 principal components:')

    print('-' * 50)

    rad_pca = pca_reduction(rad, npcomps=2)

    # fig, clf = svm_classifier(rad_pca, inds_1, inds_2)

    # fig.savefig(os.path.join(spath, 'PCA_NN.png'), dpi=300)

    # plt.close(fig)

    fig = neural_network_cross(rad_pca, label)

    # fig.savefig(os.path.join(spath, 'NN_PCA2_roc.svg'))

    # fig.savefig(os.path.join(spath, 'NN_PCA2_roc.png'), dpi=300)

    plt.close(fig)

    # with open(os.path.join(spath, 'PCA2_SVM.model'), 'wb') as f:

    #     pickle.dump(clf, f)

    # a look at variance

    print('\n')

    npcomps, fig = pca_analyis(rad)

    print('Trying NN classifier with %d principal components found from data variance:' % npcomps)

    print('-' * 50)

    rad_pca = pca_reduction(rad, npcomps)

    # fig.savefig(os.path.join(spath, 'Cumulative_var.png'), dpi=300)

    plt.close(fig)

    # fig, clf = svm_classifier(rad_pca, inds_1, inds_2)

    # fig = neural_network_cross(rad_pca, label)

    # fig.savefig(os.path.join(spath, 'NN_PCA_var_roc.svg'))

    # fig.savefig(os.path.join(spath, 'NN_PCA_var_roc.png'), dpi=300)

    # plt.close(fig)

    # with open(os.path.join(spath, 'PCA95_NN.model'), 'wb') as f:

    #     pickle.dump(clf, f)

    # Close log file

    sys.stdout = old_stdout

    log_file.close()

def mRMR(df_tumor, base_path, group='Post', area='tumor', mrmr_file=None, num_features=10, htmaps=False):

    """

    Compute mRMR feature selection. This funtion has been bundled with code which computes correlation

    maps and later calls training on classifiers for tumor recurrent.

    Args:

        df_tumor (panas dataframe):

        base_path (str): where to save mRMR output

        group (str): pre/post RT

        area (str): masked location

        mrmr_file (str): file of mRMR features to load or create

        num_features (int): number of features to use

        htmaps (bool): whether to create and save all heatmaps

    Returns:

    """

    print('\n')

    print('-'*80)

    print('Processing %s, %s' %(group, area))

    print('Number of features: %d' % num_features)

    dfn = df_tumor.loc[(df_tumor['Group'] == group + 'Cnt') | (df_tumor['Group'] == group + 'PD1'), :]

    dfn = dfn.reset_index()

    keys = df_tumor.keys().tolist()

    rm_keys = ['ID', 'Group', 'rec_days']

    [keys.remove(i) for i in rm_keys]

    lab_keys = 'recurrence'

    data_keys = keys

    data_keys.remove(lab_keys)

    # # Separate by larger pre/post

    # df_tumor = df_tumor[(df_tumor['Group'] == group + 'Cnt') | (df_tumor['Group'] == group + 'PD1')]

    # Create class in the first column

    category = dfn['recurrence']

    # Remove columns

    dfn = dfn[data_keys].copy()

    dfn['Class'] = category.astype(int)

    # Rearrange columns so class is first

    cols = dfn.columns.tolist()

    cols = [cols[-1]] + cols[:-1]

    dfn = dfn[cols]

    # Drop outliers

    dfn = dfn.loc[dfn['original_gldm_GrayLevelNonUniformity_T1'] != 5607, :]

    # Write as csv

    sname = 'tumor_%s_%s.csv' % (group, area)

    dfn.to_csv(sname, index=False, index_label=False)

    # Generate savepath

    if mrmr_file:

        if 'post' in mrmr_file.lower():

            mrmr_time = 'Post'

        elif 'pre' in mrmr_file.lower():

            mrmr_time = 'Pre'

        if 'tumor' in mrmr_file.lower():

            mode = 'tumor'

        elif 'bed' in mrmr_file.lower():

            mode = 'bed'

        else:

            mode = 'edge'

        spath = os.path.join(base_path, 'Analysis', 'mrmr_%s_%s_feat_%s_%s' % (group, area, mrmr_time, mode))

    else:

        spath = os.path.join(base_path, 'Analysis', 'mrmr_%s_%s' % (group, area))

    print('\tSaving to: ', spath)

    if not os.path.exists(spath):

        os.makedirs(spath)

    # Run mRMR

    if not mrmr_file:

        print('\tCreating a new mRMR file')

        log_file_path = os.path.join(base_path, 'Analysis', 'mRMR_results_%s_%s.txt' % (group, area))

        log_file = open(log_file_path, 'w')

        call([PATH_TO_MRMR, '-i', sname, '-t', '1', '-n', '200'],

             stdout=log_file)

        log_file.close()

        # Select out predictive features

        mrmr_file = os.path.join(base_path, 'Analysis', 'mRMR_results_%s_%s.txt' % (group, area))

    else:

        mrmr_file = os.path.join(base_path, 'Analysis', mrmr_file)

        print('\tUsing the specified mRMR file:\n\t\t%s' % mrmr_file)

    # Read mRMR features

    # res = pd.read_csv(mrmr_file, delimiter='\t', header=205, skipfooter=9, skipinitialspace=True, engine='python')

    res = pd.read_csv(mrmr_file, delimiter='\t', header=3, skipfooter=212, skipinitialspace=True, engine='python')

    cols = res.keys().tolist()

    cols = [i.strip() for i in cols]

    res.columns = cols

    features = res['Name'].tolist()

    features = [i.strip() for i in features]

    # Make heatmaps of control and pd1

    if htmaps:

        # Make heatmaps for cntr and pd1

        df_cnt = df_tumor.loc[df_tumor['Group'] == group + 'Cnt', :].copy()

        df_pd1 = df_tumor.loc[df_tumor['Group'] == group + 'PD1', :].copy()

        # Create class in the first column

        category_cnt = df_cnt['recurrence'].astype(int)

        category_pd1 = df_pd1['recurrence'].astype(int)

        # Remove columns

        df_cnt = df_cnt[features]

        df_pd1 = df_pd1[features]

        # Convert columns to float

        for key in features:

            df_cnt.loc[:, key] = pd.to_numeric(df_cnt.loc[:, key], errors='ignore')

            df_pd1.loc[:, key] = pd.to_numeric(df_pd1.loc[:, key], errors='ignore')

        # Rename to remove "original_"

        for key in df_cnt.keys():

            nkey = key.replace('original_', '')

            df_cnt = df_cnt.rename(columns={key: nkey})

            df_pd1 = df_pd1.rename(columns={key: nkey})

        # Set up colors for the animal classification

        lut = dict(zip(dfn['Class'].unique(), 'bg'))

        row_colors_cnt = category_cnt.map(lut)

        row_colors_pd1 = category_pd1.map(lut)

        row_colors_cnt.name = 'Recurrence'

        row_colors_pd1.name = 'Recurrence'

        # sns.set(font_scale=0.8)

        # g = sns.clustermap(df_cnt, figsize=(18, 18), standard_scale=1, metric='correlation', row_colors=row_colors_cnt)

        # g.savefig(os.path.join(spath, 'heatmap_animals_cnt.png'), dpi=300)

        # plt.close()

        #

        # sns.set(font_scale=0.8)

        # g = sns.clustermap(df_pd1, figsize=(18, 18), standard_scale=1, metric='correlation', row_colors=row_colors_pd1)

        # g.savefig(os.path.join(spath, 'heatmap_animals_pd1.png'), dpi=300)

        # plt.close()

        # Remove class category

        # df_cnt = df_cnt.drop('Class', axis=1)

        # df_pd1 = df_pd1.drop('Class', axis=1)

        # Compute correlation matrix

        df_cnt = df_cnt.corr()

        df_pd1 = df_pd1.corr()

        # Sort correlation matrix

        order_cnt = df_cnt.sum(axis=0).argsort()[::-1]

        order_pd1 = df_pd1.sum(axis=0).argsort()[::-1]

        # Reorder row and columns

        df_cnt = df_cnt.iloc[order_cnt, order_cnt]

        df_pd1 = df_pd1.iloc[order_cnt, order_cnt]

        # Plot correlation - control

        sns.set(font_scale=0.8)

        f, ax = plt.subplots(figsize=(18, 10))

        sns.heatmap(df_cnt, vmax=1.0, square=True, ax=ax)

        plt.yticks(rotation=0)

        ax.get_xaxis().set_ticks([])

        plt.xticks(rotation=90)

        # f.savefig(os.path.join(spath, 'heatmap_control_corr.svg'))

        f.savefig(os.path.join(spath, 'heatmap_control_corr.png'), dpi=300)

        plt.close()

        # Plot correlation - PD1

        sns.set(font_scale=0.8)

        f, ax = plt.subplots(figsize=(18, 10))

        sns.heatmap(df_pd1, vmax=1.0, square=True, ax=ax)

        plt.yticks(rotation=0)

        ax.get_xaxis().set_ticks([])

        plt.xticks(rotation=90)

        # f.savefig(os.path.join(spath, 'heatmap_pd1_corr.svg'))

        f.savefig(os.path.join(spath, 'heatmap_pd1_corr.png'), dpi=300)

        plt.close()

        # Plot correlation - difference

        sns.set(font_scale=0.8)

        f, ax = plt.subplots(figsize=(18, 10))

        sns.heatmap(np.abs(df_cnt - df_pd1), vmax=1.0, square=True, ax=ax)

        plt.yticks(rotation=0)

        ax.get_xaxis().set_ticks([])

        plt.xticks(rotation=90)

        # f.savefig(os.path.join(spath, 'heatmap_cnt_pd1_diff_corr.svg'))

        f.savefig(os.path.join(spath, 'heatmap_cnt_pd1_diff_corr.png'), dpi=300)

        plt.close()

    if htmaps:

        # Make heatmap dataframe

        htmp1 = dfn[features].copy()

        # Convert columns to float (from object)

        for key in htmp1.keys():

            htmp1.loc[:, key] = pd.to_numeric(htmp1.loc[:, key], errors='ignore')

        # Rename to remove "original_"

        for key in htmp1.keys():

            nkey = key.replace('original_', '')

            htmp1 = htmp1.rename(columns={key: nkey})

        # Set up colors for the animal classification

        lut = dict(zip(dfn['Class'].unique(), 'bg'))

        row_colors = dfn['Class'].map(lut)

        row_colors.name = 'Recurrence'

        # sns.set(font_scale=0.8)

        # g = sns.clustermap(htmp1, figsize=(18, 18), standard_scale=1, metric='correlation', row_colors=row_colors)

        # g.savefig(os.path.join(spath, 'heatmap_animals.svg'))

        # # g.savefig(os.path.join(spath, 'heatmap_animals.png'), dpi=300)

        # plt.close()

        # Compute correlation matrix - recurrence and no recurrence

        htmp_nrec = dfn.loc[dfn['Class'] == 0, features]

        htmp_rec = dfn.loc[dfn['Class'] == 1, features]

        # Convert columns to float

        for key in features:

            htmp_nrec.loc[:, key] = pd.to_numeric(htmp_nrec.loc[:, key], errors='ignore')

            htmp_rec.loc[:, key] = pd.to_numeric(htmp_rec.loc[:, key], errors='ignore')

        # Rename to remove "original_"

        for key in htmp_nrec.keys():

            nkey = key.replace('original_', '')

            htmp_nrec = htmp_nrec.rename(columns={key: nkey})

            htmp_rec = htmp_rec.rename(columns={key: nkey})

        # Compute correlation matrix

        htmp_nrec = htmp_nrec.corr()

        htmp_rec = htmp_rec.corr()

        # Sort correlation matrix

        order_nrec = htmp_nrec.sum(axis=0).argsort()[::-1]

        order_rec = htmp_rec.sum(axis=0).argsort()[::-1]

        # Reorder row and columns

        htmp_nrec = htmp_nrec.iloc[order_nrec, order_nrec]

        htmp_rec = htmp_rec.iloc[order_nrec, order_nrec]

        # Plot correlation - no recurrence

        sns.set(font_scale=0.8)

        f, ax = plt.subplots(figsize=(18, 10))

        sns.heatmap(htmp_nrec, vmax=1.0, square=True, ax=ax)

        plt.yticks(rotation=0)

        ax.get_xaxis().set_ticks([])

        plt.xticks(rotation=90)

        f.savefig(os.path.join(spath, 'heatmap_noRecur_corr.svg'))

        # f.savefig(os.path.join(spath, 'heatmap_noRecur_corr.png'), dpi=300)

        plt.close()

        # Plot correlation - recurrence

        sns.set(font_scale=0.8)

        f, ax = plt.subplots(figsize=(18, 10))

        sns.heatmap(htmp_rec, vmax=1.0, square=True, ax=ax)

        plt.yticks(rotation=0)

        ax.get_xaxis().set_ticks([])

        plt.xticks(rotation=90)

        f.savefig(os.path.join(spath, 'heatmap_Recur_corr.svg'))

        # f.savefig(os.path.join(spath, 'heatmap_Recur_corr.png'), dpi=300)

        plt.close()

        # Plot correlation - difference

        sns.set(font_scale=0.8)

        f, ax = plt.subplots(figsize=(18, 10))

        sns.heatmap(np.abs(htmp_nrec - htmp_rec), vmax=1.0, square=True, ax=ax)

        plt.yticks(rotation=0)

        ax.get_xaxis().set_ticks([])

        plt.xticks(rotation=90)

        f.savefig(os.path.join(spath, 'heatmap_Recur_diff_corr.svg'))

        # f.savefig(os.path.join(spath, 'heatmap_Recur_diff_corr.png'), dpi=300)

        plt.close()

    # Use only top 10 features

    features = features[:num_features]

    # Make heatmap dataframe - 10 features

    htmp1 = dfn[features].copy()

    # Convert columns to float (from object)

    for key in htmp1.keys():

        htmp1.loc[:, key] = pd.to_numeric(htmp1.loc[:, key], errors='ignore')

    # Rename to remove "original_"

    for key in htmp1.keys():

        nkey = key.replace('original_', '')

        htmp1 = htmp1.rename(columns={key: nkey})

    # Set up colors for the animal classification

    lut = dict(zip(dfn['Class'].unique(), 'bg'))

    row_colors = dfn['Class'].map(lut)

    row_colors.name = 'Recurrence'

    # sns.set(font_scale=0.8)

    # g = sns.clustermap(htmp1, figsize=(18, 18), standard_scale=1, metric='correlation', row_colors=row_colors)

    # g.savefig(os.path.join(spath, 'heatmap_animals.svg'))

    # # g.savefig(os.path.join(spath, 'heatmap_animals.png'), dpi=300)

    # plt.close()

    # Compute correlation matrix - recurrence and no recurrence

    htmp_nrec = dfn.loc[dfn['Class'] == 0, features]

    htmp_rec = dfn.loc[dfn['Class'] == 1, features]

    # Convert columns to float

    for key in features:

        htmp_nrec.loc[:, key] = pd.to_numeric(htmp_nrec.loc[:, key], errors='ignore')

        htmp_rec.loc[:, key] = pd.to_numeric(htmp_rec.loc[:, key], errors='ignore')

    # Rename to remove "original_"

    for key in htmp_nrec.keys():

        nkey = key.replace('original_', '')

        htmp_nrec = htmp_nrec.rename(columns={key: nkey})

        htmp_rec = htmp_rec.rename(columns={key: nkey})

    # Compute correlation matrix

    htmp_nrec = htmp_nrec.corr()

    htmp_rec = htmp_rec.corr()

    # Sort correlation matrix

    order_nrec = htmp_nrec.sum(axis=0).argsort()[::-1]

    order_rec = htmp_rec.sum(axis=0).argsort()[::-1]

    # Reorder row and columns

    htmp_nrec = htmp_nrec.iloc[order_nrec, order_nrec]

    htmp_rec = htmp_rec.iloc[order_nrec, order_nrec]

    # Plot correlation - no recurrence

    sns.set(font_scale=0.8)

    f, ax = plt.subplots(figsize=(18, 10))

    sns.heatmap(htmp_nrec, vmax=1.0, square=True, ax=ax)

    plt.yticks(rotation=0)

    ax.get_xaxis().set_ticks([])

    plt.xticks(rotation=90)

    f.savefig(os.path.join(spath, 'heatmap_noRecur_corr_10.svg'))

    # f.savefig(os.path.join(spath, 'heatmap_noRecur_corr_10.png'), dpi=300)

    plt.close()

    # Plot correlation - recurrence

    sns.set(font_scale=0.8)

    f, ax = plt.subplots(figsize=(18, 10))

    sns.heatmap(htmp_rec, vmax=1.0, square=True, ax=ax)

    plt.yticks(rotation=0)

    ax.get_xaxis().set_ticks([])

    plt.xticks(rotation=90)

    f.savefig(os.path.join(spath, 'heatmap_Recur_corr_10.svg'))

    # f.savefig(os.path.join(spath, 'heatmap_Recur_corr_10.png'), dpi=300)

    plt.close()

    # Plot correlation - difference

    sns.set(font_scale=0.8)

    f, ax = plt.subplots(figsize=(18, 10))

    sns.heatmap(np.abs(htmp_nrec - htmp_rec), vmax=1.0, square=True, ax=ax)

    plt.yticks(rotation=0)

    ax.get_xaxis().set_ticks([])

    plt.xticks(rotation=90)

    f.savefig(os.path.join(spath, 'heatmap_Recur_diff_corr_10.svg'))

    # f.savefig(os.path.join(spath, 'heatmap_Recur_diff_corr_10.png'), dpi=300)

    plt.close()

    # Create a DataFrame with mRMR features

    df_sub = df_tumor[features + rm_keys + [lab_keys]].copy()

    # Perform computations

    examine_recurrence(df_sub, spath, group=group)

    return spath

def plot_radiomics(rad, label, ids, spath):

    """

    Plot radiomic features used for recurrence classification.

    Args:

        rad (pandas dataframe): radiomic features

        label (pandas dataframe): recurrence labels

        ids (list): animal IDs

        spath (str): savepath

    Returns:

    """

    sns.set_style('whitegrid')

    # Reduce the number of features to plot to 5

    keys = rad.keys()

    rad = rad[keys[:6]]

    # Set up figure

    fig = plt.figure(figsize=(10, 6))

    ax = fig.add_axes([0.1, 0.5, 0.80, 0.47])

    # Get the list of radiomic features

    keys = rad.keys()

    # Create a dictionary (later a Dataframe) for plotting with Seaborn

    df = {'Features': [], 'Value': [], 'Recurrence': [], 'ID': []}

    # Populate dictionary, df, with features and values

    for key in keys:

        # Get feature values

        tmp = rad[key].to_numpy()

        # Normalize between -1 and 1

        tmp -= tmp.min()

        tmp /= 0.5 * tmp.max()

        tmp -= 1

        for i in range(len(tmp)):

            # Remove "original" from the plots

            nkey = key.replace('original_', '')

            # Capitalize method

            ind1 = nkey.find('_')

            method = nkey[:ind1]

            if method == 'firstorder':

                method = 'First Order'

            elif method == 'shape':

                method = method.capitalize()

            else:

                method = method.upper()

            # Get modality

            ind2 = nkey.rfind('_')

            modality = nkey[ind2 + 1:]

            if method != 'Shape':

                nkey = '%s\n%s\n%s' % (modality, method, nkey[ind1 + 1:ind2])

            else:

                nkey = '%s\n%s' % (method, nkey[ind1 + 1:ind2])

            # Append features, values, recurrence, and IDs to the dictionary

            df['Features'].append(nkey)

            df['Value'].append(tmp[i])

            if label.iloc[i]:

                df['Recurrence'].append('Yes')

            else:

                df['Recurrence'].append('No')

            df['ID'].append(ids.iloc[i])

    # Convert dictionary to Dataframe

    df = pd.DataFrame.from_dict(df)

    # Plot

    g = sns.boxplot(x='Features', y='Value', hue='Recurrence', data=df, ax=ax, palette="Set1")

    # Rotate feature names

    g.set_xticklabels(g.get_xticklabels(), rotation=60, fontsize=10)

    # g.set_yticklabels(g.get_yticklabels(), fontsize=10)

    g.set_xlabel('Features', fontsize=10)

    g.set_ylabel('Normalized Value')

    # Fix legend

    plt.legend(bbox_to_anchor=(1.01, 1), loc=2, borderaxespad=0., title='Recurrence')

    # Save figure

    fig.savefig(os.path.join(spath, 'feature_box.svg'), format='svg')

    # fig.savefig(os.path.join(spath, 'feature_box.png'), format='png', dpi=300)

def clean_key(key):

    # Remove "original" from the plots

    nkey = key.replace('original_', '')

    # Capitalize method

    ind1 = nkey.find('_')

    method = nkey[:ind1]

    if method == 'firstorder':