import os

import luigi
import pandas as pd
import scipy.stats as st
import seaborn as sns

from connectome.builder import DistributedFileStorer
from connectome.model import LL
from figures import style as style
from figures.analysispipeline.auxiliary import plot, ReplotMixin
from figures.analysispipeline.datatypes import DatasetTask
from figures.analysispipeline.folder import OutputFolder
from figures.analysispipeline.plots import AllBeePlots, AllMedianPlots, TreePlot, FeatureImportancePlot
from figures.style import color
from ..data import DataBase
from ..auxiliary import FolderPlot
import numpy as np
import figures.style.size
from figures.analysispipeline.config import config


class AllSampleOnlyPlots(DatasetTask):
    def requires(self):
        return [AllBeePlots(self.dataset),
                AllMedianPlots(self.dataset),
                #TreePlot(self.dataset),
                #FeatureImportancePlot(self.dataset)
                ]

    def output(self):
        return [AllBeePlots(self.dataset).output(),
                AllMedianPlots(self.dataset).output(),
                #TreePlot(self.dataset).output(),
                #FeatureImportancePlot(self.dataset).output()
                ]


class MeasureDistribution(ReplotMixin, DatasetTask):
    # force_replot = True
    SCALE = style.size.xs * (.95, 1)

    def requires(self):
        return {"db": DataBase(self.dataset), "folder": OutputFolder()}

    def output(self):
        return luigi.LocalTarget(os.path.join(config.dataset_output_base_folder(self.input()["db"].fn) + ".d", "measure_distribution"))

    def run(self):
        os.makedirs(self.output().path, exist_ok=True)

        db = DistributedFileStorer(self.input()["db"].fn)
        data = db.fetch_all()

        def transform_name(name):
            if name == "ERFEVER":
                return "FEVER"
            return name

        names = [transform_name(data['name']) for data in data]
        values = [data['value'] for data in data]
        measures = [value['result'] for value in values]
        df = pd.DataFrame(measures, names).reset_index().rename(columns={"index": "name"})
        df["nr_neurons"] = df["nr_exc"] + df["nr_inh"]
        df["f_i"] = df["nr_inh"] / df["nr_neurons"]


        df_tsne = df.copy()
        tsne_measures = ["relative_reciprocity_ee", 'relative_reciprocity_ei',
                         'relative_cycles_5', 'in_out_degree_correlation_exc']
        from sklearn.manifold import TSNE
        X = df_tsne[tsne_measures].to_numpy()
        tsne = TSNE()
        X_trans = tsne.fit_transform(X)
        df_tsne["x_tsne"] = X_trans[:, 0]
        df_tsne["y_tsne"] = X_trans[:, 1]

        with plot(os.path.join(self.output().path, "tsne." + config.extension)) as (fig, ax):
            for name, group in df_tsne.groupby("name"):
                ax.scatter(group["x_tsne"], group["y_tsne"],
                           color=color[name],
                           label=style.name[name])
            ax.legend(loc="center left", bbox_to_anchor=(1, .5))
            ax.set_xlabel("t-SNE dimension 1")
            ax.set_ylabel("t-SNE dimension 2")

        # NOTE(amotta): Export data points per model
        for cur_name, cur_data_points in df.groupby('name'):
            cur_out_file = '{}.xlsx'.format(cur_name.lower())
            cur_out_file = os.path.join(self.output().path, cur_out_file)
            cur_data_points.to_excel(cur_out_file)

        df.groupby("name").describe().to_excel(os.path.join(self.output().path, "description.xlsx"))

        measure_order = ["p_ee", "p_ie", "nr_neurons", "f_i", "relative_reciprocity_ee"]
        mean_std = df.groupby("name").agg(["mean", "std"])[measure_order].sort_index()

        mean_std_stacked = mean_std.stack(0)

        def format_nr(nr):
            if nr == 0:
                return "0"
            if nr < 1:
                return "{:.1e}".format(nr)
            return str(nr)

        def format_mean_std(row):
            row = row.round(4)
            return format_nr(row["mean"]) + " ± " + format_nr(row["std"])


        formatted_mean_std = mean_std_stacked.apply(format_mean_std, axis=1).unstack(1)
        formatted_mean_std.to_excel(os.path.join(self.output().path, "mean_std.xlsx"))

        def format_locator(ax):
            ax.locator_params(axis="x", nbins=3)

        def plot_column_distribution(df, column_names, fit=None, ax=None, with_legend=True, with_ylabel=True):
            for k, column_name in enumerate(column_names):
                for name in df.name.unique():
                    col = color[name]
                    sns.distplot(df[df.name == name][column_name], fit=fit, kde=False, label=None,
                                 fit_kws={"color": col,
                                          "label": style.name[name] if k == 0 else None}, color=(1, 1, 1, 0), ax=ax,
                                 hist_kws={"histtype": "step", "linewidth": style.linewidth, "alpha": 1})
                if with_legend:
                    ax.legend(loc="center left", bbox_to_anchor=(1.2, .5))
            ax.set_yscale("log")
            ax.set_ylim(10 ** -1, 10**1)
            format_locator(ax)
            sns.despine(ax=ax)
            for label in ax.get_yticklabels()[2:-2]:
                label.set_visible(False)

            if with_ylabel:
                ax.set_ylabel(style.name["PDF"], labelpad=style.default_label_pad_reduced)
            else:
                ax.set_yticklabels([])

            return ax

        ext =  config.extension
        title_offset = 1.1
        for with_legend in [True, False]:
            for with_ylabel in [True, False]:
                legend_str = "_legend" if with_legend else ""
                legend_str += "_label" if with_ylabel else ""
                for rec_type in ["ee", "ii", "ei", "ie"]:
                    with plot(os.path.join(self.output().path,
                                           "rr_{}{}.{}".format(rec_type, legend_str, ext)), scale=self.SCALE) as (fig, ax):
                        plot_column_distribution(df, ["relative_reciprocity_" + rec_type],
                                                 fit=st.norm, ax=ax, with_legend=with_legend, with_ylabel=with_ylabel)
                        xmax = int(np.ceil(ax.get_xlim()[1]))
                        ax.set_xlim(0, xmax)
                        ax.set_xticks(np.arange(xmax + 1))
                        ax.set_xticks(np.arange(0, xmax, .1), minor=True)
                        ax.set_xticklabels([0, 1] + [""] * (xmax - 2) + [xmax])
                        ax.set_xlabel(style.latex_name("relative_reciprocity_" + rec_type))

                exc_inh_conn_title = "$p_{exc}$, $p_{inh}$"
                with plot(os.path.join(self.output().path, "p_ee_p_ie{}.{}".format(legend_str,ext)), scale=self.SCALE) as (fig, ax):
                    ax = plot_column_distribution(df, ["p_ee", "p_ie"], fit=st.norm, ax=ax,
                                                  with_legend=with_legend, with_ylabel=with_ylabel)
                    ax.set_xlabel("$p_{ee}, p_{ie}$")
                    style.set_ticksmiddle_minor(ax, np.linspace(0, 1, 11), "x")

                with plot(os.path.join(self.output().path, "p_ii_p_ei{}.{}".format(legend_str, ext)), scale=self.SCALE) as (fig, ax):
                    ax = plot_column_distribution(df, ["p_ii", "p_ei"], fit=st.norm, ax=ax,
                                                  with_legend=with_legend, with_ylabel=with_ylabel)
                    ax.set_xlabel("$p_{ii}, p_{ei}$")
                    style.set_ticksmiddle_minor(ax, np.linspace(0, 1, 11), "x")

                with plot(os.path.join(self.output().path,
                                       "in_out_degree_correlation_exc{}.{}".format(legend_str, ext)), scale=self.SCALE) as (fig, ax):
                    ax = plot_column_distribution(df, ["in_out_degree_correlation_exc"], fit=st.norm, ax=ax,
                                                  with_legend=with_legend, with_ylabel=with_ylabel)
                    ax.set_xlabel(style.latex_name("in_out_degree_correlation_exc"))
                    style.set_ticksmiddle_minor(ax, np.linspace(-1, 1, 21), "x")
                    ax.set_xticks([-1, 0, 1])

                with plot(os.path.join(self.output().path, "relative_cycles_5{}.{}".format(legend_str, ext)), scale=self.SCALE) as (fig, ax):
                    ax = plot_column_distribution(df, ["relative_cycles_5"], fit=st.norm, ax=ax,
                                                  with_legend=with_legend, with_ylabel=with_ylabel)
                    ax.set_xlabel(style.latex_name("relative_cycles_5"))
                    style.set_ticksmiddle_minor(ax, np.linspace(0, 7, 8), "x")
                    #ax.set_xticks([0, 1, 10])


class LLAdjacencyMatrix(FolderPlot):
    def run(self):
        import numpy as np
        for k, n_exc_sub in enumerate([3]):
            net = LL(p_inh=.5,
                     p_exc=.2,
                     nr_neurons=2000,
                     inh_ratio=.1,
                     nr_exc_subpopulations=n_exc_sub, reciprocity_exc=.2)()
            net.binarize()
            title_offset = 1.5
            ylabel = "Pre"
            xlabel = "Post"

            def style_axes(ax):
                ax.locator_params(nbins=3)
                #ax.set_title(style.name["LL"], y=title_offset)
                ax.set_xlabel(xlabel)
                ax.set_ylabel(ylabel)
                ax.yaxis.set_ticks_position("left")
                ax.xaxis.set_ticks_position("top")

            with self.plot("n_exc_sub={n}-{k}.pdf".format(n=n_exc_sub, k=k), ncols=2,
                           despine=lambda x: None, sharey=True) as (fig, (ax1, ax2)):
                fig.set_size_inches(2.4, 1.5)
                import matplotlib.cm
                import matplotlib as mpl
                cmap = mpl.cm.get_cmap("RdBu")
                ax2.matshow(net.adjacency_matrix.copy().T, cmap=cmap)




                net.random_relabel_neurons()

                apl1 = ax1.matshow(net.adjacency_matrix.T, cmap=cmap)

                for ax in (ax1, ax2):
                    ax.set_xlim(0, 2000)
                    ax.set_ylim(0, 2000)
                    ax.set_adjustable('box-forced')
                    style.set_ticksmiddle_minor(ax, np.arange(0, net.nr_neurons+1, 200), "x")
                    style.set_ticksmiddle_minor(ax, np.arange(0, net.nr_neurons+1, 200), "y")
                    ax.set_aspect("equal")
                    ax.set_xlabel("Post", labelpad=style.default_label_pad_reduced)
                    sns.despine(ax=ax)

                ax1.set_ylabel("Pre", labelpad=style.default_label_pad_reduced)



                ax2.legend([mpl.patches.Patch(color=cmap(np.inf)), mpl.patches.Patch(color=cmap(-np.inf))],
                           ["e", "i"], loc="center left", bbox_to_anchor=(1, .5))

                fig.tight_layout(w_pad=.4)
