#coding:utf-8

import pandas as pd
import os
import subprocess as sb
import numpy as np
import json
from scipy.stats.mstats import gmean
import random
from joblib import Parallel, delayed

import mpbn
import mpbn.simulation as mpbn_sim
from tqdm import tqdm
import gc

from NORDic.UTILS.utils_state import compare_states

####################
## Spread process ##
####################

def run_experiments(network_name, spreader, gene_list, state, gene_outputs, simu_params, quiet=False):
    ## Create file
    import json
    from subprocess import call as sbcall
    perts, perts_S = [(state.loc[[g for g in lst if (g in state.index)]]+1)%2 for lst in [gene_list, spreader]]
    experiments, nb_sims = [{"name": "mpsim", "rates": simu_params.get("rates", "fully_asynchronous"), "depth": simu_params.get("depth", "constant_unitary")}], simu_params["nb_sims"]
    experiments_di = {
        "bnet_file": network_name,
        "init_active": list(state.loc[state[state.columns[0]]==1].index),
        "nb_sims": nb_sims,
        "mutants": {g+"_"+("KO" if (perts.loc[g][perts.columns[0]]==0) else "OE")+" {"+str(spreader)+"}": {gx: str(pd.concat((perts,perts_S),axis=0).loc[gx][perts.columns[0]]) for gx in [g]+spreader} for g in list(perts.index)},
        "experiments": experiments,
    }
    with open("experiments.json", "w") as f:
        json.dump(experiments_di)
    sbcall("mpbn_sim --save experiments.json", shell=True)
    
#######################################
## INFLUENCE MAXIMIZATION ALGORITHM  ##
#######################################

def compute_similarities(f, x0, A, A_WT, gene_outputs, nb_sims, experiments, unif_proba=True, repeat=1, exp_name="", quiet=False):
    '''
    Compute similarities between any attractor in WT and in mutants, weighted by their probabilities

    ...

    Parameters
    ----------
    f : Boolean Network (MPBN) object
        the mutated network
    x0 : MPBN object
        initial state
    A : Attractor list
        list of attractors in mutant
    A_WT : Attractor list
        list of attractors in WT
    gene_outputs : Python character string list
        list of node names to check
    nb_sims : Python integer
        number of iterations to compute the probabilities
    experiments : Python dictionary list
        list of experiments (different rates/depths)
    unif_proba : Python bool
    	if set to True, all reachable attractors have the same probability
    repeat : Python integer
        [default=1] : how many times should these experiments be repeated
    exp_name : Python character string
        [default=""] : printed info about the experiment (if quiet=True)
    quiet : Python bool
        [default=False] : prints out verbose

    Returns
    ----------
    sim : Python float
        change in attractors induced by the mutation
    '''
    for exp in experiments:
        if "name" not in exp:
            continue
        dists = []
        for _ in range(repeat):
            rates = getattr(mpbn_sim, f"{exp['rates']}_rates")
            depth = getattr(mpbn_sim, f"{exp['depth']}_depth")
            rates_args = exp.get("rate_args", {})
            depth_args = exp.get("depth_args", {})
            if (not quiet):
                print(exp_name+" "*int(len(exp_name)>0)+(f"- {depth.__name__}{depth_args}\t{rates.__name__}{rates_args}"))
            probs = mpbn_sim.estimate_reachable_attractors_probabilities(f, x0, A, nb_sims, depth(f, **depth_args), rates(f, **rates_args)) ##
            probs = {i: x for i,x in list(probs.items()) if (x>0)}
            if (len(probs)==0 and len(A_WT)>0 and unif_proba):
                print(f"Value of (nb_attractors={len(A_WT)})*(nb_sims={nb_sims}) is too small (nb_sims={nb_sims},#WT attractors={len(A_WT)},#attractors with proba>0={len(probs)}), selecting uniform probabilities")
                probs = {i: 1/len(A) for i in range(len(A))}
            elif ((len(probs)==0 and not unif_proba) or len(A_WT)==0):
                raise ValueError(f"Value of #WT attractors={len(A_WT)},#attractors with proba>0={len(probs)} is too small, no probability for reachable attractors can be computed. limit and nsims={nb_sims} should be increased.")
            attrs = pd.DataFrame({"MUT_%d"%ia: a for ia, a in enumerate(A)}).replace("*",np.nan).astype(float) 
            print(len(probs))
            attrs = attrs[[attrs.columns[i] for i in list(probs.keys())]]
            ## if too many attractors, select the most common ones
            if (attrs.shape[1]>45000):
                idx_common = np.argsort([probs[i] for i in range(len(probs))]).tolist()
                idx_common.reverse()
                idx_common = idx_common[:45000]
                attrs = attrs[[attrs.columns[i] for i in idx_common]]
                probs = {i:probs[i] for i in idx_common}
                sprobs = sum(list(probs.values()))
                probs = {i:probs[i]*100/sprobs for i in probs}
            probs = np.array([probs[ia]/100 for ia in probs])
            attrs_init = pd.DataFrame({"WT_%d"%ia: a for ia, a in enumerate(A_WT)}).replace("*",np.nan).astype(float)
            sims, nb_gene = compare_states(attrs, attrs_init, gene_outputs)
            assert nb_gene == len(gene_outputs)
            sims = probs.T.dot(sims)
            dt = 1-np.max(sims) #max: minimum of change in (*different*) attractors induced by the subset S
            dists.append(dt)
    return np.mean(dists) if (repeat>1) else dists[0]

def spread(network_name, spreader, gene_list, state, gene_outputs, simu_params, seednb=0, quiet=False):
    '''
    Compute the spread of each gene in gene_inputs+spreader with initial state state on genes gene_outputs. Here, the (single state) spread is defined as the indicator of the emptyness of the intersection between WT and mutant attractors

    ...

    Parameters
    ----------
    network_name : Python character string
        filename of the network in .bnet (needs to be pickable)
    spreader : Python character string list
        subset of node names
    gene_list : Python character string list
        list of node names to perturb in addition to the spreader
    state : Pandas DataFrame
        binary initial state rows/[genes] x columns/[values in {-1,0,1}]
    gene_outputs : Python character string list
        list of node names to check
    simu_params : Python dictionary
        arguments to MPBN-SIM
    seednb : Python integer
        [default=0] : random seed
    quiet : Python bool
        [default=False] : prints out verbose

    Returns
    ----------
    spds : Python float dictionary
        change in mutant attractor states for each gene in gene_list that is, the similarity between any attractor reachable from state in WT and any in mutant spreader+{g} where g in gene_list
    '''
    limit=simu_params.get("limit",0) ##
    random.seed(seednb)
    np.random.seed(seednb)
    ## 1. Load the Boolean network
    f = mpbn.load(network_name)
    ## 2. Create the initial profile
    if (not quiet):
        print("\t<NORD_PMR> Initial state %s (gene(s):%s)" % (state.columns[0], gene_list[0]))
    x0 = f.zero()
    for i in list(state.loc[state[state.columns[0]]==1].index):
        x0[i] = 1
    ## 3. Get the reachable attractors from initial state in the absence of perturbation ("wild type")
    experiments, nb_sims = [{"name": "mpsim", "rates": simu_params.get("rates", "fully_asynchronous"), "depth": simu_params.get("depth", "constant_unitary")}], simu_params["nb_sims"]
    A_WT = [a for a in tqdm(list(f.attractors(reachable_from=x0,limit=limit)))] ##
    exp = experiments[0]
    rates = getattr(mpbn_sim, f"{exp['rates']}_rates")
    depth = getattr(mpbn_sim, f"{exp['depth']}_depth")
    rates_args = exp.get("rate_args", {})
    depth_args = exp.get("depth_args", {})
    probs_WT = mpbn_sim.estimate_reachable_attractors_probabilities(f, x0, A_WT, nb_sims, depth(f, **depth_args), rates(f, **rates_args))
    probs_WT = {i: x for i,x in list(probs_WT.items()) if (x>0)}
    A_WT = [A_WT[i] for i in list(probs_WT.keys())]
    if (len(A_WT)==0):
        A_WT = [state[state.columns[0]].to_dict()] ## if no convergent states can be found, compare to the initial state
        print("Compare with initial state (initial state %s)" % state.columns[0])
    elif (not quiet):
        print("%d wild type attractors with proba > 0 (initial state %s)" % (len(A_WT), state.columns[0]))
    ## if too many attractors, select the most common ones
    if (len(A_WT)>45000):
        idx_common_WT = np.argsort([probs_WT[i] for i in range(len(probs_WT))]).tolist()
        idx_common_WT.reverse()
        idx_common_WT = idx_common_WT[:45000]
        A_WT = [A_WT[i] for i in idx_common_WT]
        if (not quiet):
            print("> reduced to %d wild type attractors (initial state %s, %d perc. of all attractors)" % (len(A_WT), np.sum([probs_WT[i] for i in idx_common_WT])))
    ## 4. Create the mutated networks
    def patch_model(f, patch):
        f = mpbn.MPBooleanNetwork(f)
        #print(patch)
        for i, fi in patch.items():
            #print((i,fi,fi=="nan", type(fi)))
            if (fi=="nan"):
                continue
            f[i] = fi
        return f
    perts, perts_S = [(state.loc[[g for g in lst if (g in state.index)]]+1)%2 for lst in [gene_list, spreader]]
    mutants = {g: {gx: str(pd.concat((perts,perts_S),axis=0).loc[gx][perts.columns[0]]) for gx in [g]+spreader} for g in list(perts.index)}
    f_mutants = {name: patch_model(f, patch) for name, patch in mutants.items()}
    ## 5. Get the reachable attractors from initial state in the presence of mutations ("mutants" KO/OE)
    ## 6. Estimate probabilities of attractors from Mutants and compute similarities
    spds = [0 if (name_g not in f_mutants) 
    	else compute_similarities(
    		f_mutants[name_g], 
    		x0, 
    		[a for a in tqdm(list(f_mutants[name_g].attractors(reachable_from=x0,limit=limit)))],  ## 
    		A_WT, 
    		[g for g in gene_outputs if (g not in [name_g]+spreader)], 
    		nb_sims, 
    		experiments, 
    		exp_name="Gene %s (%d/%d) in state %s" % (name_g, ig+1, len(gene_list), state.columns[0]),
    		quiet=quiet
    	) for ig, name_g in enumerate(gene_list)]
    return spds

def spread_multistate(network_name, spreader, gene_list, states, gene_outputs, im_params, simu_params, quiet=False):
    '''
    Compute the spread of each gene in gene_inputs+spreader with initial states in states on genes gene_outputs. Here, the (single state) spread is defined as the indicator of the emptyness of the intersection between WT and mutant attractors

    ...

    Parameters
    ----------
    network_name : Python character string
        filename of the network in .bnet (needs to be pickable)
    spreader : Python character string list
        subset of node names
    gene_list : Python character string list
        list of node names to perturb in addition to the spreader
    states : Pandas DataFrame
        binary initial state rows/[genes] x columns/[state ID]
    gene_outputs : Python character string list
        list of node names to check
    im_params : Python dictionary
        arguments to Influence Maximization
    simu_params : Python dictionary
        arguments to MPBN-SIM
    quiet : Python bool
        [default=False] : prints out verbose

    Returns
    ----------
    spds : Python float dictionary
        change in mutant attractor states for each gene in gene_list that is, the geometric mean of similarities between any attractor reachable from state in states in WT and any in mutant spreader+{g} where g in gene_list
    '''
    from multiprocessing import cpu_count
    assert simu_params.get('thread_count', 1)>=1 and simu_params.get('thread_count', 1)<=max(1,cpu_count()-2)
    if (simu_params.get('thread_count', 1)==1):
        sprds_multistate = [spread(network_name, spreader, gene_list, states[[col]], gene_outputs, simu_params, seednb=im_params.get("seed", 0)) for col in states.columns]
    else:
        if (states.shape[1]>1):
            sprds_multistate = Parallel(n_jobs=simu_params['thread_count'], backend='loky')(delayed(spread)(network_name, spreader, gene_list, states[[col]], gene_outputs, simu_params, seednb=im_params.get("seed", 0)) for col in states.columns)
        else:
            sprds_multistate = Parallel(n_jobs=simu_params['thread_count'], backend='loky')(delayed(spread)(network_name, spreader, [gene], states[[col]], gene_outputs, simu_params, seednb=im_params.get("seed", 0)) for col in states.columns for gene in gene_list)
    ## Aggregate the values across the set of initial states
    ## Genes which are not measured in the states are assigned value 0
    if (states.shape[1]>1):
        spds = [(gmean([(s[ig]+1) for s in sprds_multistate])-1) for ig, g in enumerate(gene_list)]
    else:
        spds = [sprds_multistate[0][ig] for ig, g in enumerate(gene_list)]
    return spds

def greedy(network_name, k, states, im_params, simu_params, save_folder=None, quiet=False):
    '''
    Greedy Influence Maximization Algorithm [Kempe et al., 2003]. Finds iteratively the maximum spreader and adds it to the list until the list is of size k

    ...

    Parameters
    ----------
    network_name : Python character string
        bnet network
    k : Python integer
        maximum size of the spreader
    im_params : Python dictionary or None
        [default=None] : parameters of the influence maximization
    states : Pandas DataFrame or None
        [default=None] : list of initial states to consider
    save_folder : Python character string
        [default=None] : where to save intermediary results (if None: do not save intermediary results)
    quiet : Python bool
        [default=False] : prints out verbose

    Returns
    ----------
    S, spreads : Python character string list
        nodes in the spreader set, Python dictionary: spread value associated with every tested subset of nodes
    '''
    random.seed(im_params.get("seed", 0))
    np.random.seed(im_params.get("seed", 0))
    if (im_params.get("gene_inputs", None) is None):
        with open(solution_fname, "r") as f:
            network = str(f.read())
        genes = [x.split(", ")[0] for x in network.split("\n")[:-1]]
        gene_inputs = genes
    else:
        gene_inputs = im_params["gene_inputs"]
    if (im_params.get("gene_outputs", None) is None):
        with open(solution_fname, "r") as f:
            network = str(f.read())
        gene_outputs = [x.split(", ")[0] for x in network.split("\n")[:-1] if (x.split(", ")[1] not in [x.split(", ")[0], "0", "1"])]
    else:
        gene_outputs = im_params["gene_outputs"]
    S, S_unfold, spreads, start_k = [], [], {}, 0
    if (save_folder is not None and os.path.exists(save_folder+"application_regulators.json")):
        with open(save_folder+"application_regulators.json", "r") as f:
            res = json.loads(f.read())
        if (len(res)>0):
            spreads, S, S_unfold, start_k = res["spreads"], res["S"], res["S_unfold"], res["k"]
    S_spread = -1 if (len(spreads)==0 and len(S)==0) else np.max(list(spreads[str(S)].values()))
    for current_k in range(start_k+1, k+1):
        if (not quiet):
            print("<NORD_PMR> Iteration k=%d" % (current_k))
        gene_list = [g for g in gene_inputs if (g not in S_unfold)]
        ## List of Python floats in the order gene_list
        sprds_lst = spread_multistate(network_name, S_unfold, gene_list, states, gene_outputs, im_params, simu_params, quiet=quiet)
        ## can't find spreader that strictly increases the spread of the previous spreader
        sprds_arr = np.array(sprds_lst)
        if (np.max(sprds_arr)<=S_spread):
            if (not quiet):
                print("\n<NORD_PMR> /!\ Can't find a more performant spreader for k=%d" % current_k)
            break
        ## Keep ex-aequo
        ex_aequo = np.argwhere(sprds_arr==np.max(sprds_arr)).flatten().tolist()
        S_spread, nodes = np.max(sprds_arr), [gene_list[x] for x in ex_aequo]
        if (len(nodes)>1):
            nodes = tuple(nodes)
            S_unfold += nodes
            S += [nodes]
        else:
            nodes = nodes[0]
            S_unfold += [nodes]
            S += [[nodes]]
        spreads.update({str(S):{g: sprds_lst[ig] if (g not in S_unfold) else S_spread for ig, g in enumerate(gene_list)}})
        print(f"Saving in {save_folder+'application_regulators.json'}")
        with open(save_folder+"application_regulators.json", "w") as f:
            json.dump({"spreads": spreads, "S": S, "S_unfold": S_unfold, "k": current_k}, f)
            pd.DataFrame(spreads).to_csv(save_folder+"application_regulators.csv")
        print(f"Saved in {save_folder+'application_regulators.json'}")
    spreads_df = pd.DataFrame(spreads)
    return S, spreads_df
