import itertools
import os
from itertools import product

import luigi
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import scipy as sp

from connectome.model.ll import connectivity, reciprocity
from figures import style
import figures.style.size as size
from figures.analysispipeline.auxiliary import plot
from figures.analysispipeline.data import APIOriginalVsModifiedData
from ..data import DataFrame
from ..data import FEVEROriginalVsModifiedData
from ..folder import OutputFolder
import seaborn as sns
from ..auxiliary import FolderPlot

from .poverd import ExpConnectivityOverDistance
viridis = mpl.cm.get_cmap("viridis")


def make_feature_dim_legend(ax, df, original_marker, modified_marker, color_dim):
    pseudo_data = ((1, 1),) * 2
    legend_items = [plt.Line2D(*pseudo_data, linestyle="none", marker=original_marker, color="k"),
                    plt.Line2D(*pseudo_data, linestyle="none", marker=modified_marker, color="k")]
    legend_names = ["Original", "Modified"]
    fdims = sorted(df.feature_space_dimension.unique())
    if len(fdims) > 6:
        fdims = fdims[::2]
    for dim in fdims:
        legend_names.append("${}={}$".format(style.latex_name("feature_space_dimension", dollar=False), int(dim)))
        legend_items.append(mpl.patches.Patch(color=color_dim(dim)))
    ax.legend(legend_items, legend_names, loc="center left", bbox_to_anchor=(1, .5))


def color_dim_factory(df):
    min_f, max_f = df.feature_space_dimension.min(), df.feature_space_dimension.max()

    def color_dim(dim):
        return viridis((dim - min_f) / (max_f - min_f))

    return color_dim


api_fever_marker_size = 10


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

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

    def run(self):
        from connectome.analysis import standard_analysis
        from connectome.model import ER

        def make_er(nr_neurons, inh_ratio, p_exc, p_inh):
            network = ER(nr_neurons=nr_neurons, inh_ratio=inh_ratio, p_exc=p_exc, p_inh=p_inh)()
            return standard_analysis(network)

        df = pd.DataFrame([make_er(2000, .1, float(p_exc), .6) for p_exc in sp.linspace(0, 1)])
        self.output().save(df)


class ERConnectivityReciprocityScatter(FolderPlot):
    scatter_point_size = 3
    def requires(self):
        return ERConnectivityReciprocityScatterData()

    def run(self):
        import seaborn as sns
        df = self.input().load()
        for with_annot, suffix in [(True, ""), (False, "_no_annotation")]:
            with self.plot("er_r_over_p" + suffix, scale=ExpConnectivityOverDistance.SCALE) as (fig, ax):
                df = df[df.p_ee < .5]
                ax.scatter(df.p_ee, df.reciprocity_ee, facecolor=sns.get_color_cycle()[0], edgecolor="w",
                           s=self.scatter_point_size)
                ax.locator_params(nbins=1)
                ax.set_xlabel(style.latex_name("p_ee"))
                ax.set_ylabel(style.latex_name("reciprocity_ee"))
                ax.set_xlim(0, 1)
                ax.set_ylim(0, 1)
                style.style_r_over_p(ax, "half")
                x_start = .55
                if with_annot:
                    ax.annotate(xy=(.25, .2), xytext=(x_start, .1),
                                s=style.latex_name.long("constraints"),
                                arrowprops=style.annotarrowprops, va="center")
                    ax.annotate(xy=(.32, .28), xytext=(x_start, .4), s=style.name["ER"],
                                arrowprops=style.annotarrowprops, va="center")


class APIOrginalVsAdaptedPlot(FolderPlot):
    SCALE = size.s

    def requires(self):
        return APIOriginalVsModifiedData()

    def run(self):
        df = self.input().load()

        s = .5

        measures = ["feature_space_dimension", "n_pow"]
        for with_title in [True, False]:
            for filter_, suffix, orig_annot, adap_annot in [(lambda x: x, "", True, True),
                                                           (lambda x: x[x.name == "SimpleMiller"], "orig_only", False, False),
                                                           (lambda x: x[x.name == "API"], "adap_only", False, False)]:
                for measure in measures:
                    title_str = "_title" if with_title else "_notitle"
                    with self.plot(measure + "_"  + title_str + "_" + suffix, scale=self.SCALE) as (fig, ax):
                        cbar = style.two_d_measure_scatter(filter_(df), ax, fig,
                                                    "p_ee", "reciprocity_ee", measure, lambda ax, label: ax.set_label(label),
                                                    vmin=sp.floor(df[measure].min()),
                                                    vmax=sp.ceil(df[measure].max()), s=s)

                        cbar.set_ticks([sp.floor(df[measure].min()), sp.ceil(df[measure].max())])
                        cbar.set_label(style.latex_name(measure), labelpad=style.default_label_pad_reduced)


                        style.style_r_over_p(ax)

                        x_text_pos = .41
                        if orig_annot:
                            ax.text(x_text_pos, .95, style.latex_name.long("original"))
                        if adap_annot:
                            ax.text(x_text_pos, .6, style.latex_name.long("adapted"))



class FEVEROriginalVsModifiedPLot(FolderPlot):
    def requires(self):
        return FEVEROriginalVsModifiedData()

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

        color_dim = color_dim_factory(df)
        df["color"] = df["feature_space_dimension"].apply(color_dim)

        modified_marker = "o"
        original_marker = "^"
        df["marker"] = df["name"].apply(lambda x: modified_marker if x == "ERFEVER" else original_marker)



        for measure in ["feature_space_dimension"]:
            for kind in ["", "orig_only", "adap_only"]:
                with self.plot(measure + "joint_scatter" + kind, scale=style.size.m) as (fig, ax):
                    max_val = 200 if measure == "feature_space_dimension" else 1
                    style.style_r_over_p(ax, view="dot6")
                    df_orig = df[df["name"] == "ERFEVERInitialConnectivity"]
                    df_adap = df[df["name"] == "ERFEVER"]
                    if kind == "":
                        df_plot = df
                    elif kind == "orig_only":
                        df_plot = df_orig
                    elif kind == "adap_only":
                        df_plot = df_adap
                    else:
                        raise Exception("Problem")
                    cbar = style.two_d_measure_scatter(df_plot, ax, fig,
                                                "p_ee", "reciprocity_ee", measure,
                                                lambda ax, label: ax.set_label(label),
                                                vmin=0, vmax=max_val,
                                                s=.5, tight_ax_lims=False,
                                                labelpad=style.default_label_pad_reduced-5)
                    cbar.set_ticks([0, max_val])
                    cbar.set_label(style.latex_name(measure), labelpad=style.default_label_pad_reduced-10)


                    arrow = {"arrowstyle": "-|>"}
                    if kind == "":
                        ax.annotate(xy=(.08, .52), xytext=(.25, .55), s=style.latex_name.long("original"), ha="left", va="center",
                                arrowprops=arrow)

                        ax.annotate(xy=(.3, .33), xytext=(.35, .13), s=style.latex_name.long("adapted"), ha="left", va="center",
                                    arrowprops=arrow)


        def at_single_value_plot(df, x, y, hue, title):
            with self.plot(x) as (fig, ax):
                max_val = 200 if x == "feature_space_dimension" else 1
                vmin, vmax = 0, .25
                mappable = ax.scatter(x, y, c=hue, data=df, edgecolor="w", vmin=vmin, vmax=vmax)
                cbar = style.colorbar(ax, mappable, style.latex_name(hue), labelpad=style.default_label_pad_reduced)
                cbar.set_ticks([vmin, vmax])
                ax.set_xlabel(style.latex_name(x), labelpad=style.default_label_pad_reduced)
                ax.set_ylabel(style.latex_name(y), labelpad=style.default_label_pad_reduced)
                sns.despine(ax=ax)
                #ax.locator_params(nbins=4)
                ax.set_xlim(0, max_val)
                ax.set_xticks(ax.get_xlim())
                ax.set_title(title, y=1.05)
                ax.set_ylim(0, .5)
                ax.set_yticks(ax.get_ylim())
                ax.set_yticks([.1,.2,.3,.4], minor=True)

        df = modified[(modified.p_ee < .22) & (modified.p_ee > .18) & (modified.feverization_ratio == 0.7)]
        at_single_value_plot(df, "feature_space_dimension", "reciprocity_ee", "p_ee",
                             "${} \in [0.18, 0.22], {}=0.7$".format(style.latex_name("p_ee", dollar=False),
                                                                 style.latex_name("feverization_ratio", dollar=False)))


        df = modified[(modified.p_ee < .22) & (modified.p_ee > .18) &( modified.feature_space_dimension == 120)]
        at_single_value_plot(df, "feverization_ratio", "reciprocity_ee", "p_ee",
                             "${} \in [0.18, 0.22], {}=120$".format(style.latex_name("p_ee", dollar=False),
                                                                    style.latex_name("feature_space_dimension", dollar=False))
                             )





def my_palplot(ax, pal, size=1):
    """Plot the values in a color palette as a horizontal array.

    Parameters
    ----------
    pal : sequence of matplotlib colors
        colors, i.e. as returned by seaborn.color_palette()
    size :
        scaling factor for size of plot

    """
    import numpy as np
    n = len(pal)
    ax.imshow(np.arange(n).reshape(n, 1),
              cmap=mpl.colors.ListedColormap(list(pal)),
              interpolation="nearest", aspect="auto")
    ax.set_xticks(np.arange(n) - .5)
    ax.set_yticks([-.5, .5])
    ax.set_xticklabels([])
    ax.set_yticklabels([])


HEATMAP_SCALE = .6

class LLROverP(FolderPlot):
    def run(self):
        min_nr_layers, max_nr_layers = 2, 9
        with self.plot("LL_r_over_p", scale=(1, 1)) as (fig, ax):
            colors = []
            for nr_layers in list(range(min_nr_layers, max_nr_layers+1)):
                p_min = 0
                p_hidden_range = sp.linspace(p_min, 1, 20)
                p_range = (  list(itertools.product([p_min], p_hidden_range))
                           + list(itertools.product(p_hidden_range, [1]))
                           + list(itertools.product([1], sp.flipud(p_hidden_range)))
                           + list(itertools.product(sp.flipud(p_hidden_range), [p_min]))
                           )


                LL = [[connectivity(nr_layers, *val), reciprocity(nr_layers, *val)] for val in p_range]
                # plot
                color = mpl.cm.viridis((nr_layers-min_nr_layers)/(max_nr_layers-min_nr_layers))
                colors.append(color)
                ax.fill(*zip(*LL), facecolor=color,
                        edgecolor="k", linewidth=.5, label="$n_l={}$".format(nr_layers))
                corner_points = list(product([1, p_min], repeat=2))
                corner_conn_rec = [(connectivity(nr_layers, *point), reciprocity(nr_layers, *point))
                                   for point in corner_points]

            style.make_plausible_regime_r_over_p_rectangle(ax, color="white")
            ax.set_xlim(0, 1)
            ax.set_ylim(0, 1)
            ax.set_ylabel(style.latex_name("reciprocity_ee"))
            ax.set_xlabel(style.latex_name("p_ee"))
            ax.locator_params(nbins=1)
            style.reduced_ticks_0_1(ax)
            ax2 = fig.add_axes([.78, .15, .08, .7])
            style.reduced_ticks_0_1(ax)

            labels = list(range(2, max_nr_layers+1))
            labels[::2] = [""] * len(labels[::2])
            style.make_cbar(ax2, colors, labels, style.latex_name("nr_exc_subpopulations"))


