# -*- coding: utf-8 -*-
"""
This module allows to build predictive GRNs models using regression methods
"""
from tqdm.autonotebook import tqdm
import pandas as pd
import numpy as np
import networkx as nx
from sklearn.model_selection import cross_val_score
from sklearn.linear_model import BayesianRidge
from grenadine.Evaluation.evaluation import get_top_tfs_per_tg_from_scores
import copy
import random
from deap import base
from deap import creator
from deap import tools
creator.create("FitnessMax", base.Fitness, weights=(1.0,))
creator.create("Individual", list, fitness=creator.FitnessMax)


def ranking_to_networkx(ranking_df):
    """
    This method returns a directed graph generated by networkx

    Args:
        ranking_df (pandas.DataFrame): ranking matrix.
        A ranking matrix contains a row for each possible regulatory link, it
        also contains 4 columns, namely the rank, the score, the transcription
        factor id, and the target gene id.

    Returns:
        networkx.Digraph: networkx gene regulatory network

    """
    return(nx.from_pandas_edgelist(ranking_df,source='TF',target='TG'))

def join_models(*models):
    model = models[0]
    for i in range(1,len(models)):
        model.tgs_tfs.update(models[i].tgs_tfs)
        model.M.update(models[i].M)
        model.eval.update(models[i].eval)
    return model


class Model:
    def __init__(self,
                 tgs_tfs,
                 predictor_class,
                 **predictor_parameters):
        self.tgs_tfs = tgs_tfs
        self.clf = predictor_class(**predictor_parameters)
        self.M = {tg:None for tg in self.tgs_tfs}

    def flatten_evaluation(self):
        evaluation = []
        for c in self.eval:
            for tg in self.eval[c]:
                for s in self.eval[c][tg]:
                    evaluation.append([c,tg,s])
        evaluation = pd.DataFrame(evaluation,columns=["condition","TG","score"])
        return(evaluation)

    def evaluation_tg(self,tg):
        eval_tg = pd.DataFrame({c: m2.eval[c][tg] for c in m2.eval}).T
        eval_tg_mean = eval_tg.mean(axis=1)
        return eval_tg.loc(eval_tg_mean.sort_values(ascending=False).index)

    def evaluation_avg(self):
        eval_avg = pd.DataFrame({c: {tg: m2.eval[c][tg].mean() for tg in m2.eval[c] } for c in m2.eval})
        return(eval_avg)

    def _get_coefs_tg(self,tg):
        coefs = self._get_coefs()
        return pd.DataFrame({c: {tf: coefs[c][tg][tf] for tf in coefs[c][tg] } for c in coefs})

    def fit(self, gene_expression_matrix, conditions=None,cv=3):
        if conditions is not None:
            self.M = {c:{tg:None for tg in self.tgs_tfs} for c in set(conditions)}
            self.eval = {c:{tg:None for tg in self.tgs_tfs} for c in set(conditions)}
            for c in tqdm(list(set(conditions))):
                gene_exp_local = gene_expression_matrix.loc[:,list(conditions==c)]
                for tg in self.tgs_tfs:
                    X = gene_exp_local.loc[self.tgs_tfs[tg]].T
                    y = gene_exp_local.loc[tg]
                    self.M[c][tg] = copy.deepcopy(self.clf.fit(X=X,y=y))
                    self.eval[c][tg] = cross_val_score(self.clf,X,y,cv=cv)
        else:
            self.M = {tg:None for tg in self.tgs_tfs}
            self.eval = {tg:None for tg in self.tgs_tfs}
            for tg in self.tgs_tfs:
                X = gene_expression_matrix.loc[self.tgs_tfs[tg]].T
                y = gene_expression_matrix.loc[tg]
                self.M[tg] = self.clf.fit(X=X,y=y)
                self.eval[tg] = cross_val_score(self.clf,X,y,cv=cv)

    def predict(self, gene_expression_matrix, conditions=None):
        if conditions is not None:
            results = {c:{tg:0 for tg in self.tgs_tfs} for c in set(conditions)}
            for c in list(set(conditions)):
                gene_exp_local = gene_expression_matrix.loc[:,list(conditions==c)]
                for tg in self.tgs_tfs:
                    X = gene_exp_local.loc[self.tgs_tfs[tg]].T
                    if c in self.M:
                        results[c][tg] = self.M[c][tg].predict(X=X)
                    else:
                        results[c][tg] = self.M[tg].predict(X=X)
        else:
            results = {tg:0 for tg in self.tgs_tfs}
            for tg in self.tgs_tfs:
                X = gene_expression_matrix.loc[self.tgs_tfs[tg]].T
                results[tg] = self.M[tg].predict(X=X)
        return results

    def score(self, gene_expression_matrix, conditions=None):
        if conditions is not None:
            results = {c:{tg:0 for tg in self.tgs_tfs} for c in set(conditions)}
            for c in list(set(conditions)):
                gene_exp_local = gene_expression_matrix.loc[:,list(conditions==c)]
                for tg in self.tgs_tfs:
                    X = gene_exp_local.loc[self.tgs_tfs[tg]].T
                    y = gene_exp_local.loc[tg]
                    if c in self.M:
                        results[c][tg] = self.M[c][tg].score(X=X,y=y)
                    else:
                        results[c][tg] = self.M[tg].score(X=X,y=y)
        else:
            results = {tg:0 for tg in self.tgs_tfs}
            for tg in self.tgs_tfs:
                X = gene_expression_matrix.loc[self.tgs_tfs[tg]].T
                y = gene_expression_matrix.loc[tg]
                results[tg] = self.M[tg].score(X=X,y=y)
        return results

    def _get_coefs(self):
        scores = {}
        for c in self.M:
            scores[c] = {}
            if type(self.M[c]) == dict:
                for tg in self.M[c]:
                    if "coef_" in self.M[c][tg].__dict__:
                        scores[c][tg] = {tf: self.M[c][tg].coef_[i] for i,tf in enumerate(self.tgs_tfs[tg])}
            else:
                scores[c] = {tf: self.M[c].coef_[i] for i,tf in enumerate(self.tgs_tfs[c])}
        return(scores)


def Evolve_optimal_model_tg(X,
                            tfs_full,
                            tg,
                            alpha=0.5,
                            genome_size = 30,
                            population_size=50,
                            p_init=0.1,
                            p_mut=0.1,
                            tour_size=10,
                            p_indiv_mut=0.1,
                            p_cross_over=0.5,
                            nb_generations=20,
                            model=BayesianRidge):
    def generate_bit(p):
        if random.random()<p:
            return 1
        return 0

    def map_individual(individual,tfs_full=tfs_full):
        tfs = [tfs_full[i] for i,g in enumerate(individual) if g]
        return(tfs)

    def evalOneMax(individual,alpha=alpha,tg=tg,tfs_full=tfs_full,X=X,model=model):
        if sum(individual)==0:
            return -10,
        tfs = map_individual(individual)
        tgs_tfs = {tg:tfs}
        model = Model(tgs_tfs,model)
        model.fit(X)
        p_active_genes = sum(individual)/len(individual)
        fitness = model.score(X)[tg] - alpha*p_active_genes
        return(fitness,)

    # generate model
    toolbox = base.Toolbox()
    toolbox.register("attr_bool", generate_bit, p_init)
    toolbox.register("individual", tools.initRepeat, creator.Individual, toolbox.attr_bool, genome_size)
    toolbox.register("population", tools.initRepeat, list, toolbox.individual)
    toolbox.register("evaluate", evalOneMax)
    toolbox.register("mate", tools.cxTwoPoint)
    toolbox.register("mutate", tools.mutFlipBit, indpb=p_mut)
    toolbox.register("select", tools.selTournament, tournsize=tour_size)
    pop = toolbox.population(population_size)

    # Compute fitnesses
    fitnesses = list(map(toolbox.evaluate, pop))
    for ind, fit in zip(pop, fitnesses):
        ind.fitness.values = fit
    fits = [ind.fitness.values[0] for ind in pop]

    # distributions hall-of-fame
    fitness_distrib = fits[:]
    best_fitnesses = [max(fits)]
    hall_of_fame = [pop[fits.index(best_fitnesses[0])]]
    generations = [0 for i in range(len(fits))]
    # Begin the evolution
    for g in tqdm(range(nb_generations)):
        offspring = toolbox.select(pop, len(pop))
        offspring = list(map(toolbox.clone, offspring))
        # cross over
        for child1, child2 in zip(offspring[::2], offspring[1::2]):
            if random.random() < p_cross_over:
                toolbox.mate(child1, child2)
                del child1.fitness.values
                del child2.fitness.values
        # mutation
        for mutant in offspring:
            if random.random() < p_indiv_mut:
                toolbox.mutate(mutant)
                del mutant.fitness.values
        # Evaluate new individuals
        invalid_ind = [ind for ind in offspring if not ind.fitness.valid]
        fitnesses = map(toolbox.evaluate, invalid_ind)
        for ind, fit in zip(invalid_ind, fitnesses):
            ind.fitness.values = fit
        # Replace population
        pop = offspring
        # Gather all the fitnesses in one list and print the stats
        fits = [ind.fitness.values[0] for ind in pop]
        best_fitness = max(fits)
        best_fitnesses.append(best_fitness)
        hall_of_fame.append(pop[fits.index(best_fitness)])
        fitness_distrib+=fits
        generations+=[g+1 for i in range(len(fits))]
    fitness_distrib = pd.DataFrame({"fitness":fitness_distrib,
                                    "generation": generations})
    hall_of_fame = list(map(map_individual,hall_of_fame))
    best_model = Model({tg:hall_of_fame[-1]},model)
    best_model.fit(X)
    return fitness_distrib,list(hall_of_fame),best_fitnesses,best_model

def Evolve_optimal_model(X,
                         scores,
                         nb_full_tfs=30,
                         alpha=0.5,
                         genome_size = 30,
                         population_size=50,
                         p_init=0.1,
                         p_mut=0.1,
                         tour_size=10,
                         p_indiv_mut=0.1,
                         p_cross_over=0.5,
                         nb_generations=20,
                         model=BayesianRidge):
    top_tfs = get_top_tfs_per_tg_from_scores(scores,nb_full_tfs)
    models = []
    fitnesses_full = {}
    hall_of_fame_full = {}
    best_fitnesses_full = {}
    for i,tg in enumerate(top_tfs.index):
        print(tg+" "+str(i)+"/"+str(top_tfs.shape[0]))
        results = Evolve_optimal_model_tg(X,
                                          list(top_tfs.loc[tg]),
                                          tg,
                                          alpha,
                                          genome_size ,
                                          population_size,
                                          p_init,
                                          p_mut,
                                          tour_size,
                                          p_indiv_mut,
                                          p_cross_over,
                                          nb_generations,
                                          model)
        fitnesses,hall_of_fame,best_fitnesses,best_model = results
        fitnesses_full[tg] = fitnesses
        hall_of_fame_full[tg] = hall_of_fame
        best_fitnesses_full[tg] = best_fitnesses
        models.append(best_model)
    models = join_models(*models)
    return(models,fitnesses_full,hall_of_fame_full,best_fitnesses_full)

"""
from dream5.gene_expression import load_ecoli
from dream5.tf import load_tf_ecoli
from grenadine.Inference.inference import score_links
from grenadine.Inference.regression_predictors import BayesianRidgeScore as BRS
from grenadine.Preprocessing.standard_preprocessing import z_score
from grenadine.Evaluation.evaluation import get_top_tfs_per_tg_from_scores
from grenadine.Model.model import Evolve_optimal_model_tg,Evolve_optimal_model
from grenadine.Model.model import join_models

import numpy as np
import pandas as pd
X = z_score(load_ecoli(),1)
tf = load_tf_ecoli()[1]
scores = score_links(X,BRS,tf_list=tf,progress_bar=True)#,tg_list=["iap","cspA"])
top_tfs = get_top_tfs_per_tg_from_scores(scores,40)

Evolve_optimal_model_tg(X,list(top_tfs.loc["iap"]),'iap',alpha=0.5)

model,fitnesses_full,hall_of_fame_full,best_fitnesses_full = Evolve_optimal_model(X,scores,alpha=0.3)
model.eval
model.tgs_tfs
"""
