import os
from math import sqrt

import luigi
import scipy as sp
from connectome.model import SORN, Network
from connectome.analysis import ConnectivityEstimator, ReciprocityEstimator, RelativeReciprocityEstimator
from figures import style
import figures.style.size as size
import seaborn as sns
from connectome.model.dynamicsorn import make_sorn
from .datatypes import JSON
from .folder import OutputFolder
from parallel import SGE, sge_available
import pandas as pd
from .auxiliary import FolderPlot


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

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

    def run(self):
        mapper = SGE(priority=0, name="SORN", num_threads=4).map if sge_available() else map
        p_inh = .6

        def make(pars: dict):
            pars = pars.copy()
            pars_no_type = {key: value for key, value in pars.items() if key != "type"}
            if pars["type"] == "original":
                adj = make_sorn(**pars_no_type)
                net = Network(adj)
            else:
                net = SORN(**pars_no_type)()

            for Analysis in [ConnectivityEstimator, ReciprocityEstimator, RelativeReciprocityEstimator]:
                ana_pars = Analysis(network=net)()
                pars.update(ana_pars)
            return pars

        n_samples = 70

        def rand_conn():
            return float(sp.rand() * .3 + .05)

        def rand_eta_stdp():
            return float(sp.rand() * 0.0008 + 0.0008)

        def rand_eta_int():
            return float(sp.rand() * 0.05 + 0.05)

        pars_original = [dict(p_exc=rand_conn(), nr_exc=200, nr_inh=int(0.2*200), eta_intrinsic=0.01,
                              p_inh=p_inh,
                              sigma_noise_input=sqrt(0.04), nr_time_steps=200000,
                              fixed_nr_add_attempts_overwriting_p=.2,
                              eta_equal=0, eta_plus=0.004, eta_minus=-0.004, type="original")
                         for _ in range(n_samples)]

        pars_modified = [dict(p_exc=rand_conn(), nr_neurons=2000, inh_ratio=.1, p_inh=.6, type="modified",
                              eta_stdp=rand_eta_stdp(), eta_intrinsic=rand_eta_int())
                         for _ in range(n_samples)]

        result = list(mapper(make, pars_original + pars_modified))
        self.output().save(result)


class SORNOriginalVsModified(FolderPlot):
    scatter_size = 5

    def requires(self):
        return SORNOriginalVsModifiedData()

    def run(self):
        data = pd.DataFrame(self.input().load())
        original = data[data.type == "original"]
        modified = data[data.type == "modified"]

        for orig_only, suffix in [(True, "_orig_only"), (False, "")]:
            with self.plot("sorn_ree_pee_scatter" + suffix, scale=size.s) as (fig, ax):
                ax.scatter(original.p_ee, original.reciprocity_ee, label=style.latex_name.long("original"),
                        c=sns.utils.get_color_cycle()[0], zorder=1000, clip_on=False,
                        s=self.scatter_size)
                if not orig_only:
                    ax.scatter(modified.p_ee, modified.reciprocity_ee, label=style.latex_name.long("adapted"),
                               c=sns.utils.get_color_cycle()[1],
                               clip_on=False,
                               s=self.scatter_size)
                if not orig_only:
                    ax.legend(loc="center left", bbox_to_anchor=(.8, .5))
                ax.set_zorder(-1000000)
                style.style_r_over_p(ax, "half")
                ax.spines["bottom"].set_position(("outward", 3))



class SORNROverP(FolderPlot):
    def requires(self):
        return SORNOriginalVsModifiedData()

    def run(self):
        df = pd.DataFrame(self.input().load())
        df = df[df.type == "modified"]

        for measure in ["eta_intrinsic", "eta_stdp"]:
            with self.plot(measure) as (fig, ax):
                style.two_d_measure_scatter(df, ax, fig, "p_ee", "reciprocity_ee", measure,
                                            lambda ax, label: ax.set_label(label), vmin=None, vmax=None)
                style.style_r_over_p(ax)