import os

import luigi
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scipy as sp
import seaborn as sns

from abcsmc.loader import ABCLoader, SQLDataStore
from connectome.model import EXP
from figures import memory
from figures import style
import figures.style.linestyle
from figures.analysispipeline.auxiliary import plot, ReplotMixin
from connectome.function import UnsynchronizedActivityTask
from connectome.model import SORN, default_mean_par
from figures.analysispipeline.datatypes import DatasetTask
from ..data import DataBase
from ...style import size
from figures.analysispipeline.config import config
from figures.analysispipeline.auxiliary import FolderPlot

def subs_noise_title(ax, subsampling, noise):
    if subsampling != 1:
        ax.set_title("${}={}, {}={}$".format(style.latex_name("subsampling", dollar=False), subsampling,
                                           style.latex_name("noise", dollar=False), noise), y=.95)
    else:
        if noise != 0:
            ax.set_title("${}={}$".format(style.latex_name("noise", dollar=False), noise), y=1)


def plot_heatmap(ax, data, cbar_label, annot=True):
    cax = ax.figure.add_axes([.93, .4, .1, .2])
    sns.heatmap(data.convert_objects(), annot=annot, annot_kws={'size': config.heatmap_annot_size},
                cmap="Greys", vmin=0, vmax=1, ax=ax, cbar_ax=cax, yticklabels=True,
                xticklabels=True)
    cbar = ax.collections[0].colorbar
    cbar.set_ticks([0, 1])
    cbar.ax.set_ylabel(cbar_label)


def label_axes(ax):
    ax.set_xlabel(style.latex_name.full("noise"))
    ax.set_ylabel(style.latex_name.full("subsampling"))


def plot_confusion_matrix(output_path, LABEL_ROTATION_PARS, plotter):
    for index, confusion_matrix in plotter.confusion_matrix_dict().items():
        title = ", ".join("{}={}".format(prop, label) for prop, label in zip(index.names, index))
        conf_pdf = os.path.join(output_path, f"confusion_{title}.{config.extension}")
        with plot(conf_pdf) as (fig, ax):
            confmat = confusion_matrix.rename(columns=style.name, index=style.name)
            sns.heatmap(confmat.applymap(lambda x: round(x, 2)),
                        annot=False, annot_kws={'size': 5}, vmin=0, vmax=1, cmap="Greys", ax=ax,
                        xticklabels=True,
                        yticklabels=True,
                        square=True)
            cbar = ax.collections[0].colorbar
            style.style_cbar_0_1(cbar, "Posterior probability $p$")
            plt.setp(ax.xaxis.get_majorticklabels(), **LABEL_ROTATION_PARS)
            plt.setp(ax.yaxis.get_majorticklabels(), rotation=0)
            ax.set_xlabel("Estimated model")
            ax.set_ylabel("Original model")
            subsampling = index.subsampling_fraction
            try:
                noise = index.noise
            except AttributeError:
                noise = index.fraction_remove_and_add
            subs_noise_title(ax, subsampling, noise)

            with open(conf_pdf + ".txt", "w") as f:
                f.write(str(confmat))


class PlotPosteriorDistributionMatrix(ReplotMixin, DatasetTask):
    only_matrix = luigi.Parameter(default=False)
    LABEL_ROTATION_PARS = dict(rotation=75, rotation_mode="anchor", horizontalalignment="right")
    MAP_LABEL = "Average\nMAP accuracy"
    AVG_PROB_LABEL = "Average\nprobability"
    matrix_scale = 1.1
    matrix_dim = (2, 1)

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

    def output(self):
        path = os.path.join(config.dataset_output_base_folder(self.input().path) + ".d", "posterior_matrices")
        return luigi.LocalTarget(path)

    def run(self):
        try:
            os.makedirs(self.output().path)
        except FileExistsError:
            pass

        plotter = ABCLoader(SQLDataStore(self.input().fn))


        with plot(os.path.join(self.output().path, f"map_{plotter.noise_type}.{config.extension}"),
                  scale=(1, .5)) as (fig, ax):
            #fig.set_size_inches(self.matrix_dim)
            plot_heatmap(ax, plotter.maximum_a_posteriori(), self.MAP_LABEL)
            label_axes(ax)

        plot_pdf = os.path.join(self.output().path, f"probability_{plotter.noise_type}.{config.extension}")
        with plot(plot_pdf, scale=(1, .5)) as (fig, ax):
            #fig.set_size_inches(self.matrix_dim)
            avg = plotter.average_mass_at_tround_truth()
            plot_heatmap(ax, avg, self.AVG_PROB_LABEL)
            label_axes(ax)


        if not self.only_matrix:
            plot_confusion_matrix(self.output().path, self.LABEL_ROTATION_PARS, plotter)


class PosteriorNoParameterGrouping(ReplotMixin, DatasetTask):
    def requires(self):
        return DataBase(self.dataset)

    def output(self):
        path = os.path.join(config.dataset_output_base_folder(self.input().path) + ".d")
        return luigi.LocalTarget(path)

    def run(self):
        ordered_models = ["ER-ESN", "EXP-LSM", "LAYERED",
                          "SYNFIRE", "FEVER", "API", "STDP-SORN"]
        db = "sqlite:///" + self.input().fn
        abc_smc = pd.read_sql("abc_smc", db)
        populations = pd.read_sql("populations", db)
        models = pd.read_sql("models", db)
        models["name"] = models["name"].map(lambda x: style.name[x])

        gt_model_ids = models.merge(populations, left_on="population_id", right_on="id",
                     suffixes=("", "_population"))[["id", "t"]].query("t == -1")


        models.loc[models.id.isin(gt_model_ids.id),"p_model"] = sp.NAN


        gt = (abc_smc
            .merge(populations[populations.t == -1], left_on="id",
                   right_on="abc_smc_id")
            .merge(models, left_on="id_y", right_on="population_id")
        [["abc_smc_id", "name"]])

        max_t = populations.groupby("abc_smc_id").max()[
            "t"].reset_index().rename(columns={"t": "max_t"})

        joined = (
            populations.merge(max_t, left_on=["abc_smc_id", "t"],
                              right_on=["abc_smc_id", "max_t"])
                .merge(models, left_on="id", right_on="population_id")[
                ["abc_smc_id", "name", "p_model"]])

        confusion = (joined
                .merge(gt, left_on="abc_smc_id", right_on="abc_smc_id",
                       suffixes=("", "_gt"))
                .set_index(["name", "name_gt"])
                .unstack(0)
            ["p_model"]
            [ordered_models].T[ordered_models].T)


        with plot(os.path.join(self.output().path,
                               f"{self.dataset}.{config.extension}"),
                  scale=1) as (fig, ax):
            # fig.set_size_inches(self.matrix_dim)
            sns.heatmap(confusion, vmin=0, vmax=1, cmap="Greys", ax=ax,
                        square=True,
                        xticklabels=True,
                        yticklabels=True,
                        )
            ax.set(xlabel="Posterior estimate", ylabel="Original model")
            style.style_cbar_0_1(ax.collections[0].colorbar, "$p$")
            plt.setp(ax.xaxis.get_majorticklabels(),
                     **PlotPosteriorDistributionMatrix.LABEL_ROTATION_PARS)



class OneDPosteriors(ReplotMixin, DatasetTask):
    #force_replot = True
    posterior_models = luigi.Parameter()
    y_lim_min_log_scale = 10 ** -3
    y_lim_max = 1.05
    half_width = .4

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

    def output(self):
        path = os.path.join(self.input().fn + ".d", "1d_posterior")
        return luigi.LocalTarget(path)


    def run(self):
        try:
            os.makedirs(self.output().fn)
        except FileExistsError:
            pass

        plotter = ABCLoader(SQLDataStore(self.input().fn))


        results = plotter.results()
        models = list(plotter.model_names)

        assert len(results.subsampling_fraction.unique()) == 1, "Subsampling not well defined"
        subsampling = results.subsampling_fraction[0]


        try:
            assert len(results.noise.unique()) == 1
            noise = results.noise.unique()[0]
        except AttributeError:
            assert len(results.fraction_remove_and_add.unique()) == 1
            noise = results.fraction_remove_and_add.unique()[0]


        posterior_models_to_plot = self.posterior_models
        results_model_filtered = results[results.gt_model_name.isin(posterior_models_to_plot)].copy()
        results_model_filtered["Type"] = results.gt_model_name.map(lambda x: x)
        posterior = results_model_filtered[models + ["Type"]]
        prior = pd.DataFrame({model: [1 / len(models)] for model in models})
        prior["Type"] = "Prior"
        distribution = posterior.append(prior)

        for ax_scale in ["log", "linear"]:
            with plot(os.path.join(self.output().path, f"1d_posterior_{ax_scale}.{config.extension}"),
                      scale = style.size.m * (.8, .65)) as (fig, ax):
                x_pos_left = np.arange(len(models))
                x_lim = -self.half_width, x_pos_left[-1] + 3 * self.half_width
                # All posteriors
                legend_items = []
                legend_labels = []
                for index, data in distribution.iterrows():
                    y = list(data[models])
                    if data.Type != "Prior":
                        color = style.color[data.Type]
                        linewidth = style.bar_line_width
                        ax.bar(x_pos_left, y, facecolor=(1, 1, 1, 0),
                               edgecolor=color, linewidth=linewidth)
                        legend_items.append(plt.Line2D((0,0), (0,0),
                                            color=color, linewidth=linewidth))
                        legend_labels.append(style.name[data.Type])

                # The prior
                prior_values = distribution[distribution.Type == "Prior"][models].iloc[0]
                prior_pos = x_pos_left.copy().astype(float)
                prior_pos[0] = x_pos_left[0]
                prior_pos[-1] = x_pos_left[-1] + 2 * self.half_width
                color = style.color["Prior"]

                ax.plot(prior_pos, prior_values, dashes=style.linestyle.dashed, color=color)
                legend_items.append(plt.Line2D((0,0), (0,0), dashes=style.linestyle.dashed, color=color))
                legend_labels.append("Prior")
                ax.legend(legend_items, legend_labels, loc="center left", bbox_to_anchor=(1, .5),
                          title="Posteriors")
                ax.set_yscale(ax_scale)
                y_lim_min = self.y_lim_min_log_scale if ax_scale == "log" else 0
                ax.set_ylim(y_lim_min, self.y_lim_max)
                ax.set_xticks(x_pos_left + self.half_width)
                ax.set_xticklabels([style.name[m] for m in models])
                ax.set_xlim(*x_lim)
                ax.set_ylabel("$p$")
                subs_noise_title(ax, subsampling, noise)
                ax.set_yticks([0, 1])
                style.reduced_ticks_0_1(ax, axis="y")
                plt.setp(ax.xaxis.get_majorticklabels(), **PlotPosteriorDistributionMatrix.LABEL_ROTATION_PARS)
                #zsns.despine(ax=ax)


class ModifiedExp(ReplotMixin, DatasetTask):
    def requires(self):
        return DataBase(self.dataset)

    def output(self):
        return {"bar": luigi.LocalTarget(os.path.join(config.dataset_output_base_folder(self.input().path) + ".d", f"barplot.{config.extension}")),
                "stair": luigi.LocalTarget(os.path.join(config.dataset_output_base_folder(self.input().path) + ".d", f"stair.{config.extension}")),
                "decay": luigi.LocalTarget(os.path.join(config.dataset_output_base_folder(self.input().path) + ".d", f"decay.{config.extension}"))}

    def run(self):
        self.output()["bar"].makedirs()
        loader = ABCLoader(SQLDataStore(self.input().path))
        loader.group_parameters_from_noise = []
        loader.group_parameters_from_pars = ["decay"]
        res = loader.results()
        models = ["API", "ER", "ERFEVER", "EXP", "LL", "SORN", "SYN"]
        res = res[models + ["decay"]]
        sums = res[models + ["decay"]].sum(axis=0)
        nonzero = sums[sums != 0].index
        res = res[nonzero]

        with plot(self.output()["bar"].path) as (fig, ax):
            ax = res.rename(columns=style.name).plot(x="decay", kind="bar",
                                              stacked=True,
                                              color=[style.color[name] for name in nonzero],
                                              linewidth=0,
                                              edgecolor="none",
                                              ax=ax)
            ax.legend(loc="center left", bbox_to_anchor=(1, .5))
            ax.set_xlabel(style.latex_name("decay"), labelpad=style.default_label_pad_reduced)
            ax.set_ylabel("Posterior", labelpad=style.default_label_pad_reduced)
            ax.set_yticks([0, .2, .4, .6, .8, 1])
            ax.set_yticklabels([0, "", "", "", "", 1])
            style.reduced_ticks_0_1(ax, axis="y")
            ax.set_xticklabels([res.decay.iloc[0]] + [""] * (len(res) - 2) + [res.decay.iloc[-1]])
            plt.setp(ax.xaxis.get_majorticklabels(), rotation=0)
            height = -.4
            ax.text(0, height, "ER ←", ha="center", transform=ax.transAxes)
            ax.text(1, height, "→ EXP", ha="center", transform=ax.transAxes)
            sns.despine(ax=ax)

        with plot(self.output()["stair"].path) as (fig, ax):
            for name in nonzero:
                if name != "decay":
                    ax.step(res.decay, res[name], color=style.color[name], where="mid")
            ax.legend(loc="center left", bbox_to_anchor=(1, .5))
            ax.set_xlabel(style.latex_name("decay"))
            ax.set_ylabel("$p$")
            ax.set_title("Posterior distribution")
            plt.setp(ax.xaxis.get_majorticklabels(), rotation=90)
            sns.despine(ax=ax)

        distances = sp.linspace(0, 1)
        decays = sp.arange(.1, 1.05, .1)

        @memory.cache()
        def decay_line(decay, distance):
            exp = EXP(nr_neurons=2000, inh_ratio=.1, p_exc=.2, p_inh=.6, decay=decay)
            exp()
            p = exp.connection_probability(distance, exp.alpha_exc_, exp.p_0_exc)
            return p

        ps = map(lambda x: decay_line(x, distances), decays)

        with_labels = decays[::3]
        with plot(self.output()["decay"].path) as (fig, ax):
            ax.plot([0, 1], [.2, .2],
                    label="0 ER (ESN)",
                    color=style.cmap(0))
            for p, decay in zip(ps, decays):
                ax.plot(distances, p,
                        label="{:.1f}".format(decay) + (" EXP (LSM)" if decay == 1 else "")
                        if decay in with_labels else None,
                        color=style.cmap(decay))
                ax.set_xlabel(style.latex_name.long("distance"))
                ax.set_ylabel("$ " + style.latex_name("p_exc", dollar=False) + ", " + style.latex_name("p_inh", dollar=False) + " $",
                              labelpad=style.default_label_pad_reduced)
                style.reduced_ticks_0_1(ax, axis="both")

                sns.despine(ax=ax)
            fig2, ax2 = plt.subplots()
            cbar_proxy = ax2.matshow(sp.zeros((0, 0)), vmin=0, vmax=1)
            cbar = fig.colorbar(cbar_proxy)
            cbar.outline.set_visible(False)
            cbar.set_ticks([0, 1])
            cbar.set_ticklabels(["0 " + style.name["ER"], "1 " + style.name["EXP"]])
            cbar.set_label(style.latex_name("decay"))


@memory.cache
def run_sorn(nr_patterns):
    sorn = SORN(nr_neurons=2000, inh_ratio=.1, p_exc=.2, p_inh=.6,
                eta_stdp=default_mean_par("SORN", "eta_stdp"),
                eta_intrinsic=default_mean_par("SORN", "eta_intrinsic"),
                nr_patterns=nr_patterns)

    task = UnsynchronizedActivityTask()
    res = task(sorn)
    return res


class ModifiedSORN(ReplotMixin, DatasetTask):
    MAX_ALLOWED_CORRELATION = 1

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

    def output(self):
        return {"bar": luigi.LocalTarget(os.path.join(config.dataset_output_base_folder(self.input().path) + ".d", f"barplot.{config.extension}")),
                "median_correlation": luigi.LocalTarget(os.path.join(
                    config.dataset_output_base_folder(
                        self.input().path) + ".d",
                    f"median_correlation.{config.extension}"))}

    def run(self):
        self.output()["bar"].makedirs()
        loader = ABCLoader(SQLDataStore(self.input().fn))
        loader.group_parameters_from_noise = []
        loader.group_parameters_from_pars = ["nr_patterns"]
        res = loader.results()
        models = ["API", "ER", "ERFEVER", "EXP", "LL", "SORN", "SYN"]
        res = res[models + ["nr_patterns"]]
        sums = res[models + ["nr_patterns"]].sum(axis=0)
        nonzero = sums[sums != 0].index
        res = res[nonzero]

        correlations = sp.array([run_sorn(int(nr_patterns))["correlation"] for nr_patterns in res.nr_patterns])
        res = res[correlations < self.MAX_ALLOWED_CORRELATION]
        res["nr_patterns"] = res["nr_patterns"].apply(lambda x: int(x) if x != -1 else "$\infty$")

        with plot(self.output()["bar"].path,
                  despine=lambda x: None, sharex=True, scale=.7) as (fig, ax):
            ax = res.rename(columns=style.name).plot(x="nr_patterns",
                                                     kind="bar",
                                              stacked=True,
                                              color=[style.color[name] for name in nonzero],
                                              linewidth=0,
                                              edgecolor="none",
                                              ax=ax)
            ax.legend(loc="upper center", bbox_to_anchor=(.5, 1.5))
            ax.set_xlabel(style.latex_name("nr_patterns"))
            ax.set_ylabel("Posterior", labelpad=style.default_label_pad_reduced)
            ax.set_yticks([0, .2, .4, .6, .8, 1])
            ax.set_yticklabels([0, "", "", "", "", 1])
            plt.setp(ax.xaxis.get_majorticklabels(), rotation=60)
            sns.despine(fig)

        with plot(self.output()["median_correlation"].path,
                  despine=lambda x: None, sharex=True, scale=.7) as (fig, ax2):
            ax2.plot(correlations, ".-")
            ax2.set_xticks(sp.arange(correlations.size))
            ax2.set_xticklabels(res["nr_patterns"])
            ax2.set_ylabel(style.latex_name("median_correlation"))
            ax2.set_xlabel(style.latex_name("nr_patterns"))
            plt.setp(ax2.xaxis.get_majorticklabels(), rotation=60)
            sns.despine(fig)


class CombinedPosteriorDistributionMatrix(DatasetTask, FolderPlot):
    def requires(self):
        return [DataBase(d) for d in self.dataset]

    def run(self):
        """
        Warning
        ------

        The code here assumes that adding the content of the ABC databass
        per noise/subsampling and averaging makes sense.
        This assumes that all the files have the same number of runs
        per noise/subsampling pair.
        """
        plotters = [ABCLoader(SQLDataStore(inp.path)) for inp in self.input()]
        probabilities = [p.average_mass_at_tround_truth() for p in plotters]
        maps = [p.maximum_a_posteriori() for p in plotters]

        p_agg = pd.DataFrame()
        cnt = pd.DataFrame()
        for p in probabilities:
            p_agg = p_agg.add(p, fill_value=0)
            cnt = cnt.add(p != 0, fill_value=0)

        map_agg = pd.DataFrame()
        for m in maps:
            map_agg = map_agg.add(m, fill_value=0)

        with self.plot("map", scale=(1.3, .5)) as (fig, ax):
            plot_heatmap(ax, map_agg / cnt, PlotPosteriorDistributionMatrix.MAP_LABEL)
            label_axes(ax)

        with self.plot("probability", scale=(1.3, .5)) as (fig, ax):
            plot_heatmap(ax, p_agg / cnt, PlotPosteriorDistributionMatrix.AVG_PROB_LABEL)
            label_axes(ax)


class AverageConfusion(ReplotMixin, luigi.Task):
    runs = luigi.Parameter()
    output_folder = luigi.Parameter()

    def output(self):
        return luigi.LocalTarget(os.path.join(os.environ["PLOT_PATH"],
                                              self.output_folder + "-average"))

    def run(self):
        data = [ABCLoader(SQLDataStore(os.path.join(os.environ["PLOT_PATH"], r)))
                for r in self.runs]
        conf_mat_dct = [d.confusion_matrix_dict() for d in data]

        for c in conf_mat_dct:
            assert len(c) == 1, "More than one run in file"

        conf_mat = [list(c.values())[0] for c in conf_mat_dct]

        avg = sum(conf_mat) / len(conf_mat)

        try:
            os.makedirs(self.output().path)
        except FileExistsError:
            pass

        order = ["ER", "EXP", "LL", "SYN", "ERFEVER", "API", "SORN"]

        with plot(os.path.join(self.output().path, "average." + config.extension)) as (fig, ax):
            sns.heatmap(avg.loc[order][order], ax=ax, cmap="Greys", vmin=0, vmax=1, xticklabels=True,
                        yticklabels=True)
            plt.setp(ax.xaxis.get_majorticklabels(), rotation=-90)

        avg.to_excel(os.path.join(self.output().path, "average.xlsx"))
