import sys
import ray
from ray.util.queue import Queue
import numpy as np
import pandas as pd
from time import sleep
import socket
import os
import pickle
import seaborn as sns
import scipy as sp
from sklearn.model_selection import KFold
from pingouin import partial_corr
from nilearn import plotting
from hcpsuite import load_list_from_file, get_nimg_data, printv, nifti_dim_to_cifti_dim, symmetrize_matrices, symmetrize, timer

def get_behav_data(fname, ids=None):
    """
    Read behavioural data for specified IDs from CSV and return dataframe
    """

    behav_data = pd.read_csv(fname, dtype={'Subject': str}, low_memory = False)
    behav_data.set_index('Subject', inplace=True)
    if ids:
        behav_data = behav_data.loc[ids] # Keep only data from specified subjects

    return behav_data


def convert_matrices_to_dataframe(array, subj_ids):
    """
    Takes a NumPy array (subjects:parcels:parcels) and converts it into a Pandas dataframe fit
    for downstream CPM analyses
    """
    assert array.shape[0] == len(subj_ids), "Number of subject IDs is not equal to number of subjects in neuroimage file"
    fc_data = {}
    n = 0
    for id in subj_ids:
        printv("Flattening matrix of subject {} ({} of {}...)".format(id, n+1, len(subj_ids)), update=True)
        tmp = array[n, :, :] # Get matrix of a single subject
        fc_data[id] = tmp[np.triu_indices_from(tmp, k=1)] # Only use upper triangle of symmetric matrix
        n += 1
    printv("\nCreating DataFrame from matrices...")
    fc_data = pd.DataFrame.from_dict(fc_data, orient='index')

    return fc_data


def create_ccc_upper_reshaped(matrix):
    """
    Takes a matrix in the shape of parcels:parcels:1:subjects (e.g. NIfTI style)
    or subjects:parcels:parcles (i.e. CIFTI style) and converts it into a
    matrix of subjects:parcels:parcels, and cleans it of edges of non-interest,
    i.e. cerebrocerebral edges.
    """
    if len(matrix.shape) == 4: # Assume NIfTI style
        matrix = nifti_dim_to_cifti_dim(matrix)
    matrix[:, 99:512, 99:512] = 0
    matrix = np.triu(matrix)

    return matrix


def get_consistent_edges(masks, tail='pos', thresh = 1., color='red', coords=None, plot=True, **plot_kwargs):
    """
    Extracts edges which are consistent in a defined
    percentage of folds and optionally plots them
    """

    masks_tail_array = np.array([v[tail] for v in masks.values()]) # Arrange tail masks in array
    edge_frac = (masks_tail_array.sum(axis=0))/(masks_tail_array.shape[0])
    print("{} {} edges were selected in at least {}% of folds".format((edge_frac>=thresh).sum(), tail, thresh*100))
    edge_frac_square = sp.spatial.distance.squareform(edge_frac)

    node_mask = np.amax(edge_frac_square, axis=0) >= thresh # find nodes that have at least one edge that passes the threshold
    node_size = edge_frac_square.sum(axis=0)*node_mask*10 # size nodes based on how many suprathreshold edges they have

    if plot:
        plotting.plot_connectome(adjacency_matrix=edge_frac_square, edge_threshold=thresh,
                    node_color = color,
                    node_coords=coords, node_size=node_size,
                    display_mode= 'lzry',
                    edge_kwargs={"linewidth": 1, 'color': color},
                    **plot_kwargs)

    return (edge_frac_square >= thresh).astype(int), node_mask


def plot_top_n_edges(matrix, top_n, img_base=None, img_naming_scheme='label', img_suffix='png', labels=None, color='red', check_for_symmetry=False):
    """
    Takes a connectivity/correlation/prediction matrix (e.g. a meaned and masked r_mat from
    pcorr_dfs_to_rmat_pmat(), and plots the top n edges as NetworkX graphs. img_naming_scheme
    is either 'label' (in which case a list of labels has to be supplied) or 'number'
    """
    import networkx as nx
    from matplotlib import pyplot as plt

    if img_naming_scheme == 'label' and not labels:
        raise ValueError("A list of labels needs to be specified if img_namingscheme is 'label'")
    elif labels:
        assert len(labels) == matrix.shape[0], "Number of labels does not match number of nodes"

    if check_for_symmetry: # is_symmetric() is somewhat unreliable, therefore it is disabled by default
        if is_symmetric(matrix):
            matrix = np.tril(matrix) # Clear upper triangle in symmetric matrix to avoid duplicate edges
    else:
        matrix = np.tril(matrix) # Clear upper triangle in symmetric matrix to avoid duplicate edges

    top_n_indices, top_n_values = get_top_n_edges(matrix, top_n)
    figures = []
    for n in range(top_n):
        G = nx.Graph()
        nodes = []
        nodes.append(top_n_indices[0][n])
        nodes.append(top_n_indices[1][n])
        for node in nodes:
            G.add_node(node)
            if labels:
                G.nodes[node]['label'] = labels[node]
            if img_base:
                if img_naming_scheme == 'label':
                    fname = os.path.join(img_base, "{}.{}".format(labels[node], img_suffix))
                elif img_naming_scheme == 'number':
                    fname = os.path.join(img_base, "{}.{}".format(node+1, img_suffix)) # NumPy arrays are 0-indexed, nodes are not
                else:
                    raise ValueError("img_naming_scheme must be either 'label' or 'number'")
                img = plt.imread(fname)
                G.nodes[node]['image'] = img
        G.add_edge(nodes[0], nodes[1], weight=round(top_n_values[n],6))

        # The next section relies heavily on https://stackoverflow.com/a/53968787
        pos = nx.planar_layout(G)

        fig = plt.figure(figsize=(8,4))
        ax = plt.subplot(111)
        ax.set_aspect('equal')
        edge_labels = nx.get_edge_attributes(G, 'weight')
        nx.draw_networkx_edges(G, pos, ax=ax, edge_color=color)
        nx.draw_networkx_edge_labels(G, pos, ax=ax, edge_labels=edge_labels)

        plt.xlim(-0.9,0.9)
        plt.ylim(0.5,-0.5)

        trans = ax.transData.transform
        trans2 = fig.transFigure.inverted().transform

        img_size = 0.5 # this is the image size
        img_size_12 =img_size/2.0
        for n in G:
            xx, yy = trans(pos[n]) # figure coordinates
            xa, ya = trans2((xx,yy)) # axes coordinates
            a = plt.axes([xa-img_size_12,ya-img_size_12, img_size, img_size])
            a.set_aspect('equal')
            a.imshow(G.nodes[n]['image'])
            a.text(0.5, -0.2, G.nodes[n]['label'], transform=a.transAxes, horizontalalignment='center')
            a.axis('off')
        #ax.text(0, 0.1, G.)
        ax.axis('off')
        figures.append(fig)

    return figures


def get_top_n_nodes(results, n, perm=-1, thresh=1, labels=None, names=None, coords=None):
    """
    Takes CPM results and returns top n nodes according to weighted degree
    (weights of edges connecting nodes are summed).
    Returns a dictionary with top n positive and negative nodes (keys 'pos'
    and 'neg' respectively) as a Pandas DataFrame.
    Optionally supply a list of labels, the corresponding label will be added
    as a column.
    If n = None, all nodes will be returned
    """
    import networkx as nx

    r_mat, p_mat = pcorr_dfs_to_rmat_pmat(results['fselection_results'][perm])
    r_mat_mean = np.mean(r_mat, axis=0)

    masks = get_suprathr_edges_new(results['fselection_results'][perm], p_thresh_pos = 0.05, p_thresh_neg = 0.05) # Only use relevant edges

    degrees_sorted = {}
    for tail in ['pos', 'neg']:
        m_cons = get_consistent_edges(masks, tail, thresh=thresh, plot=False) # Only use consistent edges
        r_mat_mean_masked = abs(r_mat_mean * m_cons[0]) # abs() to make summing work on negative tails

        lower = np.tril(r_mat_mean_masked) # Convert into asymmetric matrix and remove NaNs so NetworkX doesn't get confused
        lower[np.isnan(lower)] = 0

        G = nx.from_numpy_array(lower, parallel_edges=False) # parallel_edges=False effectively means weighted edges
        sorted_degrees = sorted(G.degree(weight='weight'), key=lambda x: x[1] if not np.isnan(x[1]) else 0, reverse=True) # Nan breaks sorting

        df_columns = ['Node', 'Degree']
        if labels is not None: # Add label to node tuple
            sorted_degrees = [degree_tuple + (labels[degree_tuple[0]],) for degree_tuple in sorted_degrees]
            df_columns.append('Label')

        if names is not None: # Add name to node tuple
            sorted_degrees = [degree_tuple + (names[degree_tuple[0]],) for degree_tuple in sorted_degrees]
            df_columns.append('Name')

        if coords is not None: # Add coords to node tuple
            sorted_degrees = [degree_tuple + (' '.join(map(str, np.around(coords[degree_tuple[0]], decimals=2))),) for degree_tuple in sorted_degrees]
            df_columns.append('MNI coordinates')

        sorted_degrees_df = pd.DataFrame.from_records(sorted_degrees, columns=df_columns)
        sorted_degrees_df.loc[:, 'Node'] += 1 # Downstream operations (e.g. plotting of nodes) are not 0-based
        degrees_sorted[tail] = sorted_degrees_df[:n]

    return degrees_sorted, G


def plot_predictions(predictions, tail="glm", save=False, title=None, fname='predictions.svg', color='gray'):
    x = predictions['observed'].astype(float)
    y = predictions[tail].astype(float)

    g = sns.regplot(x=x.T.squeeze(), y=y.T.squeeze(), color=color)
    ax_min = min(min(g.get_xlim()), min(g.get_ylim()))
    ax_max = max(max(g.get_xlim()), max(g.get_ylim()))
    g.set_xlim(ax_min, ax_max)
    g.set_ylim(ax_min, ax_max)
    g.set_aspect('equal', adjustable='box')
    g.set_title(title)

    r = get_r_value(x, y)
    g.annotate('r = {0:.2f}'.format(r), xy = (0.7, 0.1), xycoords = 'axes fraction')

    if save:
        fig = g.get_figure()
        fig.savefig(fname, bbox_inches="tight")

    return g


def plot_permutation_results(results_dict_or_list_of_files, plot_real_r=True, title="", **histplot_kwargs):
    """
    Takes either permutation results as a dictionary or a list of saved permutation results
    and plots it as a histography plot after checking for completion.
    Also prints and returns very basic statistics.
    """
    import copy
    if type(results_dict_or_list_of_files) is dict:
        results_dict = copy.deepcopy(results_dict_or_list_of_files.copy())
    elif type(results_dict_or_list_of_files) is list:
        results_dict = {}
        n = 0
        for file in results_dict_or_list_of_files:
            printv("Reading file {}...".format(file))
            results = np.load(file, allow_pickle=True).item()['prediction_results']
            try:
                results_dict[-1] = results[-1]
                del results[-1]
            except KeyError:
                pass
            for perm in results.values():
                results_dict[n] = perm # For permutations we don't need to care about preserving permutation order
                n += 1
    else:
        raise TypeError("Either provide permutation results as a Python dictionary or a list of saved permutation results")

    try:
        if plot_real_r:
            plot_real_r = get_r_value(results_dict[-1]['observed'].astype(float), results_dict[-1]['glm'].astype(float))
        del results_dict[-1]
    except KeyError:
        pass

    # Check results for incomplete permutations
    incomplete_list = check_for_completion(results_dict, print_info=False)
    if len(incomplete_list) > 0:
        for perm in incomplete_list:
            del results_dict[perm]

    r_list = []
    for perm in results_dict.values():
        x = perm['observed'].astype(float)
        y = perm['glm'].astype(float)
        r = get_r_value(x, y)
        r_list.append(r)

    max_value = np.max(r_list)
    min_value = np.min(r_list)
    count = len(r_list)
    min_p = 1/count
    info = "Minimum value: {}\nMaximum value: {}\nNumber of permutations: {} (minimum p value: {})".format(min_value, max_value, count, min_p)

    plot = sns.histplot(r_list, kde="True", **histplot_kwargs)
    plot.set(title=title, xlabel="Pearson’s r")
    if plot_real_r < 1:
        plot.axvline(x = plot_real_r, color = 'y', label = "True r")
        plot.legend()
        info = info + "\nReal r: {}".format(plot_real_r)

    print(info)
    plot.perm_info = info # Save in plot to automatically save info file with safesave()

    return plot, min_value, max_value, count, min_p


def get_r_value(x, y):
    x[np.isnan(x)] = 0
    y[np.isnan(y)] = 0

    r = sp.stats.pearsonr(x, y)[0]

    return r


def get_kfold_indices(subs_list, k):
    """
    Returns a dictionary of train and test indices of length k
    """
    subs_array = np.array(subs_list)
    kf = KFold(n_splits=k, shuffle=True)
    kfold_indices = {}
    kfold_indices['train'] = []
    kfold_indices['test'] = []
    for train_index, test_index in kf.split(subs_list):
        kfold_indices['train'].append(subs_array[train_index])
        kfold_indices['test'].append(subs_array[test_index])

    return kfold_indices


def clean_kfold_indices(kfold_indices, behav_data, noneb_group='Mother_ID'):
    """
    Cleans the train part of kfold indices of entries sharing the same group.
    This was written to exclude closely related subjects in HCP data by using
    'Mother_ID' as the noneb_group (read: non-exchangeability group).

    Returns the cleaned kfold_indices dictionary.
    """
    kfold_indices_clean = {} # Rebuilding from scratch seems easier than using copy.deepcopy()
    kfold_indices_clean['train'] = []
    kfold_indices_clean['test'] = []
    for fold in range(len(kfold_indices['train'])):
        noneb_values = behav_data.loc[kfold_indices['test'][fold], noneb_group]
        mask = behav_data.loc[kfold_indices['train'][fold], noneb_group].isin(noneb_values)
        kfold_indices_clean['train'].append(kfold_indices['train'][fold][~mask])
        kfold_indices_clean['test'].append(kfold_indices['test'][fold])

    return kfold_indices_clean


def pcorr_dfs_to_rmat_pmat(pcorr_dfs):
    """
    Takes a list of correlation DataFrames (post-feature selection), extracts their
    p and r values and puts them in corresponding matrices.

    Input: List of DataFrames
    Returns: r_mat and p_mat
    """
    n_mats = len(pcorr_dfs) # Number of individual matrices
    n_edges = len(pcorr_dfs[0]['r'])
    shape_mats = sp.spatial.distance.squareform(np.zeros(n_edges)).shape # Get squared dimension of matrix
    r_mat = np.zeros((n_mats, shape_mats[0], shape_mats[1]))
    p_mat = np.zeros((n_mats, shape_mats[0], shape_mats[1]))

    for n in range(n_mats):
        r_mat[n] = sp.spatial.distance.squareform(pcorr_dfs[n]['r'])
        p_mat[n] = sp.spatial.distance.squareform(pcorr_dfs[n]['p-val'])

    return r_mat, p_mat


def get_suprathr_edges_new(df_dict, p_thresh_pos=None, p_thresh_neg=None, r_thresh_pos=None, r_thresh_neg=None, percentile_neg=None, percentile_pos=None, top_n_pos=None, top_n_neg=None, fdr_correction=False):
    folds_list = list(df_dict.keys())
    n_edges = len(df_dict[folds_list[0]])
    masks_dict = {}

    for fold in folds_list:
        pcorr_df = df_dict[fold].copy()
        n_edges = len(df_dict[fold])
        masks_dict[fold] = {}
        suprathr_edges_mask = {}
        if p_thresh_pos and p_thresh_neg:
            if fdr_correction:
                a_corr = np.array(pcorr_df['p-val']).copy()
                mask_corr = np.isfinite(a_corr)
                a_corr[mask_corr] = sp.stats.false_discovery_control(a_corr[mask_corr])
                pcorr_df.loc[:, 'p-val'] = a_corr

            suprathr_edges_mask['pos'] = (pcorr_df['r'] > 0) & (pcorr_df['p-val'] <= p_thresh_pos)
            suprathr_edges_mask['neg'] = (pcorr_df['r'] < 0) & (pcorr_df['p-val'] <= p_thresh_neg)
        elif r_thresh_pos and r_thresh_neg:
            suprathr_edges_mask['pos'] = pcorr_df['r'] > r_thresh_pos
            suprathr_edges_mask['neg'] = pcorr_df['r'] < -abs(r_thresh_neg) # r_thresh_neg can be both given as a positive or a negative value
        elif percentile_pos and percentile_neg:
            r_thresh_pos = np.nanpercentile(pcorr_df['r'], percentile_pos)
            r_thresh_neg = np.nanpercentile(pcorr_df['r'][pcorr_df['r'] < 0], 100 - percentile_neg)
            suprathr_edges_mask['pos'] = pcorr_df['r'] > r_thresh_pos
            suprathr_edges_mask['neg'] = pcorr_df['r'] < -abs(r_thresh_neg)
        elif top_n_pos and top_n_neg:
            suprathr_edges_mask['pos'] = np.zeros(pcorr_df.shape[0])
            suprathr_edges_mask['neg'] = np.zeros(pcorr_df.shape[0])
            suprathr_edges_mask['pos'][np.argpartition(pcorr_df['r'][pcorr_df['r'].notna()], -top_n_pos)[-top_n_pos:]] = 1
            suprathr_edges_mask['neg'][np.argpartition(pcorr_df['r'][pcorr_df['r'].notna()], top_n_neg)[:top_n_neg]] = 1
        else:
            raise TypeError('Either p_thresh_{neg, pos} or r_thresh_{neg, pos} or percentile_{neg, pos} or top_n_{pos, neg} needs to be defined.')

        printv("Fold {}: Pos/neg suprathreshold edges (max r pos/max r neg): {}/{} ({}/{})".format(fold+1, suprathr_edges_mask['pos'].sum(), suprathr_edges_mask['neg'].sum(), pcorr_df['r'].max(), pcorr_df['r'].min()))
        for tail in ('pos', 'neg'):
            masks_dict[fold][tail] = np.zeros(n_edges)
            masks_dict[fold][tail][:] = suprathr_edges_mask[tail].astype(bool)

    return masks_dict


def clean_dfs_of_nan_rows(dfs, col):
    """
    Function to delete rows (e.g. subjects) in one or multiple dataframes.
    Useful for removing subjects with incomplete behavioral data from
    behavioral and functional data dataframes.

    Input: dataframe or list of dataframes to clean, name of column with incomplete data
    Output: returns cleaned dataframe(s)

    """
    results = []
    reference_df = None # Dummy for "if not type" check if finding one fails
    try:
        iterator = iter(dfs)
        is_iter=True
    except TypeError:
        is_iter=False
        dfs = [dfs] # Make df iterable for "for df in dfs"

    for df in dfs:
        try:
            check_col = df[col]
            reference_df = df.copy()
            break
        except KeyError:
            continue

    if not type(reference_df) == pd.core.frame.DataFrame:
        return reference_df
        raise TypeError("Provide iterable of or a single Pandas DataFrame containing the column to check for NaNs.")

    to_delete = np.where(np.isnan(reference_df[col]))[0][::-1]

    for df in dfs:
        df = df.loc[reference_df.index]
        results.append(df.drop(reference_df.index[to_delete]))

    if is_iter:
        return results
    else:
        return results[0]


class RayHandler:
    def __init__(self, fc_data, behav_data, behav, covars, n_perm=0, **ray_kwargs):
        self.behav_data = behav_data # For adding kfold_indices
        self.fselection_results = {}
        self.fselection_results[-1] = {} # Create sub-dictionary for original (i.e. non-permuted) data
        self.prediction_results = {}
        self.fold_info = {}

        ray.shutdown() # Make sure Ray is only initialised once
        self.ray_info = ray.init(**ray_kwargs)

        self.actors_list = []
        self.job_list = []

        self.data_dict = {}
        self.data_dict['behav'] = behav
        self.data_dict['covars'] = covars
        self.data_dict['n_perm'] = n_perm
        self.data_dict['data'] = fc_data
        self.data_dict['edges'] = self.data_dict['data'].columns.astype(str) # Save edges columns before adding behavioral columns
        # Pingouin needs all the data (edges, behav, and covars) in a single DataFrame
        if covars:
            self.data_dict['data'][covars] = behav_data[covars]
        if n_perm > 0:
            # It seems to be more efficient to create a separate df and concat later;
            # .to_frame() converts Pandas series into a DataFrame on-the-fly
            behav_df = behav_data[behav].to_frame()
            for perm in range(n_perm):
                behav_df["{}-perm-{}".format(behav, perm)] = np.random.permutation(behav_df[behav])
            behav_df = behav_df.copy() # TODO: try to put this in the for-loop (no warning but maybe performance price)
            # To avaid fragmentation (and the corresponding warning), consolidate into a
            # new DataFrame)
            self.data_dict['data'] = pd.concat([self.data_dict['data'], behav_df], axis=1)
        else:
            self.data_dict['data'][behav] = behav_data[behav]
        self.data_dict['data'].columns = self.data_dict['data'].columns.astype(str)
        self.data_dict['config'] = {'edges_thresh': {"p_thresh_pos": 0.05, "p_thresh_neg": 0.05}, "edges_fdr_correction": False,
                                   'debug': False}

        # Start status and results actor preferrably (status actor) or forcibly
        # (results actor) on head node
        self.status_actor = StatusActor.options(
          scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
          node_id = ray.get_runtime_context().get_node_id(),
          soft = True)).remote()
        self.results_actor = ResultsActor.options(
          scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
          node_id = ray.get_runtime_context().get_node_id(),
          soft = False)).remote()


    def add_kfold_indices(self, n_folds, clean=True):
        subject_ids = self.data_dict['data'].index
        kfold_indices = get_kfold_indices(subject_ids, n_folds)
        if clean:
            kfold_indices = clean_kfold_indices(kfold_indices, self.behav_data)
        self.data_dict['kfold_indices'] = kfold_indices
        printv("You need to (re-) upload data after this operation.")


    def upload_data(self):
        self.data_object = ray.put(self.data_dict)
        printv("Data uploaded.")


    def start_actors(self, job_list):
        printv("Starting actors for {} jobs ...".format(len(job_list)))
        self.job_list.extend(job_list)
        self.actors = [RayActor.remote(self.data_object, job, self.status_actor, self.results_actor) for job in job_list]


    def process_fselection_results(self, result_list):
        for result in results:
            fold = result[0]
            perm = result[1]
            df = result[2]
            self.fselection_results[perm][fold] = df


    def process_prediction_result(self, results_dict):
        if results_dict['perm'] not in self.prediction_results:
            self.prediction_results[results_dict['perm']] = pd.DataFrame()
            self.prediction_results[results_dict['perm']]['observed'] = self.data_dict['data'][self.data_dict['behav']]
            self.fold_info[results_dict['perm']] = {} # Assume that it also does not exist
        for tail in ('pos', 'neg', 'glm'):
            self.prediction_results[results_dict['perm']].loc[results_dict['test_IDs'], [tail]] = results_dict[tail]
        if results_dict['perm'] == -1: # Preserve fselection_results and masks for plotting of edges etc.
            self.fselection_results[results_dict['perm']][results_dict['fold']] = results_dict['fselection_result']
        self.fold_info[results_dict['perm']][results_dict['fold']] = results_dict


    def get_results(self, **get_kwargs):
        results = [1] # Make len > 0
        while len(results) > 0: # Loop e.g. for n > 0 in **get_kwargs, i.e. if you don't want to get results in total from results actor
            N = ray.get(self.results_actor.get_size.remote())[0]
            if 'n' in get_kwargs.keys():
                n = get_kwargs['n']
            else:
                n = N
            printv("Getting {}/{} results...".format(n if n < N else N, N))
            results = ray.get(self.results_actor.get_results.remote(**get_kwargs))
            if 'compress' in get_kwargs.keys() and get_kwargs['compress'] == True:
                import lz4.frame
                results = pickle.loads(lz4.frame.decompress(results))
            for result in results:
                if type(result) == dict:
                    self.process_prediction_result(result)
                if type(result) == list:
                    self.process_fselection_result(result)

        results_dict = {} # Return a results_dict as a consolidated (more explicit, easier to save and load) result representation
        results_dict['prediction_results'] = self.prediction_results
        results_dict['fselection_results'] = self.fselection_results
        results_dict['fold_info'] = self.fold_info
        results_dict['config'] = self.data_dict['config']
        return results_dict


    def status(self, verbose=True):
        n = 1
        self.status_dict = ray.get(self.status_actor.get_status.remote())
        completed_timedeltas = []
        for pid, info in self.status_dict.items():
            if(info['timestamp_stop'] == 0): # Only print alive actors
                print("Actor {} [{}@{}]: {}".format(n, pid, info['node'], info['msg']))
                n += 1
            else:
                completed_timedeltas.append(info['timestamp_stop'] - info['timestamp_start'])

        print("\n")
        n_held, n_done = ray.get(self.results_actor.get_size.remote())
        n_total = len(self.job_list)
        n_remaining = n_total - n_done
        percent_done = (n_done / n_total) * 100

        # Calculate ETA based on the last 100 completed jobs
        # How long did it take to complete the last 100 jobs?
        # -> Get the 100th to last completed actor's timestamp and calculate delta
        if len(completed_timedeltas) < 100:
            eta_msg = "Not yet available"
        else:
            delta_mean = np.mean(completed_timedeltas)
            total_seconds = delta_mean * n_remaining
            relative_seconds = total_seconds / n # Relative to number of actors
            eta_delta = pd.Timedelta(round(relative_seconds), unit='s')
            timestamp_now = pd.Timestamp.now().timestamp()
            eta_absolute = pd.to_datetime(round(timestamp_now) + round(relative_seconds), unit='s', origin='unix')

            eta_msg = "{} ({})".format(eta_delta, eta_absolute)


        print("Jobs done/remaining: {}/{} ({} %)".format(n_done, n_remaining, round(percent_done, 2)))
        print("ETA: {}".format(eta_msg))
        print("\nResults held in distributed storage: {}".format(n_held))

        return n_done, n_remaining, n_held

    def terminate(self):
        ray.shutdown()


@ray.remote
class StatusActor:
    def __init__(self):
        self.status_dict = {}

    def send(self, status):
        pid = status[0]
        node = status[1]
        msg = status[2]
        timestamp_start = status[3]
        timestamp_stop = status[4]
        self.status_dict[pid] = {"msg": msg, "node": node,
                                 "timestamp_start": timestamp_start,
                                 "timestamp_stop": timestamp_stop}

    def exit(self):
        ray.actor.exit_actor()

    def get_status(self):
        return self.status_dict


@ray.remote
class ResultsActor:
    def __init__(self):
        # Create dictionaries to keep results (it makes sense to do this class-wide to add results on-the-fly
        # and for later reference e.g. if get results functions are called too early)
        self.results = []
        self.n_total = 0

    def send(self, results):
        """
        Receives result from worker, determines type of result and calls the appropriate
        processing function
        """
        self.results.append(results)
        self.n_total += 1

    def get_results(self, n=-1, compress=False):
        n_results = len(self.results)
        if n_results < n or n < 0: # a) get all results when n=-1 (default); b) when n > 0 and less than requested results are available to receive, get only available results to avoid error
            n = n_results
        to_return = self.results[0:n]
        del self.results[0:n]
        if compress:
            import lz4.frame
            to_return = lz4.frame.compress(pickle.dumps(to_return))
        return to_return

    def get_size(self):
        current_size = len(self.results)
        return current_size, self.n_total

    def exit(self):
        ray.actor.exit_actor()


@ray.remote(num_cpus=1)
class RayActor:
    def __init__(self, data, job, status_actor, results_actor):
        self.data = data
        self.status_actor = status_actor
        self.results_actor = results_actor

        self.fold = job[1]
        self.perm = job[2]

        self.node = socket.gethostname()
        self.pid = os.getpid()
        self.timestamp_start = pd.Timestamp.now().timestamp()
        self.timestamp_stop = 0

        self.start_job(job)

    def exit(self):
        self.timestamp_stop = pd.Timestamp.now().timestamp()
        # The following status_update is mainly for debugging purposes, as completed
        # jobs are not printed by default
        self.status_update("Completed fold {} of permutation {}.".format(self.fold+1, self.perm+1))
        ray.actor.exit_actor()

    def start_job(self, job):
        job_type = job[0]
        try:
            obj = job[3] # This was added to optionally supply train_subs manually (for debugging or learning purposes); when predicting, a mask _has_ to be supplied
        except IndexError:
            obj = None

        if job_type == 'fselection':
            fselection_result = self.do_fselection(self.fold, self.perm, train_subs=obj)
            self.results_actor.send.remote([self.fold, self.perm, fselection_result])
            self.exit() # Exit so memory gets freed up and no substantial memory leak happens
        elif job_type == 'prediction':
            prediction_result = self.do_prediction(self.fold, self.perm, mask=obj)
            self.results_actor.send.remote(prediction_result)
            self.exit()
        elif job_type == 'fselection_and_prediction':
            fselection_result = self.do_fselection(self.fold, self.perm, train_subs=obj)
            df_dict = {self.fold: fselection_result}
            self.status_update("Thresholding edges...")
            mask = self.get_suprathr_edges_new(df_dict)
            prediction_result = self.do_prediction(self.fold, self.perm, mask=mask[self.fold])
            if self.perm == -1: # Preserve fselection_results and masks for plotting of edges etc.
                prediction_result['fselection_result'] = fselection_result # Include this in result for plotting of edges etc.
            self.results_actor.send.remote(prediction_result)
            self.exit()
        else:
            raise TypeError('Ill-defined job type.')

    def get_suprathr_edges_new(self, df_dict):
        # We need to duplicate this here or "no module found" error will occur
        folds_list = list(df_dict.keys())
        n_edges = len(df_dict[folds_list[0]])
        masks_dict = {}

        for fold in folds_list:
            pcorr_df = df_dict[fold].copy()
            n_edges = len(df_dict[fold])
            masks_dict[fold] = {}
            suprathr_edges_mask = {}
            if self.data['config']['edges_thresh']['p_thresh_pos'] and self.data['config']['edges_thresh']['p_thresh_neg']:
                if self.data['config']['edges_fdr_correction']:
                    a_corr = np.array(pcorr_df['p-val']).copy()
                    mask_corr = np.isfinite(a_corr)
                    a_corr[mask_corr] =  sp.stats.false_discovery_control(a_corr[mask_corr])
                    pcorr_df.loc[:, 'p-val'] = a_corr.astype(np.float64)
                    df_dict[fold]['p-val_fdr'] = a_corr.astype(np.float64)

                suprathr_edges_mask['pos'] = (pcorr_df['r'] > 0) & (pcorr_df['p-val'] <= self.data['config']['edges_thresh']['p_thresh_pos'])
                suprathr_edges_mask['neg'] = (pcorr_df['r'] < 0) & (pcorr_df['p-val'] <= self.data['config']['edges_thresh']['p_thresh_neg'])
            elif self.data['config']['edges_thresh']['r_thresh_pos'] and self.data['config']['edges_thresh']['r_thresh_neg']:
                suprathr_edges_mask['pos'] = pcorr_df['r'] > self.data['config']['edges_thresh']['r_thresh_pos']
                suprathr_edges_mask['neg'] = pcorr_df['r'] < -abs(self.data['config']['edges_thresh']['r_thresh_neg']) # r_thresh_neg can be both given as a positive or a negative value
            elif self.data['config']['edges_thresh']['percentile_pos'] and self.data['config']['edges_thresh']['percentile_neg']:
                r_thresh_pos = np.nanpercentile(pcorr_df['r'], self.data['config']['edges_thresh']['percentile_pos'])
                r_thresh_neg = np.nanpercentile(pcorr_df['r'][pcorr_df['r'] < 0], 100 - self.data['config']['edges_thresh']['percentile_neg'])
                suprathr_edges_mask['pos'] = pcorr_df['r'] > r_thresh_pos
                suprathr_edges_mask['neg'] = pcorr_df['r'] < -abs(r_thresh_neg)
            elif self.data['config']['edges_thresh']['top_n_pos'] and self.data['config']['edges_thresh']['top_n_neg']:
                suprathr_edges_mask['pos'] = np.zeros(pcorr_df.shape[0])
                suprathr_edges_mask['neg'] = np.zeros(pcorr_df.shape[0])
                suprathr_edges_mask['pos'][np.argpartition(pcorr_df['r'][pcorr_df['r'].notna()], -self.data['config']['edges_thresh']['top_n_pos'])[-self.data['config']['edges_thresh']['top_n_pos']:]] = 1
                suprathr_edges_mask['neg'][np.argpartition(pcorr_df['r'][pcorr_df['r'].notna()], self.data['config']['edges_thresh']['top_n_neg'])[:self.data['config']['edges_thresh']['top_n_neg']]] = 1
            else:
                raise TypeError('Either p_thresh_{neg, pos} or r_thresh_{neg, pos} or percentile_{neg, pos} or top_n_{pos, neg} needs to be defined in the config dictionary.')

            for tail in ('pos', 'neg'):
                masks_dict[fold][tail] = np.zeros(n_edges)
                masks_dict[fold][tail][:] = suprathr_edges_mask[tail].astype(bool)

        return masks_dict

    def do_fselection(self, fold, perm, train_subs=None):
        if not train_subs: # Get train_subs from uploaded database if not supplied
            train_subs = self.data['kfold_indices']['train'][fold]
        fselection_result = self.edgewise_pcorr(train_subs, fold, perm)

        return fselection_result

    def do_prediction(self, fold, perm, mask=None):
        if not mask:
            raise ValueError("A mask needs to be supplied in the job description.")
        train_subs = self.data['kfold_indices']['train'][fold]
        test_subs = self.data['kfold_indices']['test'][fold]
        self.status_update("Predicting fold {}...".format(fold+1)) # No detailed update needed, so once is sufficient
        prediction_result = self.predict(mask, train_subs, test_subs, perm)
        prediction_result['fold'] = fold # Preserve fold info

        return prediction_result

    def status_update(self, msg):
        self.status_actor.send.remote([self.pid, self.node, msg, self.timestamp_start, self.timestamp_stop])

    def edgewise_pcorr(self, train_subs, fold, perm, method='pearson'):
        corr_dfs = [] # Appending to list and then creating dataframe is substantially faster than appending to dataframe
        empty_df = pd.DataFrame({'r': {'pearson': np.nan}, 'p-val': {'pearson': np.nan}}) # Handle all-zero edges
        train_data = self.data['data'].loc[train_subs]
        train_data.columns = train_data.columns.astype(str)

        N = len(self.data['edges'])
        n = 1
        percent = round((n/N)*100)

        for edge in self.data['edges']: # Edge-wise correlation
            if (train_data[edge] != 0).any(): # All-zero columns will raise a ValueError exception. This is _way_ faster than try: except:
                if perm >= 0:
                    y = "{}-perm-{}".format(self.data['behav'], perm)
                else:
                    y = self.data['behav']
                if self.data['covars']:
                    pcorr = partial_corr(data=train_data, x=edge, y=y, covar=self.data['covars'], method=method)[['r', 'p-val']] # Taking only the necessary columns speeds this up a few %
                    pcorr['covars'] = True # Debug, remove later
                else: # We could also use pcorr from Pingouin on the entire df, but this is a prohibitively memory-intensive operation; edge-wise like this works just fine. This was introduced to test implausibly good results for the above operation by Pengouin's partial_corr, by setting covars=None we can use SciPy's implementation of Pearson's r
                    pcorr = empty_df.copy()
                    pcorr[['r', 'p-val']] = sp.stats.pearsonr(train_data.loc[:, edge], train_data.loc[:, y]) # We are basically reproducing Pingouin's output format here for unified downstream processing
                    pcorr['covars'] = False # Debug, remove later
            else:
                pcorr = empty_df
            corr_dfs.append(pcorr)
            percent_new = round((n/N)*100)
            if perm >= 0:
                fold_msg = "{} of permutation {}".format(fold+1, perm+1)
            else:
                fold_msg = fold+1
            if percent_new > percent:
                self.status_update("Computing fold {} ({} %)...".format(fold_msg, percent_new))
                percent = percent_new
            n += 1
        self.status_update("Assembling data frame...")
        combined_corr_dfs = pd.concat(corr_dfs) # Assembling df before .put() seems to avoid awfully slow pickling of data through queue (or whatever, it is orders of magnitude faster that way)
        return combined_corr_dfs

    def predict(self, mask_dict, kfold_indices_train, kfold_indices_test, perm):
        train_vectors = pd.DataFrame(self.data['data'].loc[kfold_indices_train, self.data['edges']])
        test_vectors = pd.DataFrame(self.data['data'].loc[kfold_indices_test, self.data['edges']])

        if perm >= 0:
            behav = "{}-perm-{}".format(self.data['behav'], perm)
        else:
            behav = self.data['behav']

        train_behav = pd.DataFrame(self.data['data'].loc[kfold_indices_train, behav])

        model = self.build_models(mask_dict, train_vectors, train_behav)
        prediction_results = self.apply_models(mask_dict, test_vectors, model)
        prediction_results['test_IDs'] = kfold_indices_test
        prediction_results['train_IDs'] = kfold_indices_train
        prediction_results['perm'] = perm
        if perm == -1:
            prediction_results['mask'] = mask_dict
        if self.data['config']['debug']:
            prediction_results["model"] = model # Debug
            prediction_results["mask"] = mask_dict # Debug
            prediction_results["train_IDs"] = kfold_indices_train # Debug

        return prediction_results

    def build_models(self, mask_dict, train_vectors, train_behav):
        """
        Takes a feature mask, sums all edges in the mask for each subject, and uses simple linear
        regression to relate summed network strength to behavior; returns a dictionary with the
        model
        """
        from sklearn.linear_model import LinearRegression # For some weird reason, we had to suddenly move this here after a python package upgrade
        assert train_vectors.index.equals(train_behav.index), "Row indices of functional vectors and behaviour do not match."
        lr = LinearRegression()
        model_dict = {}

        glm = np.zeros((len(train_vectors.index), len(mask_dict.items())))

        t = 0
        # Loop through pos and neg tails
        for tail, mask in mask_dict.items():
            summed_edges = train_vectors.values[:, mask.astype(bool)].sum(axis=1)
            glm[:, t] = summed_edges
            lr.fit(summed_edges.reshape(-1, 1), train_behav)
            model_dict[tail] = (lr.coef_, lr.intercept_)
            t += 1

        glm = np.c_[glm, np.ones(glm.shape[0])]
        model_dict["glm"] = tuple(np.linalg.lstsq(glm, train_behav, rcond=None)[0])

        return model_dict

    def apply_models(self, mask_dict, test_vectors, model_dict):
        """
        Applies a previously trained linear regression model to a test set to generate
        predictions of behavior.
        """

        prediction_results = {}

        glm = np.zeros((test_vectors.shape[0], len(mask_dict.items())))

        # Loop through pos and neg tails
        t = 0
        for tail, mask in mask_dict.items():
            X = test_vectors.loc[:, mask.astype(bool)].sum(axis=1)
            glm[:, t] = X

            slope, intercept = model_dict[tail]
            prediction_results[tail] = slope[0]*X + intercept
            t+=1

        glm = np.c_[glm, np.ones(glm.shape[0])]
        glm_results = np.dot(glm, model_dict["glm"])
        prediction_results["glm"] = glm_results[:, 0] # Transforms 2d array into 1d array

        return prediction_results


def get_overlap(list_of_results_dictionaries, p_thresh_pos=0.05, p_thresh_neg=0.05, thresh=1, top_n_nodes=50, top_n_cerebellar_nodes=-1, top_n_edges=None, odd_one_out=None, coords=None, plot_top_nodes=True, perm=-1, **nodes_kwargs):
    '''
    Take list of results (in dictionary form with
    'fselection_results' as key of primary interest here)
    and return overlap mask between all edges.

    odd_one_out means that the opposite tail of the specified result
    will be used to compare against. E.g. set to 0 to compare 'pos' tail
    of first result with 'neg' tail of remaining results.
    '''

    import networkx as nx
    m_cons = {'pos': [], 'neg': []}
    r_mat_mean_list = [] # Hold averaged (over all folds) r_mat of each result for top_n_edges later

    for result in list_of_results_dictionaries:
        masks = get_suprathr_edges_new(result['fselection_results'][-1], p_thresh_pos=p_thresh_pos, p_thresh_neg=p_thresh_neg)

        r_mat, p_mat = pcorr_dfs_to_rmat_pmat(result['fselection_results'][-1])
        r_mat_mean_tmp = np.mean(r_mat, axis=0)
        r_mat_mean_list.append(np.nan_to_num(r_mat_mean_tmp))

        for tail, cons_edges_list in m_cons.items():
            cons_edges = get_consistent_edges(masks, tail, thresh=thresh, plot=False)
            cons_edges_list.append(cons_edges[0]) # Take first element i.e. the 0/1 integer version of the mask

    overlap = {}
    degrees_sorted = {}
    plot_dict = {}
    top_n_edges_dict = {}
    #if odd_one_out is int:
    #    tmp_pos = m_cons['pos'][odd_one_out].copy() # Quick non-generic solution, but this function is supposed to be non-generic, replaced with generic version down below
    #    m_cons['pos'][odd_one_out] = m_cons['neg'][odd_one_out].copy()
    #    m_cons['neg'][odd_one_out] = tmp_pos.copy()
    for tail, cons_edges_list in m_cons.items():
        if type(odd_one_out) is int: # Inappropriately generic way to do this, but it is kind of elegant: https://stackoverflow.com/questions/61058709/return-next-key-of-a-given-dictionary-key-python-3-6?
            keys = iter(m_cons.keys())
            tail in keys
            next_key = next(keys)
            tmp = m_cons[next_key][odd_one_out].copy()
            m_cons[next_key][odd_one_out] = m_cons[tail][odd_one_out].copy() # Switch keys in tails
            m_cons[tail][odd_one_out] = tmp.copy()
            printv("\nSwitched tails {} and {} of result number {}".format(tail, next_key, odd_one_out))
            odd_one_out = None # Only perform switch once

        overlap[tail] = np.prod(cons_edges_list, axis=0)
        r_mat_mean_list = np.abs(r_mat_mean_list) # For odd_one_out we need to use absolute values
        r_mat_mean = np.mean(r_mat_mean_list, axis=0)
        overlap_weighted = r_mat_mean * overlap[tail]
        overlap['{}_weighted'.format(tail)] = overlap_weighted

        lower = np.tril(overlap['{}_weighted'.format(tail)]) # Convert into asymmetric matrix and remove NaNs so NetworkX doesn't get confused
        lower[np.isnan(lower)] = 0

        G = nx.from_numpy_array(lower, parallel_edges=False) # parallel_edges=False effectively means weighted edges
        sorted_degrees = sorted(G.degree(weight='weight'), key=lambda x: x[1] if not np.isnan(x[1]) else 0, reverse=True) # Nan breaks sorting

        df_columns = ['Node', 'Degree']
        if labels is not None: # Add label to node tuple
            sorted_degrees = [degree_tuple + (labels[degree_tuple[0]],) for degree_tuple in sorted_degrees]
            df_columns.append('Label')

        if names is not None: # Add name to node tuple
            sorted_degrees = [degree_tuple + (names[degree_tuple[0]],) for degree_tuple in sorted_degrees]
            df_columns.append('Name')

        if coords is not None: # Add coords to node tuple
            sorted_degrees = [degree_tuple + (' '.join(map(str, np.around(coords[degree_tuple[0]], decimals=2))),) for degree_tuple in sorted_degrees]
            df_columns.append('MNI coordinates')

        sorted_degrees_df = pd.DataFrame.from_records(sorted_degrees, columns=df_columns)
        sorted_degrees_df.loc[:, 'Node'] += 1 # Downstream operations (e.g. plotting of nodes) are not 0-based
        sorted_degrees_df = sorted_degrees_df[sorted_degrees_df['Degree'] > 0]
        degrees_sorted[tail] = sorted_degrees_df
        printv("{} overlapping edges in {}. network".format(len(G.edges()), tail))

        if coords is not None:
            plot_dict[tail] = {}
            r_mat, p_mat = pcorr_dfs_to_rmat_pmat(result['fselection_results'][perm])
            reduced_dict = reduce_matrix_by_threshold(np.abs(overlap_weighted), coords=coords)
            color = 'red'
            viewer_cmap = 'autumn'
            if tail == 'neg':
                color = 'blue'
                reduced_dict['reduced_matrix'] = np.abs(reduced_dict['reduced_matrix']) * -1
                viewer_cmap = 'cool'

            plot_dict[tail]['connectome_viewer'] = plotting.view_connectome(reduced_dict['reduced_matrix'], reduced_dict['reduced_coords'], symmetric_cmap=False, edge_cmap = viewer_cmap, colorbar=True)

            plot_dict[tail]['connectome'] = plotting.plot_connectome(adjacency_matrix=overlap[tail], node_size=np.abs(overlap[tail]).sum(axis=0)*10,
                                                            node_color=color, edge_cmap='bwr',
                                                            node_coords=coords,
                                                            display_mode= 'lzry',
                                                            edge_kwargs={"linewidth": 1.3, "color": color}, title="Overlap ({})".format(tail))

            if top_n_edges is not None:
                top_n_indices, top_n_values = get_top_n_edges(overlap_weighted, top_n_edges*2)
                top_n_edges_dict[tail] = np.zeros(overlap_weighted.shape)
                top_n_edges_dict[tail][top_n_indices] = overlap_weighted[top_n_indices]
                top_n_edges_mask = np.where(top_n_edges_dict[tail] != 0, 1, 0)
                overlap_masked = top_n_edges_mask * overlap[tail]
                plot_dict[tail]['top_n_edges'] = plotting.plot_connectome(adjacency_matrix=top_n_edges_dict[tail], node_size=np.abs(overlap_masked).sum(axis=0)*10, # node size according to mask to make them not too tiny
                                                                node_color=color, edge_cmap='bwr',
                                                                node_coords=coords,
                                                                display_mode= 'lzry',
                                                                edge_kwargs={"linewidth": 1.3, "color": color}, title="Top {} edges ({})".format(top_n_edges, tail))

            if plot_top_nodes is True:
                figures_dict, viewer, img = plot_nodes(degrees_sorted[tail], "../hcp-suite/data/parcellations/RenTianGlasser_MNI.nii.gz", None, n_slices=6)
                plot_dict[tail]['nodes_img'] = img
                split_nodes = {}
                if top_n_cerebellar_nodes > 1: # Assume splitting in cerebellum/non-cerebellum is wanted
                    regions = ['cerebellar', 'noncerebellar']
                    split_nodes['noncerebellar_nodes'] = degrees_sorted[tail][degrees_sorted[tail]['Node'] > 99]
                    split_nodes['cerebellar_nodes'] = degrees_sorted[tail][degrees_sorted[tail]['Node'] < 100]
                else:
                    regions = ['wholebrain']
                    split_nodes['wholebrain_nodes'] = degrees_sorted[tail]

                for region in regions:
                    if region == 'cerebellar':
                        plot_n_nodes = top_n_cerebellar_nodes
                    else:
                        plot_n_nodes = top_n_nodes
                    figures_dict, viewer, img = plot_nodes(split_nodes['{}_nodes'.format(region)][0:plot_n_nodes], "../hcp-suite/data/parcellations/RenTianGlasser_MNI.nii.gz", "../MNI152_T1_0.5mm_brain.nii.gz", n_slices=6, plot_degree=False, colorbar=False, cmap='gist_ncar', resampling_interpolation="continuous")
                    plot_dict[tail]['{}_nodes'.format(region)] = figures_dict
                    plot_dict[tail]['{}_nodes_viewer'.format(region)] = viewer

    return overlap, degrees_sorted, top_n_edges_dict, plot_dict, m_cons


def check_for_completion(results, print_info=True):
    '''
    Takes a results or prediction results dictionary and
    checks if data are complete, i.e. all jobs including
    permutations have run to completion.

    Returns a list of incomplete analyses.
    '''

    # Accept both results dictionaries and prediction
    # results dictionaries
    if 'prediction_results' in results:
        perms = results['prediction_results']
    elif 'glm' in results[0]:
        perms = results
    else:
        raise TypeError("Either provide an entire results dictionary or the 'permutation_results' sub-dictionary.")

    nan_list = []
    for perm_n, perm in perms.items():
        if np.any(np.isnan(perm['glm'])):
            nan_list.append(perm_n)

    if print_info:
        if len(nan_list) > 0:
            print("There are {} incomplete permutations:\n\t{}".format(len(nan_list), nan_list))
            print("\nUse this command to regenerate the job list for the incomplete permutations:")
            print("\tjob_list = [[\"fselection_and_prediction\", fold, perm] for perm in {} for fold in range(n_folds)]".format(nan_list))
        else:
            print("All analyses complete.")

    return nan_list
