import json
import os

from ..auxiliary import FolderPlot
import scipy as sp
from scipy import stats as st

from figures import  style
from figures.analysispipeline.auxiliary import plot
import matplotlib.pyplot as plt
from figures.analysispipeline.config import config
import seaborn as sns

xlimits_dict = {
            # format is min, max, minor_step_size
            ("API", "feature_space_dimension"): (0, 70, 10),
            ("API", "n_pow"): (0, 8, 1),
            ("ERFEVER", "feature_space_dimension"): (0, 250, 50),
            ("ERFEVER", "feverization_ratio"): (0, 1, .1),
            ("EXP", "decay"): (0, 2, .2),
            ("joint", "inh_ratio"): (0, 1, .1),
            ("joint", "nr_neurons"): (0, 2500, 100),
            ("joint", "p_exc"): (0, 1, .1),
            ("joint", "p_inh"): (0, 1, .1),
            ("LL", "nr_exc_subpopulations"): (0, 5, 1),
            ("LL", "reciprocity_exc"): (0, 1, .1),
            ("model", "model"): (0, 6, 1),
            ("RemoveAddNoiseAndSubsample", "fraction_remove_and_add"): (0, 1, .1),
            ("RemoveAddNoiseAndSubsample", "subsampling_fraction"): (0, 1, .1),
            ("SORN", "eta_intrinsic"): (0, .4, .1),
            ("SORN", "eta_stdp"): (0, 0.002, .0001),
            ("SORN", "nr_patterns"): (-5000, 12000, 1000),
            ("SYN", "pool_size"): (0, 400, 100)
        }

override_ticks = {
    ("model", "model"): ((0, 6), ("ER", "API")),
}

def reduced_ticks(ax, axis="both", despine=True, label_pad_reduce=-5):
    if despine:
        sns.despine(ax=ax)

    if axis == "x" or axis == "both":
        config.modify_labelpad(ax.xaxis, label_pad_reduce)

    if axis == "y" or axis == "both":
        config.modify_labelpad(ax.yaxis, label_pad_reduce)
        ax.set_ylim(bottom=0)


def get_distribution(parameter_json):
    """
    Try to get the distribution if possible.
    If the value is a fixed nr, return an approximate delta distribution
    centered at that numbber. This should look good enough for plotting purposes.
    """
    try:
        dist = getattr(st, parameter_json["type"])(*parameter_json["args"])
    except TypeError:
        dist = st.uniform(parameter_json, .001)
    return dist


def plot_dist(ax, model_name, parameter_name, dist):
    *x_limits, x_step_size = xlimits_dict[(model_name, parameter_name)]
    x = sp.linspace(*x_limits, num=1000)

    if isinstance(dist.dist, st.rv_discrete):
        if dist.dist.name == 'randint':
            x = [x_limits[0], dist.args[0], dist.args[1], x_limits[1]]
            x = sp.asarray(x)
            y = dist.pmf(x)

            x = x - 0.5
            x = sp.hstack(([x[0]], sp.repeat(x[1:-1], 2), [x[-1]]))
            y = sp.repeat(y[:-1], 2)
        else:
            assert(False, 'Code path not tested')
            y = dist.pmf(x)

    elif isinstance(dist.dist, st.rv_continuous):
        if dist.dist.name == 'uniform':
            eps = dist.args[1] / 1000
            x = [x_limits[0], dist.args[0],
                 dist.args[0] + dist.args[1] + eps, x_limits[1]]
            x = sp.asarray(x)
            y = dist.pdf(x)

            x = sp.hstack(([x[0]], sp.repeat(x[1:-1], 2), [x[-1]]))
            y = sp.repeat(y[:-1], 2)
        else:
            assert(False, 'Code path not tested')
            y = dist.pdf(x)

    ax.fill_between(x, y, color=config.maybe_to_black(style.color[model_name]))

    try:
        ax.get_xaxis().get_major_formatter().set_useOffset(False)
    except AttributeError:
        pass

    y_limits = y[sp.logical_not(sp.isinf(y))]

    if y_limits.size == 0:
        y_limits = 1.0
    else:
        y_limits = y_limits.max()

    y_step = 10 ** sp.floor(sp.log10(y_limits))
    y_limits = y_step * sp.ceil(y_limits / y_step)
    y_limits = (0, y_limits)

    ax.set_ylim(y_limits)
    ax.set_ylabel(style.name["PDF"])
    reduced_ticks(ax, axis="y")

    ax.set_xlim(x_limits)
    ax.set_xticks(x_limits)
    ax.set_xticks(sp.arange(x_limits[0], x_limits[1] + x_step_size // 10, x_step_size), minor=True)

    if (model_name, parameter_name) in override_ticks:
        ticks, labels = override_ticks[(model_name, parameter_name)]
        ax.set_xlim(ticks)
        ax.set_xticks(ticks)
        ax.set_xticklabels(labels)
        ax.set_xticks(sp.arange(ticks[0], ticks[1] + 1), minor=True)

    # no break in the noise CDF in fig 5
    delimiter = " " if len(style.latex_name.long(parameter_name)) < 6 else " \n "
    ax.set_xlabel(style.latex_name.full(parameter_name))


def plot_parameter(ax, kind, name, parameter):
    dist = get_distribution(parameter)
    plot_dist(ax, kind, name, dist)


class SubPlotPlotContext:
    extension = config.extension

    def __init__(self, path, nrows=1, ncols=1):
        self.path = path
        self.nrows = nrows
        self.ncols = ncols
        self.fig, self.axes = plt.subplots(nrows=nrows, ncols=ncols)
        self.axes_iter = iter(self.axes.flatten())

    def __call__(self, *args, **kwargs):
        self.current_ax = next(self.axes_iter)
        return self

    def __enter__(self):
        return self.fig, self.current_ax

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    def set_title(self, row, col, title):
        self.axes[row, col].set_title(title + " " * 3,
                                      loc="left",
                                      fontdict={'fontsize': 12,
                                                'fontweight': "bold",
                                                'horizontalalignment': 'right'})

    def save(self, *args, **kwargs):
        for ax in self.axes_iter:
            ax.axis("off")
        for ax in self.axes[:,1:].flatten():
            ax.set_ylabel("")

        self.fig.tight_layout()
        self.fig.savefig(self.path + "." + self.extension, bbox_inches="tight", transparent=True, *args, **kwargs)


class Priors(FolderPlot):
    """
    The prior plots in the ABC posterior figures should not have any title.
    """

    scale = .5

    @property
    def pars(self):
        from connectome.model import default_definitions
        pars = default_definitions()
        del pars["models"]["model_specific_parameters"]["LL"]["reciprocity_exc"]
        del pars["models"]["model_specific_parameters"]["EXP"]["decay"]
        del pars["models"]["model_specific_parameters"]["SORN"]["nr_patterns"]
        return pars


    def run(self):
        self.make_folder()

        self.plot_model_classes(self.plot)
        self.plot_shared(self.plot)
        self.plot_noise(self.plot)
        self.plot_specific(self.plot)
        self.plot_delta_0(self.plot)

        shared_con = SubPlotPlotContext(os.path.join(self.output_folder(), "combined_priors"), 2, 7)
        shared_con.fig.set_size_inches(7.5, 1.2 * 2)
        self.plot_model_classes(shared_con)
        self.plot_shared(shared_con)
        self.plot_specific(shared_con)
        shared_con.fig.tight_layout(h_pad=1, w_pad=1.5)
        shared_con.set_title(0, 0, "a")
        shared_con.set_title(0, 1, "b")
        shared_con.set_title(0, 5, "c")
        shared_con.set_title(0, 6, "d")
        shared_con.set_title(1, 0, "e")
        shared_con.set_title(1, 2, "f")
        shared_con.set_title(1, 4, "g")
        shared_con.save()


    def plot_specific(self, plot_context):
        for model_name in style.ordered_models:
            parameters = self.pars["models"]["model_specific_parameters"][model_name]
            for parameter_name, par_value in parameters.items():
                with plot_context(model_name + "_" + parameter_name, scale=self.scale) as (fig, ax):
                    plot_parameter(ax, model_name, parameter_name, par_value)

    def plot_shared(self, plot_context):
        for parameter_name, value in self.pars["models"]["joint_parameters"].items():
            model_name = "joint"
            with plot_context(model_name + "_" + parameter_name, scale=self.scale) as (fig, ax):
                plot_parameter(ax, model_name, parameter_name, value)

    def plot_noise(self, plot_context):
        for parameter_name, value in self.pars["noise"]["parameters"].items():
            model_name = self.pars["noise"]["type"]
            for with_indicator in [True, False]:
                suffix = "_indicator" if with_indicator == True else ""
                with plot_context(model_name + "_" + parameter_name + suffix, scale=self.scale) as (fig, ax):
                    plot_parameter(ax, model_name, parameter_name, value)
                    if with_indicator:
                        ax.axvline(0.15, linestyle="dashed", color=sns.utils.get_color_cycle()[4])  # 0.15 is the selected noise level

    def plot_delta_0(self, plot_context):
        for with_indicator in [True, False]:
            suffix = "_indicator" if with_indicator == True else ""
            with plot_context("delta_0" + suffix, scale=self.scale) as (fig, ax):
                ax.fill_between([-1, 0, 0, 1], [0, 0, 1, 1],
                                color=config.maybe_to_black(style.color["RemoveAddNoiseAndSubsample"]))
                ax.set_ylim([0, 1.05])
                ax.get_xaxis().get_major_formatter().set_useOffset(False)
                reduced_ticks(ax, axis="y")
                ax.set_ylabel(style.name["CDF"])
                ax.set_xlim(-1, 1)
                ax.set_xticks([-1, 0, 1])
                ax.set_xlabel(style.latex_name.full("noise"))
                if with_indicator:
                    ax.axvline(0.15, linestyle="dashed",
                               color=sns.utils.get_color_cycle()[4])  # 0.15 is the selected noise level

    def plot_model_classes(self, plot_context):
        with plot_context("model" + "_" + "model", scale=self.scale) as (fig, ax):
            nr_models = len(self.pars["models"]["model_specific_parameters"])
            model_prior = st.randint(0, nr_models)
            plot_dist(ax, "model", "model", model_prior)


