import itertools
import os
import pickle
import random
import warnings

import luigi
import pandas as pd
import scipy as sp
import sklearn as skl
from sqlalchemy import create_engine
from parallel import SGE, sge_available
from connectome.analysis import ConnectivityEstimator, ReciprocityEstimator, RelativeReciprocityEstimator
from figures import style as style
from figures.analysispipeline.datatypes import DatasetTask, JSON, DataFrame
from figures.analysispipeline.folder import OutputFolder
from connectome.model import SimpleMiller, API, ERFEVER, ERFEVERInitialConnectivity
from pydokuwiki import download_file, get_page, extract_json_parameters
import threading
import numpy.linalg as la

INITIAL_SEED = 42


def make_analyse(model, kwargs):
    """
    Create a network and analyse it.


    Parameters
    ----------
    model: Network model
        The model which to be created.

    kwargs: dict
        The parameters are passed to the model

    Returns
    -------
    analysis: dict
        Analysis of the created networks
    """
    print("Start analysis API")
    net = model(**kwargs)
    analysis = ConnectivityEstimator(network=net)()
    analysis.update(ReciprocityEstimator(network=net)())
    analysis.update(RelativeReciprocityEstimator(network=net)())
    analysis.update(kwargs)
    analysis["name"] = model.name
    return analysis


def make_sampler(results):
    lock = threading.Lock()

    def sample(model, pars, key):
        mapper = SGE(priority=0, name="model_sample").map if sge_available() else map
        res = list(mapper(lambda x: make_analyse(model, x), pars))
        with lock:
            results[key] = res

    return sample


class DataBase(DatasetTask):
    def requires(self):
        return OutputFolder()

    def output(self):
        return luigi.LocalTarget(os.path.join(self.input().path, str(self.dataset).split(":")[-1]))

    def run(self):
        download_file(str(self.dataset) + ".xz", self.output().path + ".xz")


class JSONConfiguration(DatasetTask):
    def requires(self):
        return OutputFolder()

    def output(self):
        return JSON(os.path.join(self.input().fn, os.path.splitext(str(self.dataset).split(":")[-1])[0] + ".json"))

    def run(self):
        remote_id = ":".join(str(self.dataset).split(":")[:-1]) + ":start"
        page = get_page(remote_id)
        parsed = extract_json_parameters(page)
        self.output().save(parsed)


class Table(DatasetTask):
    table_name = luigi.Parameter()

    def requires(self):
        return DataBase(self.dataset)

    def output(self):
        return DataFrame(os.path.join(self.input().fn + ".d", self.table_name))

    def run(self):
        engine = create_engine("sqlite:///" + self.input().fn)
        table = pd.read_sql_table(self.table_name, engine)
        self.output().save(table)
        engine.dispose()


class SumStatsPlt(DatasetTask):
    def requires(self):
        return {name : Table(self.dataset, table_name=name)
                for name in ['models_sample_only', 'particles_sample_only',
                             'summary_sample_only', 'parameters_sample_only']}

    def output(self):
        return DataFrame(os.path.join(os.path.dirname(self.input()['models_sample_only'].fn), "sum_stats_plt"))

    def run(self):
        models_sample_only_orig = self.input()['models_sample_only'].load()
        particles_sample_only = self.input()['particles_sample_only'].load()
        summary_sample_only = self.input()['summary_sample_only'].load()
        parameters_sample_only = self.input()['parameters_sample_only'].load()
        models_sample_only = models_sample_only_orig

        #print(parameters_sample_only)
        parameters_sample_only.p_name = parameters_sample_only.p_name.map(lambda x: 'noise' if x == "noise.p" else x)
        parameters_sample_only.p_name = parameters_sample_only.p_name.map(lambda x: 'subsampling' if x == "noise.f" else x)
        warnings.warn("Bad hack here", DeprecationWarning)
        #print(parameters_sample_only)
        #### BAD HACK END
        noise = pd.DataFrame(parameters_sample_only.pivot('particle_id', 'p_name', 'p_value')['noise']
                         .fillna(0)).reset_index()
        #print(noise)
        sum_stats = summary_sample_only.pivot('particle_id', 'ss_name', 'ss_value').reset_index()
        subsampling_levels = (parameters_sample_only[parameters_sample_only.p_name == 'sample_only_fraction_observed']
                              [['particle_id', 'p_value']].rename(columns={'p_value': 'subsampling'}))

        sum_stats = sum_stats.merge(noise)

        if len(subsampling_levels) > 0:
            sum_stats = sum_stats.merge(subsampling_levels)
        else:
            sum_stats['subsampling'] = 1
            warnings.warn("Bad hack here", DeprecationWarning)

        sum_stats_plt = models_sample_only.merge(particles_sample_only).merge(sum_stats)
        sum_stats_plt = sum_stats_plt[['name', 'noise', 'subsampling'] + list(summary_sample_only.ss_name.unique())]

        self.output().save(sum_stats_plt)


class TreeTainingData(DatasetTask):
    def requires(self):
        return SumStatsPlt(self.dataset)

    def output(self):
        return luigi.LocalTarget(os.path.join(os.path.dirname(self.input().fn), "tree", "training_data.npy"))

    def run(self):
        self.output().makedirs()
        sum_stats_plt = pd.read_pickle(self.input().fn)
        sum_stats_plt_for_learning = sum_stats_plt[(sum_stats_plt.noise == 0) & (sum_stats_plt.subsampling == 1)]
        if len(sum_stats_plt_for_learning) == 0:
            sum_stats_plt_for_learning = sum_stats_plt
            warnings.warn("Bad hack here", DeprecationWarning)
        X = sp.asarray(sum_stats_plt_for_learning[style.features_for_classification])
        name_to_number = {name: nr for nr, name in enumerate(style.internal_names)}
        y = sp.asarray(sum_stats_plt_for_learning.name.map(name_to_number))
        with open(self.output().fn, "wb") as f:
            pickle.dump({"X": X, "y" : y}, f)


class DecisionTree(DatasetTask):
    def requires(self):
        return TreeTainingData(self.dataset)

    def output(self):
        return luigi.LocalTarget(os.path.join(os.path.dirname(self.input().fn), "decisiontree.pickle"))

    def run(self):
        self.output().makedirs()
        with open(self.input().fn, "rb") as f:
            input_dict = pickle.load(f)
        X = input_dict["X"]
        y = input_dict["y"]
        nr_samples = X.shape[0]
        nr_classes = len(style.internal_names)

        decision_tree = skl.tree.DecisionTreeClassifier(max_leaf_nodes=len(style.internal_names),
                                                        min_samples_split=nr_samples/nr_classes)

        random.seed(INITIAL_SEED)
        sp.random.seed(INITIAL_SEED)
        decision_tree.fit(X, y)

        with open(self.output().fn, "wb") as f:
            pickle.dump(decision_tree, f)


class AbcSmcNoise(DatasetTask):
    def requires(self):
        return {name: Table(self.dataset, table_name=name) for name in
                ['abc_smc', 'populations', 'models', 'parameters',
                 'particles']}

    def output(self):
        return DataFrame(os.path.join(os.path.dirname(self.input()['abc_smc'].fn), "abc_smc_noise"))

    def run(self):
        abc_smc = self.input()["abc_smc"].load()
        abc_smc.rename(columns={'id': 'abc_smc_id'}, inplace=True)

        populations = self.input()['populations'].load()
        populations.rename(columns={'id': 'population_id'}, inplace=True)

        models = self.input()['models'].load()
        models.rename(columns={'id': 'model_id'}, inplace=True)

        parameters = self.input()['parameters'].load()
        parameters.rename(columns={'id': 'parameter_id', 'name': 'parameter_name'}, inplace=True)

        particles = self.input()['particles'].load()
        particles.rename(columns={'id': 'particle_id'}, inplace=True)

        try:
            abc_smc_noise = pd.DataFrame(abc_smc.merge(populations[populations.t == -1], on='abc_smc_id')
                                 .merge(models).merge(particles).merge(parameters)
                                 .pivot('abc_smc_id', 'parameter_name', 'value')[['noise', 'subsampling_fraction']]
                                ).fillna(0).reset_index()
            self.output().save(abc_smc_noise)
        except KeyError:
            self.output().save(pd.DataFrame())


class APIOriginalVsModifiedData(luigi.Task):
    def requires(self):
        return OutputFolder()

    def output(self):
        return DataFrame(os.path.join(self.input().fn, "API_original_vs_modified.pd"))

    def run(self):
        random.seed(INITIAL_SEED)
        sp.random.seed(INITIAL_SEED)
        fixed_pars = {"nr_neurons": 2000, "inh_ratio": .1, "p_inh": .6}

        original = SimpleMiller(**fixed_pars)
        modified = API(**fixed_pars)

        feature_space_dimension_list = list(map(float, sp.arange(5, 51, 5)))

        pars = [{"p_exc": float(sp.rand() * .3 + .1),
                 "n_pow": float(sp.rand() * 2 + 4),
                 "feature_space_dimension": int(fdim)}
                for _ in range(50)
                for fdim in feature_space_dimension_list]

        results = {}
        sample = make_sampler(results)
        sample(original, pars, "original")
        sample(modified, pars, "modified")

        data = pd.DataFrame(results["original"]).append(pd.DataFrame(results["modified"]))
        self.output().save(data)


class FEVEROriginalVsModifiedData(luigi.Task):
    def requires(self):
        return OutputFolder()

    def output(self):
        return JSON(os.path.join(self.input().fn, "FEVER_original_vs_modified.json"))

    def run(self):
        random.seed(INITIAL_SEED)
        sp.random.seed(INITIAL_SEED)
        nr_neurons = 2000
        p_inh = .6
        inh_ratio = .1
        original = ERFEVERInitialConnectivity(nr_neurons=nr_neurons, p_inh=p_inh, inh_ratio=inh_ratio,
                                              p_initial_graph_exc=0, p_initial_graph_inh=0)

        modified = ERFEVER(p_inh=p_inh, nr_neurons=nr_neurons, inh_ratio=inh_ratio)
        feature_space_dimension_list = sp.arange(80, 200, 20)

        nr_exc_points = 50
        p_exc_min = .05

        original_data_pars = [dict(p_exc=float(p_exc), feature_space_dimension=int(fdim))
                              for fdim in feature_space_dimension_list
                              for p_exc in p_exc_min + sp.rand(nr_exc_points) * (min(fdim / nr_neurons, 1) - p_exc_min)]

        modified_data_pars = [dict(p_exc=float(p_exc), feature_space_dimension=int(fdim),
                                   feverization_ratio=float(feverization_ratio))
                              for fdim in feature_space_dimension_list
                              for p_exc in p_exc_min + sp.rand(nr_exc_points) * (.4 - p_exc_min)
                              for feverization_ratio in sp.linspace(.5, .9, 5)]

        results = {}

        sample = make_sampler(results)

        orig_thread = threading.Thread(target=sample, args=(original, original_data_pars, "original"))
        mod_thread = threading.Thread(target=sample, args=(modified, modified_data_pars, "modified"))

        orig_thread.start()
        mod_thread.start()
        orig_thread.join()
        mod_thread.join()

        self.output().save(results)


class FEVEROriginalVsModifiedRuleSatisfaction(luigi.Task):
    REPETITIONS = 25

    def requires(self):
        return OutputFolder()

    def output(self):
        return DataFrame(os.path.join(self.input().fn, "FEVER_original_vs_modified_rule_satisfaction.pd"))

    def run(self):

        n_neurons = 2000

        def simulate(p_exc, p_inh, feverization_ratio):

            f = ERFEVER(p_exc=p_exc, p_inh=p_inh, nr_neurons=n_neurons,
                        feverization_ratio=feverization_ratio,
                        inh_ratio=0,
                        feature_space_dimension=180)

            net = f()

            adj = net.adjacency_matrix

            fv = f.get_feature_vectors()

            norm = la.norm(fv - fv @ adj, 2)
            print(f)
            print("norm", norm)
            print("init conn", f.fever_initial_connectivity_)
            return {"pe": p_exc,
                    "pi": p_inh,
                    "nr_neurons": n_neurons,
                    "norm": norm,
                    "fdim": f.feature_space_dimension,
                    "p_init": f.fever_initial_connectivity_.p_initial_graph_exc}

        res = []
        for k in range(self.REPETITIONS):
            print(f"FEVER Repetition {k}/{self.REPETITIONS}")
            r = simulate(.2, .5, .7)
            res.append(r)
            r = simulate(.03, .03, 1)
            res.append(r)

        df = pd.DataFrame(res)

        self.output().save(df)
