#========================================================================#
# Script Name: corr_mat.py                                               #
#                                                                        #
# Description: Script creates correlation matrices with full and partial #
#              correlations from timecourses generated by parent script  #
#              via ConCorr                                           #
#                                                                        #
# Authors:      Jen Burrell (Mar 20th, 2023)                             #
#========================================================================#
# To Do List:
# sliding window
#    tapered and weighted edges
#    overlap between windows (1 TR to window length)
#    30 s, 60 s, 120 s, and 240 s windows

import sys
import subprocess
import re
import itertools
import os
import warnings
import math
import glob
from sklearn import metrics
import numpy as np
import pandas as pd
import seaborn as sns
import pingouin as pg
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from scipy.cluster import hierarchy
from sympy import symbols, solve
from misc import *
from itertools import repeat
from multiprocessing import Pool


#            python3 "$SDIR"/correlation_matrix.py $dir_path/ROI_timeseries/${file_root}_timeseries.txt $i_sub $corr_out $file_root $atlas_lut_file

#--------------------------------------#
#       Correlation Calculations       #
#          and Matrix Building         #
#--------------------------------------#
def corr_mat(*dicts):
    user_info = {}
    for dict in dicts:
        user_info.update(dict)
    print(user_info)
        
    func_data = user_info['func_data']
    
    # - get path for all participants - #
    for char in func_data:
        if char.isdigit():
            func_data = func_data.replace(char, '*')
            
    dir_path_list = sorted(glob.glob(os.path.dirname(func_data))) # list of all files
    
    # - multiprocessing - #
    pool = Pool() # Create a multiprocessing Pool
    pool.starmap(sub_corr, zip(dir_path_list, repeat(user_info))) # process gets masks
    # Close pool, keep kids safe
    pool.close()
    pool.join()

def sub_corr(dir_path, user_info):
    # - define variables in use - #
    WDIR = user_info['WDIR']
    corr_coef = user_info['corr_coef']
    p_adjust = user_info['p_adjust']
    partial = user_info['partial']
    atlas = user_info['atlas'] # name of chosen atlas
    parce_name = user_info['root'] # name of atlas
    atlas_lut = user_info['lut']
    
    # - load in data - #
    if not file_exists(f"{dir_path}/ROI_timeseries/{parce_name}_timeseries.txt"):
        print(f"{dir_path}/ROI_timeseries/{parce_name}_timeseries.txt")
        exit(f"Please run timeseries generation for {parce_name} and try again")
    
    d1 = f"{dir_path}/ROI_timeseries/{parce_name}_timeseries.txt"
    i_sub = list(filter(lambda x: "sub-" in x, d1.split("/")))[0].split('-')[1]
    corr_out = f"{WDIR}/ConCorr/sub-{i_sub}/"
    mkdir(corr_out)
    
    task_name = re.sub(f"bold.*/ROI_timeseries/{parce_name}_timeseries.txt", "", d1) # get rid of path after run specific information
    task_name = task_name[task_name.find(i_sub):].split("/")[-1] # remove folders preceeding task information taking only task specific information

    ### --- Read in data --- ###
    print("Reading in the data")
    d1 = pd.read_csv(d1,sep='\t') # read in data
    roi_list = list(d1.columns.values)  # get list of roi_names
    
    if "Schaefer" in atlas:
        atlas_lut = pd.read_table(atlas_lut, sep=' ', header=None)
        net_names= atlas_lut.iloc[: , -1: ]
        net_names_raw = net_names.rename(columns={4:'Networks'})
        net_names = net_names[4].str.split('_', n=3, expand=True) # breaks string of network names
        net_names['Network'] = net_names[1] + "_" + net_names[2] # stitches together hemispheric networks
    #    net_names['Network'] = net_names[2] # just networks
        net_names = net_names.pop('Network') # gets network names
        
    warnings.simplefilter(action='ignore', category=FutureWarning) # pingouin needs to update their code so ignore warnings
    if len(d1.axes[0]) < len(d1.axes[1]): # if less data points than covariates, partial correlations are invalid
        print("Partial Correlations will not be calculated, less data points than covariates")

    ### --- stats --- ###
    print("Starting stats")
    mkdir(corr_out, f"Stats/{parce_name}")
    corr1_name = corr_out + 'Stats/' + parce_name + '/' + task_name + 'full_corr_stats.csv'
    if len(d1.axes[0]) > len(d1.axes[1]) and partial == 'Yes':
        corr2_name = corr_out  + 'Stats/' + parce_name + '/' + task_name + 'partial_corr_stats.csv'
    
    ### --- sliding window (temporal decomposition) --- ###
    if user_info['slide'] == 'Yes'
        window = user_info['window']
        TR = user_info['TR']
        shift = user_info['shift'] # sliding bar from TR to window length
        edges = user_info['edges'] # rectangle, hamming (np.hamming()), tukey (scipy.signal.windows.tukey())
        if edges == 'rectangle':
            
        elif edges == 'hamming':
        elif edges == 'tukey':
    
    
    
    
    
    
    
    
    
    
    # - correlations - #
    corr1 = d1.corr() # full correlation
    warnings.simplefilter(action='ignore', category=RuntimeWarning) # dividing by zero results in error
    corr1_z_plot = np.arctanh(corr1) # fisher's z transformation for graph
    if partial == "Yes": # partial correlation
        corr2 = d1.pcorr()
        corr2_z_plot = np.arctanh(corr2) # fisher's z transformation for graph
    
    # - full correlations metrics - #
    if not file_exists(corr1_name):
        print("Computing Full Correlations")
        corr1_stats = d1.pairwise_corr(method=corr_coef, padjust=p_adjust) # stats for full correlation corrected for multiple comparison
        corr1_z = np.arctanh(corr1_stats.loc[:,"r"].values) # fisher's z transformation for table
        corr1_stats['Fishers Z'] = np.array(corr1_z)
        corr1_stats.to_csv(corr1_name, compression=None) # save full correlation stats table
    else:
        corr1_stats = pd.read_csv(corr1_name)
    # - partial correlations metrics - #
    if partial == "Yes":
        if len(d1.axes[0]) > len(d1.axes[1]) and partial == 'Yes':
            if not file_exists(corr2_name):
                p_corr_stats = pd.DataFrame()
                for ind, subset in enumerate(itertools.combinations(roi_list, len(roi_list) - 2)):
                    corr2_stats = d1.pairwise_corr(covar=list(subset), method=corr_coef)
                    p_corr_stats = pd.concat([p_corr_stats, corr2_stats], ignore_index=True) # stats for all partial comparisons
                    print("Computing partial correlation (",ind,"/",math.comb(len(roi_list), len(roi_list) - 2),")")
                reject, adj_p_val = pg.multicomp(p_corr_stats.loc[:,"p-unc"], method=p_adjust) # calculate adjusted p-values
                adj_p_val = list(adj_p_val)
                p_corr_stats['p-corr'] = np.array(adj_p_val) # add adjusted p-values to stats structure
                corr2_z = np.arctanh(p_corr_stats.loc[:,"r"].values) # fisher's z transformation for table
                p_corr_stats['Fishers Z'] = np.array(corr2_z)
                p_corr_stats.to_csv(corr2_name, compression=None) # save partial correlation stats table
            else:
                p_corr_stats = pd.read_csv(corr2_name)

    ### --- masks --- ###
    print("Making Masks")
    mask1 = np.triu(np.ones_like(corr1_z_plot, dtype=bool)) # mask for upper triangle (data appears in lower triangle)
    mask2 = np.tril(np.ones_like(corr2_z_plot, dtype=bool)) # mask for lower triangle (data appears in upper triangle)

    #---------------------------------------#
    #            Generate Figure            #
    #---------------------------------------#
    mkdir(corr_out, f"Images/{parce_name}") # makes directory
    if len(d1.axes[0]) > len(d1.axes[1]) and partial == "Yes": # if less data points than covariates, partial correlations are invalid
        if len(corr1) <= 30 and user_info['stars'] == 'Yes':
            ### --- create matrix of astriks for p-values --- ###
            star_mat_full = star_map(corr1_stats['p-corr'])
            star_mat_part = star_map_p(p_corr_stats['p-corr'])
            
            ### --- figure generation with star maps --- ###
            print("Building Figure")
            sns.set_theme(style="white")
            fig, ax = plt.subplots(figsize=(10, 10))

            sns.heatmap(corr1_z_plot, mask=mask1, cmap='Blues', vmax=1, vmin=-1,
                        square=True, linewidths=.1, cbar_kws={"shrink": .66, "pad":.02, "label": 'Full Correlations'}, ax=ax, annot=star_mat_full, fmt='') # full correlation in lower triangle
            sns.heatmap(corr2_z_plot, mask=mask2, cmap='YlOrBr', vmax=1, vmin=-1, square=True, linewidths=.1, cbar_kws={"shrink": .66, 'label': 'Partial Correlations'}, ax=ax, annot=star_mat_part, fmt='') # partial correlation in upper triangle
            ax.patch.set_facecolor('black') # diagonals
            
    #        plt.show() # uncomment to see image at generation
            plt.savefig(corr_out + 'Images/' + parce_name + '/' + task_name + 'corr-mat_z.png', bbox_inches="tight", pad_inches=0.3) # saves image folder of subject folder in specified output directory
        else:
            ### --- figure generation --- ###
                print("Building Figure")
                labels = [*range(0, len(corr1_z_plot)+1, (len(corr1_z_plot))//10)]
                sns.set_theme(style="white")
                fig, ax = plt.subplots(figsize=(10, 10))

                ax = sns.heatmap(corr1_z_plot, mask=mask1, cmap='Blues', vmax=1, vmin=-1, square=True, linewidths=0.0, rasterized=True, xticklabels=labels, yticklabels=labels, cbar_kws={"shrink": .66, "pad":.02, "label": 'Full Correlations'}) # full correlation in lower triangle
                ax = sns.heatmap(corr2_z_plot, mask=mask2, cmap='YlOrBr', vmax=1, vmin=-1, square=True, linewidths=0.0, rasterized=True, xticklabels=labels, yticklabels=labels, cbar_kws={"shrink": .66, "pad":.02, "label": 'Partial Correlations'} ) # partial correlation in upper triangle
                ax.patch.set_facecolor('black') # diagonals
                ax.set(title=f"{task_name} using {parce_name}", xlabel='Nodes', ylabel='Nodes')
                ax.set_xticks(labels, labels=labels)
                ax.set_yticks(labels, labels=labels)

    #            plt.show() # uncomment to see image at generation
                os.makedirs(corr_out + 'Images/' + parce_name + '/', exist_ok=True) # makes directory, or moves on
                plt.savefig(corr_out + 'Images/' + parce_name + '/' + task_name + 'corr-mat_z.png', bbox_inches="tight", pad_inches=0.3) # saves image folder of subject folder in specified output directory
    else:
        if len(corr1) <= 30 and user_info['stars'] == 'Yes':
            ### --- create matrix of astriks for p-values --- ###
            star_mat_full = star_map(corr1_stats['p-corr'])

            ### --- figure generation with star maps --- ###
            print("Building Figure")
            sns.set_theme(style="white")
            fig, ax = plt.subplots(figsize=(10, 10))

            sns.heatmap(corr1_z_plot, cmap="viridis", vmax=1, vmin=-1, square=True, linewidths=.1, cbar_kws={"shrink": .66, "pad":.02}, ax=ax, annot=star_mat_full, fmt='') # full correlation in lower triangle
            ax.patch.set_facecolor('black') # diagonals

    #        plt.show() # uncomment to see image at generation
            plt.savefig(corr_out + 'Images/' + parce_name + '/' + task_name + 'corr-mat_z.png', bbox_inches="tight", pad_inches=0.3) # saves image folder of subject folder in specified output directory
        else:
            ### --- figure generation --- ###
            print("Building Figure")
            labels = [*range(0, len(corr1_z_plot)+1, (len(corr1_z_plot))//10)]
            sns.set_theme(style="white")
            fig, ax = plt.subplots(figsize=(10, 10))

            ax = sns.heatmap(corr1_z_plot, cmap="viridis", vmax=1, vmin=-1, square=True, linewidths=0.0, rasterized=True, xticklabels=labels, yticklabels=labels, cbar_kws={"shrink": .66, "pad":.02})
            ax.patch.set_facecolor('black') # diagonals
            ax.set(title=f"{task_name} using {parce_name}", xlabel='Nodes', ylabel='Nodes')
            ax.set_xticks(labels, labels=labels)
            ax.set_yticks(labels, labels=labels)

    #        plt.show() # uncomment to see image at generation
            plt.savefig(corr_out + 'Images/' + parce_name + '/' + task_name + 'corr-mat_z.png', bbox_inches="tight", pad_inches=0.3) # saves image folder of subject folder in specified output directory

    if "Schaefer" in atlas and user_info['cluster'] == 'Yes':
        ### --- cluster figure generation --- ###
        print("Building Figure")
        labels = [*range(0, len(corr1_z_plot)+1, (len(corr1_z_plot))//10)]
        sns.set_theme(style="white")
        # - get colours for the unique networks present - #
        row_colors = np.random.randint(5, size=len(net_names.unique()))
        pal1 = sns.color_palette('GnBu', n_colors=len(net_names.unique())//2) # uncomment for hemispheric networks
        pal2 = sns.color_palette('BuGn', n_colors=len(net_names.unique())//2) # uncomment for hemispheric networks
        palette = pal1 + pal2 # uncomment for hemispheric networks
        #pal1 = sns.color_palette('GnBu', n_colors=len(net_names.unique())) # uncomment for whole networks
        #palette = pal1 # uncomment for whole networks
        lut = dict(zip(net_names.unique(), sns.color_palette(palette, n_colors=len(net_names.unique()))))
        row_colors = net_names.map(lut)

        # - figure - #
        g = sns.clustermap(corr1.reset_index(drop=True), row_colors=row_colors, cmap="viridis", linewidths=0.0, xticklabels=False, yticklabels=False, figsize=(12, 10))
        ax = g.ax_heatmap
        g.fig.subplots_adjust(right=0.7,top=0.96)
        g.ax_cbar.set_position((0.74, .2, .03, .4))
        ax.set(xlabel='Nodes', ylabel='Nodes')
        g.fig.suptitle(f"Clustering of networks from {task_name} using {parce_name}")
        handles = [Patch(facecolor=lut[name]) for name in lut]
        plt.legend(handles, lut, title='Networks', bbox_to_anchor=(0.91, 0.5), bbox_transform=plt.gcf().transFigure, loc='center')

        # - get the clustering for networks from dendrogram - #
        link_net = hierarchy.fcluster(g.dendrogram_col.linkage, len(net_names.unique())*2, criterion='maxclust')
        link_order = g.dendrogram_row.reordered_ind
        net = dict(sorted(zip(link_order, link_net)))
        net = pd.DataFrame.from_dict(net, orient='index')

        # - compare cluster to parce_name - #
        rand = round(metrics.rand_score(net_names, net[0]), 5)
        string = 'Rand score: ' + str(rand)
        plt.text(1, 0, string, transform=ax.transAxes)

    #    plt.show() # uncomment to see image at generation
        plt.savefig(corr_out + 'Images/' + parce_name + '/' + task_name + 'clust-mat.png', bbox_inches="tight", pad_inches=0.3) # saves image folder of subject folder in specified output directory



#***************************************#
#               Star Maps               #
#***************************************#
# - full corr - #
def star_map(p_values):
    corr_star_full = pd.Series(np.empty(len(p_values)))
    for r_ind, corr in enumerate(p_values):
        if corr <= 0.05:
            corr_star_full.loc[r_ind] = '*'
            if corr <= 0.01:
                corr_star_full.loc[r_ind] = '**'
                if corr <= 0.001:
                    corr_star_full.loc[r_ind] = '***'
        else:
            corr_star_full.loc[r_ind] = ''
    x = symbols('x')
    expr = x**2 - x - 2*len(corr_star_full)
    shape = solve(expr)[1]
    star_mat_full = np.chararray((shape, shape), itemsize=3, unicode=True)
    triu = np.triu_indices(len(star_mat_full), k = 1)
    tril = np.tril_indices(len(star_mat_full), -1)
    star_mat_full[triu] = corr_star_full
    star_mat_full[tril] = star_mat_full.T[tril]
    return star_mat_full

# - partial corr - #
def star_map_p(p_values):
    corr_star_part = pd.Series(np.empty(len(p_values)))
    for r_ind, corr in enumerate(p_values):
        if corr <= 0.05:
            corr_star_part.loc[r_ind] = '*'
            if corr <= 0.01:
                corr_star_part.loc[r_ind] = '**'
                if corr <= 0.001:
                    corr_star_part.loc[r_ind] = '***'
        else:
            corr_star_part.loc[r_ind] = ''
    corr_star_part = corr_star_part[::-1] # ROIs tested backwards compared to full correlation, so correct by flipping values
    x = symbols('x')
    expr = x**2 - x - 2*len(corr_star_part)
    shape = solve(expr)[1]
    star_mat_part = np.chararray((shape, shape), itemsize=3, unicode=True)
    triu = np.triu_indices(len(star_mat_part), k = 1)
    tril = np.tril_indices(len(star_mat_part), -1)
    star_mat_part[triu] = corr_star_part
    star_mat_part[tril] = star_mat_part.T[tril]
    return star_mat_part


if __name__ == '__main__':
    user_info = {'WDIR': '/Users/labmanager/Resilience', 'sub_count': 4, 'func_data': '/Users/labmanager/Resilience/derivatives/sub-01/ses-1/func/sub-01_ses-1_task-RS_run-1_bold.feat/filtered_func_data_denoised_standard.nii.gz', 'atlas': 'Schaefer', 'corr_coef': 'pearson', 'p_adjust': 'holm', 'partial': 'Yes', 'stars': 'Yes', 'cluster': 'Yes', 'dir': '/Users/labmanager/JnJ_analyses/ConCorr/Atlases/Schaefer', 'file': 'Schaefer2018_100Parcels_7Networks_order_FSLMNI152_2mm.nii.gz', 'root': 'Schaefer2018_100Parcels_7Networks_order', 'lut': '/Users/labmanager/JnJ_analyses/ConCorr/Atlases/Schaefer/fsleyes_lut/Schaefer2018_100Parcels_7Networks_order.lut', 'max': '100.000000'}
    corr_mat(user_info)
