import os

import luigi
import scipy as sp
import seaborn as sns
from figures.analysispipeline.config import config
import figures.style.linestyle
from figures.analysispipeline.datatypes import DatasetTask
from figures.analysispipeline.auxiliary import ReplotMixin
import figures.style.size as size
from connectome.analysis import ConnectivityEstimator, ReciprocityEstimator
from connectome.model import EXP
from ..data import DataBase
from .model_eval_data import model_eval_data
from connectome.model import SYN
from connectome.model.syn import _nr_pools
from figures import  memory
from figures import style
from ..data import FEVEROriginalVsModifiedRuleSatisfaction
from figures.analysispipeline.auxiliary import plot, FolderPlot
from figures.analysispipeline.folder import OutputFolder
from ..datatypes import JSON
from matplotlib.lines import Line2D
from connectome.model.syn import PoolConnector
from connectome.model import Network
from connectome.analysis import standard_analysis
import pandas as pd
import matplotlib.pyplot as plt
from ..auxiliary import  no_despine
from collections import namedtuple
import matplotlib as mpl
from functools import partial
from connectome.function.task.statistics import pair_correlation
from itertools import cycle
from .posterior import PlotPosteriorDistributionMatrix
from mpl_toolkits.mplot3d import Axes3D  #### this is actually neeeded to have 3d plotting working althoug it seems unused
import numpy as np
import figures.style.formatter
from .originalvsmodified import APIOrginalVsAdaptedPlot
from .poverd import ExpConnectivityOverDistance
from .originalvsmodified import ERConnectivityReciprocityScatter
import scipy.linalg as la

setminus_unicode = "∖"  # this looks like an empty string in pycharm but it is not empty

LL_SYN_HEATMAP_SCALE = config.large_heatmap_scale

def analyse(net):
    p_ee = ConnectivityEstimator(network=net)()['p_ee']
    r_ee = ReciprocityEstimator(network=net)()['reciprocity_ee']
    return p_ee, r_ee


def despine(ax):
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')


class SynConnectivityReciprocityNrPools(FolderPlot):
    SCALE =.9


    def run(self):
        sp.random.seed(42)
        p_inh = .6
        syn = SYN(nr_neurons=2000, inh_ratio=.1, p_inh=p_inh)

        def transpose(results):
            return [sp.asarray(res) for res in zip(*results)]

        def make(connectivity, pool_size):
            net = syn(p_exc=connectivity, pool_size=pool_size)
            return analyse(net)

        @memory.cache
        def syn_pool_size_sweep(connectivity, step_size):
            pool_sizes = sp.arange(80, 300, step_size)
            p_ee, r_ee = transpose([make(connectivity, pool_size) for pool_size in pool_sizes])
            nr_pools = [_nr_pools(connectivity, pool_size, syn.nr_exc) for pool_size in pool_sizes]
            return p_ee, r_ee, pool_sizes, nr_pools

        with self.plot("syn_r_over_pool_size", scale=self.SCALE) as (fig, ax):
            palette = sns.color_palette()
            for color, conn in zip(palette, [.1, .2, .3]):
                p_ee, r_ee, pool_sizes, nr_pools = syn_pool_size_sweep(conn, 2)
                ax.scatter(pool_sizes, r_ee, marker="o",
                           s=5,
                           color=color)
                x, y = pool_sizes[-1], r_ee[-1]

                ax.annotate("${}={}$".format(style.latex_name("p_ee", dollar=False), conn), xy=(x, y), xytext=(400, y),
                            arrowprops=style.annotarrowprops)

            for val in style.plausible_regime["r"]:
                ax.axhline(val, dashes=style.linestyle.dashed, zorder=-20)
            ax.set_xlabel("{} {}".format(style.latex_name.long("pool_size"),
                                         style.latex_name.short("pool_size")), labelpad=style.default_label_pad_reduced)
            ax.set_ylabel(style.latex_name("reciprocity_ee"), labelpad=style.default_label_pad_reduced)
            style.set_ticksmiddle_minor(ax, [0, .1, .2, .3, .4, .5], "y")
            style.set_ticksmiddle_minor(ax, list(range(0, 401, 100)), "x")


class SubsamplingMatrix(FolderPlot):
    def requires(self):
        return OutputFolder()

    def run(self):
        sp.random.seed(42)
        n_neurons = 2000
        subsampling = .3  ## bad name, this is actually the stuff that is removed
        adj = sp.zeros((n_neurons, n_neurons))

        with self.plot("subsampling_matrix", despine=lambda x: None, scale=style.size.sq*.6) as (fig, ax):
            rem = sp.random.choice(n_neurons, int(n_neurons*(1-subsampling)))
            adj[rem, :] = sp.nan
            adj[:, rem] = sp.nan
            ax.matshow(adj, cmap="gray")
            ax.set_xlabel("Post")
            ax.set_ylabel("Pre")
            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_title("$ " + style.latex_name("subsampling", dollar=False) + "={}".format(subsampling) + " $", y=.95)
            ax.legend([mpl.patches.Patch(facecolor="black", edgecolor="black"),
                       mpl.patches.Patch(facecolor="white", edgecolor="black")],
                      ["Measured", "Not measured"], loc="upper center",
                      bbox_to_anchor=(.5, -.1))


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

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

    def run(self):
        connectivities = sp.linspace(.03, .3, num=20)
        exp = EXP(p_inh=.3, nr_neurons=2000, inh_ratio=.1)

        def run_exp(p):
            net = exp(p_exc=p)
            alpha = exp.alpha_exc_
            p_ee, r_ee = analyse(net)
            return p_ee, r_ee, alpha

        p_ee, r_ee, alpha = zip(*[run_exp(p) for p in connectivities])
        self.output().save({"connectivity": p_ee,
                            "reciprocity": r_ee,
                            "alpha": alpha})


class ExpConnectivityReciprocityScatter(FolderPlot):
    def requires(self):
        return ExpConnectivityReciprocityScatterData()

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

        with self.plot("exp_r_over_p", scale=ExpConnectivityOverDistance.SCALE) as (fig, ax):
            ax.scatter(data['connectivity'], data['reciprocity'],
                       color=sns.get_color_cycle()[0], edgecolor="w",
                       s=ERConnectivityReciprocityScatter.scatter_point_size)
            ax.set_xlabel(style.latex_name("p_exc"), labelpad=style.default_label_pad_reduced)
            ax.set_ylabel(style.latex_name("reciprocity_ee"), labelpad=style.default_label_pad_reduced)
            style.style_r_over_p(ax, "half")



from connectome.model.ll import reciprocity, connectivity


class LLMathPlots(FolderPlot):
    n_layers = 3
    title_offset = 1.07
    other_isoline_style = {"colors": "k", "linestyles": "dashed"}

    def make_labels(self, ax):
        ax.set_xlabel(style.latex_name("p_lateral"))
        ax.set_ylabel(style.latex_name("p_forward"))


    def run(self):
        x = sp.linspace(0.01, 1)
        y = x
        lateral, forward = sp.meshgrid(x, y)
        conn = connectivity(self.n_layers, lateral, forward)
        rec = reciprocity(self.n_layers, lateral, forward)

        for with_title in [True, False]:
            suffix = "" if with_title else f"_n_layers={self.n_layers}"
            with self.plot("connectivity" + suffix, scale=LL_SYN_HEATMAP_SCALE) as (fig, ax):
                mappable = ax.pcolormesh(lateral, forward, conn, vmin=0, vmax=1)
                ax.contour(lateral, forward, rec, levels=style.plausible_regime["r"], **self.other_isoline_style)
                countour_set = ax.contour(lateral, forward, conn, levels=style.plausible_regime["p"], colors="w")
                ax.clabel(countour_set, fmt="%.2f", manual=[(0,0), (1,1)])
                self.make_labels(ax)

                if with_title:
                    ax.set_title("${}={}$".format(style.latex_name("nr_exc_subpopulations", dollar=False),
                                                  self.n_layers),
                                 y=self.title_offset)

                cbar = fig.colorbar(mappable)
                cbar.outline.set_color("none")
                style.style_cbar_0_1(cbar, style.latex_name("expected_p_ee"))

                style.reduced_ticks_0_1(ax, despine=False)
                ax.yaxis.set_ticks_position("left")

        with self.plot("reciprocity", scale=LL_SYN_HEATMAP_SCALE) as (fig, ax):
            mappable = ax.pcolormesh(lateral, forward, rec, vmin=0, vmax=1)
            ax.contour(lateral, forward, conn, levels=style.plausible_regime["p"], **self.other_isoline_style)
            countour_set = ax.contour(lateral, forward, rec, levels=style.plausible_regime["r"], colors="w")
            ax.clabel(countour_set, fmt="%.2f", manual=[(.5, .5), (.2, .5)])
            self.make_labels(ax)


            cbar = fig.colorbar(mappable)
            cbar.outline.set_color("none")
            style.style_cbar_0_1(cbar, style.latex_name("expected_r_ee"))

            # reciprocity plot wihout title!!!!!!!

            style.reduced_ticks_0_1(ax, despine=False)
            ax.yaxis.set_ticks_position("left")




@memory.cache
def make_syn(nr_pools, pool_size, repetition):
    net = Network(nr_exc=1800, nr_inh=200)
    pool_connector = PoolConnector(p_exc=-2, pool_size=pool_size, network=net)
    pool_connector.use_fixed_nr_pools = nr_pools -1
    pool_connector()
    ana = standard_analysis(net)
    assert int(nr_pools) == len(pool_connector.pools_exc), "Pool number problem {} != {}".format(nr_pools-1, len(pool_connector.pools_exc))
    ana.update({"pool_size": pool_size,
                "nr_pools": nr_pools -1,
                "nr_exc": net.nr_exc,
                "nr_inh": net.nr_inh})
    return ana


@memory.cache
def make_syn_list(pars):
    return list(map(lambda x: make_syn(*x), pars))

class SYNRPHeatmap(FolderPlot):
    def run(self):
        sp.random.seed(42)
        nr_pools = sp.arange(10, 201, 10)
        pool_sizes = sp.arange(10, 201, 10)
        repetitions = sp.arange(3)
        from itertools import product
        pars = product(nr_pools, pool_sizes, repetitions)

        res = make_syn_list(pars)
        df = pd.DataFrame(res)
        df["nr_neurons"] = df.nr_exc + df.nr_inh
        assert len(df.nr_neurons.unique()) == 1, "More than one network size"
        nr_neurons = df.nr_neurons.iloc[0]


        for measure in ["p_ee", "reciprocity_ee"]:
            heat_df, x, y = self.get_data(df, measure)
            heat_df_other, x_other, y_other = self.get_data(df, "p_ee" if measure == "reciprocity_ee" else "reciprocity_ee")
            for with_title in [True, False]:
                with self.plot(measure + ("_title" if with_title else "_no_title"),
                          scale=LL_SYN_HEATMAP_SCALE) as (fig, ax):
                    mappable = ax.pcolor(*sp.meshgrid(x, y), heat_df, vmin=0, vmax=1)
                    levels = style.plausible_regime["p" if measure == "p_ee" else "r"]
                    levels_other = style.plausible_regime["r" if measure == "p_ee" else "p"]
                    ax.contour(*sp.meshgrid(x_other, y_other), heat_df_other, levels=levels_other, **LLMathPlots.other_isoline_style)
                    label_set = ax.contour(*sp.meshgrid(x, y), heat_df, colors="w", levels=levels)
                    ax.clabel(label_set, fmt="%.2f")
                    ax.yaxis.set_ticks_position("left")
                    cbar = fig.colorbar(mappable)
                    style.style_cbar_0_1(cbar, style.latex_name(measure))
                    cbar.outline.set_linewidth(0)
                    ax.set_xlim([10, 200])
                    style.set_ticksmiddle_minor(ax, list(range(10, 201, 10)), "x")
                    ax.set_xticks(ax.get_xlim())
                    ax.set_ylim([10, 200])
                    style.set_ticksmiddle_minor(ax, list(range(10, 201, 10)), "y")
                    ax.set_xlabel(style.latex_name(x.name), labelpad=style.default_label_pad_reduced)
                    ax.set_ylabel(style.latex_name(y.name), labelpad=style.default_label_pad_reduced-5)

                    if with_title:
                        ax.set_title("N={}".format(nr_neurons))

    def get_data(self, df, measure):
        heat_df = df.groupby(["pool_size", "nr_pools"])[measure].mean().unstack(0)
        x = heat_df.index
        y = heat_df.columns
        return heat_df, x, y


def diag_to_nan(arr):
    diag = np.diagonal(arr)
    diag.flags.writeable = True
    diag[:] = np.NaN


def interploate_plottable_line(correlation, probability, kind="e"):
    transform = lambda x: x if kind == "e" else lambda x: sp.fliplr(sp.flipud(x))
    correlation = transform(correlation.copy())
    probability = transform(probability.copy())
    diag_to_nan(correlation)
    diag_to_nan(probability)

    correlation = correlation.flatten()
    probability = probability.flatten()

    valid = ~np.isnan(correlation)
    correlation = correlation[valid]
    probability = probability[valid]

    args = sp.argsort(correlation)
    correlation = correlation[args]
    probability = probability[args]

    f = sp.interpolate.interp1d(correlation, probability)
    correlation_plot_points = sp.linspace(correlation.min(), correlation.max(), 1000)
    probability_plot_values = f(correlation_plot_points)
    return correlation_plot_points, probability_plot_values


@memory.cache
def p_over_corr(model_name, n_pow, fdim, population, n):
    from connectome.model import API, SimpleMiller
    Model = API if model_name == "API" else SimpleMiller
    api = Model(nr_neurons=n, inh_ratio=.1, n_pow=n_pow, feature_space_dimension=fdim, p_exc=.2, p_inh=.5)
    api()
    correlation = api.correlations_
    probability = api.probabilities_
    if population == "e":
        correlation = correlation[:, :api.nr_exc]
        probability = probability[:, :api.nr_exc]
    else:
        correlation = correlation[:, api.nr_exc:]
        probability = probability[:, api.nr_exc:]

    correlation_plot_points, probability_plot_values = interploate_plottable_line(correlation, probability)
    return correlation_plot_points, probability_plot_values


class APIProbabilityOverCorrelation(FolderPlot):
    fdim = 20
    def run(self):
        def break_ln(s):
            if s[0] == "$":
                return s
            return s + "\n"

        def maybe_colon(s):
            if s != "exc." and s != "inh.":
                return ", " + s
            return ""

        for with_legend in [True, False]:
            with self.plot("correlation_to_probability_" + ("legend" if with_legend else "no_legend"),
                           scale=APIOrginalVsAdaptedPlot.SCALE) as (fig, ax):
                n = 2000
                n_pows = [5]
                fdim_list = [self.fdim]
                SM_SATURATION = 1
                models, linestyles = ["SM", "API"], [figures.style.linestyle.dashed, figures.style.linestyle.solid]
                for population, color in zip(["e", "i"],  sns.get_color_cycle()):
                    for n_pow in n_pows:
                        for fdim in fdim_list:
                            for model, linestyle in zip(models, linestyles):
                                c_plot, val_plot = p_over_corr(model, n_pow, fdim, population, n)
                                color_mod = style.mimic_alpha(color, SM_SATURATION, (0, 0, 0)) if model == "SM" else color
                                ax.plot(c_plot, val_plot,  color=color_mod, dashes=linestyle,
                                        clip_on=False,
                                        label=(style.latex_name.long("adapted") if model == "API" else style.latex_name.long("original"))
                                              + maybe_colon((style.latex_name.long("excitatory") if population == "e" else style.latex_name.long("inhibitory"))))
                                ax.set_xticks([-.8, -.6, -.4, -.2, .2, .4, .6, .8], minor=True)
                                ax.set_yticks(sp.arange(0, 1.1, .2))
                                ax.set_yticklabels([0, "", "", "", "", 1])
                                sns.despine(ax=ax)
                                ax.set_ylabel(break_ln(style.latex_name("p_exc") + ", ") + style.latex_name("p_inh"),
                                              labelpad=style.default_label_pad_reduced)
                                ax.set_xlabel(style.latex_name("cosine_similarity") + "$(FV_k, FV_l)$")
                                ax.set_xticks([-1, 0, 1])
                                ax.spines["bottom"].set_position(("data", -.05))
                                style.reduced_ticks_0_1(ax, "y")

                                if with_legend:
                                    ax.legend(loc="center left", bbox_to_anchor=(.9, .5))
                                    ax.set_title(break_ln(f'{style.latex_name("n_pow")}$={n_pow}$, ') +
                                                 f'{style.latex_name("feature_space_dimension")}$={fdim}$')


class FeatureVectorCorrelationHistogram(FolderPlot):
    def run(self):
        from connectome.model.featurevector import uniform_spherical_feature_fectors
        from connectome.model.correlationbasednetwork import feature_vector_correlations
        import matplotlib as mpl

        def corr_list(nr_neurons, dimension):
            fv = uniform_spherical_feature_fectors(nr_neurons, dimension)
            corr = feature_vector_correlations(fv)
            diag_to_nan(corr)
            corr = corr.flatten()
            corr = corr[~np.isnan(corr)]
            return corr

        width, height = APIOrginalVsAdaptedPlot.SCALE
        height /= 2
        for invert, suffix, dim_list in [(True, "", [APIProbabilityOverCorrelation.fdim]),
                                         (False, "_upright",
                                          [2, 5, APIProbabilityOverCorrelation.fdim, 40])]:
            with self.plot("hist" + suffix, scale=(width, height), despine=lambda x: None) as (fig, ax):
                legend_items = []
                for dim, color in zip(dim_list, sns.get_color_cycle()):
                    corr = corr_list(2000, dim)
                    ax.hist(corr, bins="auto", histtype="step", normed=True, color=color, linewidth=style.linewidth)
                    legend_items.append(mpl.lines.Line2D((0,0), (0,0), color=color))
                ax.legend(legend_items,
                          [f"${dim}$" for dim in dim_list],
                          loc="center left", bbox_to_anchor=(1, .4),
                          labelspacing=.1,
                          title=style.latex_name("feature_space_dimension"),
                          handlelength=.2,
                          handletextpad=.2)
                #sns.despine(ax=ax, offset=.5)
                if invert:
                    ax.spines["top"].set_visible(True)
                    ax.spines["bottom"].set_visible(False)
                    ax.spines["right"].set_visible(False)
                else:
                    ax.spines["top"].set_visible(False)
                    ax.spines["right"].set_visible(False)
                ax.set_xticks([-1, 0, 1])
                ax.set_xticks([-.8, -.6, -.4, -.2, .2, .4, .6, .8], minor=True)
                style.ticks_to_major_minor(ax, "y")
                ax.set_xlabel(style.latex_name("cosine_similarity") + "$(FV_k, FV_l)$")
                ax.set_ylabel("PDF", labelpad=0)
                if invert:
                    ax.invert_yaxis()
                    ax.xaxis.tick_top()
                    ax.xaxis.set_label_position('top')


def sorn_data_over_time():
    from connectome.model import SORN, Network
    from connectome.analysis import ConnectivityEstimator, ReciprocityEstimator, RelativeReciprocityEstimator
    import pandas as pd

    sorn = SORN(nr_neurons=2000, inh_ratio=.1, p_exc=.2, p_inh=.5,
                eta_stdp=.0012, eta_intrinsic=.75)

    class RPRecorder:
        def __init__(self):
            self.data = []
            self.activity = []
            self.nr_synapses_added = []

        def record(self, t, name, value):
            if name == "adjacency_matrix":
                if t % 10 == 0:
                    print("t={}: {:.2f}".format(t, t/sorn.NR_TIME_STEPS))
                net = Network(value)
                analysis = {}
                for Analysis in [ConnectivityEstimator, ReciprocityEstimator, RelativeReciprocityEstimator]:
                    analysis.update(Analysis(network=net)())
                analysis["t"] = t
                self.data.append(analysis)
            if name == "current_active_neurons":
                self.activity.append(value.copy())

            if name == "nr_synapses_added":
                self.nr_synapses_added.append(value)

    sorn.recorder = RPRecorder()
    sorn()
    df = pd.DataFrame(sorn.recorder.data)
    return df, sp.asarray(sorn.recorder.activity), sorn.nr_exc, sorn.recorder.nr_synapses_added


def visualize_activity(fig, ax, activity, nr_exc, with_legend, color, colored_legend):
    corr = pair_correlation(activity)
    show_n_neurons = 25
    show_n_time_steps = 100
    n_neurons = activity.shape[1]
    activity[:, nr_exc:] *= -1
    activity = activity[-show_n_time_steps:, ::n_neurons // show_n_neurons]

    #cmap = mpl.cm.get_cmap('Greys', 3)
    cmap = mpl.colors.LinearSegmentedColormap.from_list("mycm", ["w", color, "k"])
    ax.matshow(activity.T, aspect="auto", cmap=cmap)

    ax.set_xlabel(style.latex_name("time"), labelpad=style.default_label_pad_reduced)


    ax.locator_params(nbins=3, axis="y")
    ax.set_xticks([0, 100])
    ax.set_xlim(0, 100)
    #ax.set_title("${}={:.2f}$".format(SORNSaneActivity.POPULATION_CORRELATION_LABEL[1:-1], corr))

    if not with_legend:
        ax.set_yticklabels([""] * len(ax.get_yticks()))

    if with_legend:
        ax.set_ylabel("Neuron")
        edgecolor = color if colored_legend else "gray"
        ax.legend([mpl.patches.Patch(facecolor="black", edgecolor=edgecolor, linewidth=1.5),
                   mpl.patches.Patch(facecolor="white", edgecolor=edgecolor, linewidth=1.5)],
                  ["E", "I"],
                  loc="center left",
                  bbox_to_anchor=(1, .5),
                  handlelength=.5)


class SORNDataOverTime(FolderPlot):
    def output(self):
        return luigi.LocalTarget(os.path.join(self.output_folder(), "data.pickle"))

    def run(self):
        import pickle
        res = sorn_data_over_time()
        self.make_folder()
        with open(self.output().path, "wb") as f:
            pickle.dump(res, f)


class SORNOverTime(FolderPlot):
    def requires(self):
        return SORNDataOverTime()

    def run(self):
        import pickle
        with open(self.input().path, "rb") as f:
            inp = pickle.load(f)

        df, activity, nr_exc, nr_synapses_added = inp
        self.plot_r_p_over_time(df)
        #self.plot_activity(activity, nr_exc)
        self.plot_p_add(nr_synapses_added, nr_exc)

    def plot_r_p_over_time(self, df):

        for length, linestyle, scale in [(50, ".", 4), (10000, ".", 1)]:
            with self.plot("rp-length={}".format(length)) as (fig, ax):
                fig.set_size_inches(1, .5)
                df_plot = df[:length]
                ax.plot("t", "p_ee", data=df_plot, label=style.latex_name("p_ee"),
                           color=sns.get_color_cycle()[0])
                ax.plot("t", "reciprocity_ee", data=df_plot, label=style.latex_name("reciprocity_ee"),
                           color = sns.get_color_cycle()[1])
                ax.set_xlabel(style.latex_name("time"))
                ax.set_ylabel("Connectivity,\nreciprocity")
                sns.despine(ax=ax)
                ax.legend(loc="center left", bbox_to_anchor=(1, .5))
                ax.set_xticks([0, length])
                ax.set_xticklabels(["0", r"$t_{\mathrm{end}}$=" + str(length)])
                style.set_ticksmiddle_minor(ax, range(0, length + 1, length // 10), "x")
                style.set_ticksmiddle_minor(ax, [.15, .2, .25], "y")

    def plot_p_add(self, nr_synapses_added, nr_exc):
        for start, end in [(0, 10000), (9900, 10000)]:
            for with_legend in [True, False]:
                with self.plot("p_add_{}_legend_{}".format((start, end), with_legend),
                               scale=size.s) as (fig, ax):
                    ax.plot(sp.arange(start, end),
                            sp.array(nr_synapses_added[start:end]) / nr_exc**2, label=style.latex_name.long("adapted"))
                    ax.plot([start, end], [.1/200**2]*2, label=style.latex_name.long("original"))

                    ax.set_xticks(sp.arange(start, end, (start-end) // 10), minor=True)
                    ax.set_xticks([start, end])
                    #ax.xaxis.set_major_formatter(style.formatter.ExpFormatter())
                    ax.set_xticklabels(["$\mathdefault{{t_{{\\mathrm{{end}}}}}}-{}$".format(end-start),
                                        "$\mathdefault{t_{\\mathrm{end}}}$"])
                    ax.set_xlim(start, end)
                    ax.set_xlabel(style.latex_name("time") + "$(\Delta t)$", labelpad=style.default_label_pad_reduced)


                    #ax.set_yticks([1e-7, 1e-2])

                    ax.set_yscale("log")
                    ax.set_yticks([1e-6, 1e-5, 1e-4, 1e-3, 1e-2])
                    ax.set_yticklabels(["$10^{-6}$", "", "", "", "$10^{-2}$"])
                    ax.set_yticks([], minor=True)
                    ax.set_ylabel(style.latex_name("synapse_addition_probability"), labelpad=style.default_label_pad_reduced)

                    labels = ax.get_yticklabels()
                    for lab in labels[2:-3]:   # really weird but it works
                        lab.set_visible(False)
                    if with_legend:
                        ax.legend(loc="center left", bbox_to_anchor=(1, .5))


    def plot_activity(self, activity, nr_exc):
        with self.plot("activity") as (fig, ax):
            visualize_activity(fig, ax, activity, nr_exc, True, "w")

@memory.cache
def make_sorn_zero_after_time(nr_time_steps):
    from connectome.model import SORN
    sorn = SORN(nr_neurons=2000, inh_ratio=.1, p_exc=.2, p_inh=.5,
                eta_stdp=.0012, eta_intrinsic=.75)
    sorn.NR_TIME_STEPS = nr_time_steps
    net = sorn()
    return net


class SORNZeroAndInfinitTimeWeightDistribution(FolderPlot):
    def run(self):
        net_0 = make_sorn_zero_after_time(0)["E","E"]
        net_inf = make_sorn_zero_after_time(10000)["E","E"]

        with self.plot("weight_at_0_and_10000", scale=size.s) as (fig, ax):
            #ax.axvline(1/1800, linestyle=":")
            ax.annotate(xy=(1/1800, 26500), xytext=(.004, 30000*1.1),
                        s=r"$w_{\mathrm{add}}$", arrowprops=dict(facecolor="black", arrowstyle="-|>"))


            # having color number 2 and 3 of the color cycle is essential b/c it appears
            # on the side of a plot with colors 0 and 1 with different meaning
            ax.hist(net_0[net_0 != 0].flatten(), bins="auto", label="0", histtype="step",
                    linewidth=style.linewidth,
                    color=sns.get_color_cycle()[2])
            ax.hist(net_inf[net_inf != 0].flatten(), bins="auto", label="10000", histtype="step",
                    linewidth=style.linewidth, color=sns.get_color_cycle()[3])



            style.set_ticksmiddle_minor(ax, np.linspace(0, .01, 11), "x")
            style.set_ticksmiddle_minor(ax, np.linspace(0, 30000, 4), "y")


            ax.legend([mpl.lines.Line2D((0,0), (0,0), linewidth=style.linewidth, color=color)
                       for color in sns.get_color_cycle()[2:]],
                      [style.Tex("t_0"), style.Tex(r"t_{\mathrm{end}}")],
            loc="center left", bbox_to_anchor=(.4, .5))
            ax.set_xlabel(style.latex_name("weight"), labelpad=style.default_label_pad_reduced)
            ax.set_ylabel(style.latex_name("synapse_density"), labelpad=style.default_label_pad_reduced-5)


            from figures.style.formatter import ExpFormatter
            ax.yaxis.set_major_formatter(ExpFormatter())






class SynPiOverInhPoolSize(FolderPlot):
    def run(self):
        with self.plot("pi_over_inh_pool_size", scale=SynConnectivityReciprocityNrPools.SCALE*.8) as (fig, ax):
            from connectome.model import SYN
            from connectome.analysis import standard_analysis
            syn = SYN(nr_neurons=2000, inh_ratio=.1, p_exc=.2, p_inh=.5)
            pool_sizes = sp.arange(80, 300, 10)

            @memory.cache
            def pi_inh_pool_size(pool_size):
                net = syn(pool_size=pool_size)
                ana = standard_analysis(net)
                inh_size = syn.pool_connector_.inh_pool_size()
                ana.update({"inh_pool_size": inh_size})
                return ana

            df = pd.DataFrame(list(map(pi_inh_pool_size, pool_sizes)))
            measure_list = ["p_ie", "p_ei"]
            for measure, color in zip(measure_list, sns.get_color_cycle()):
                ax.scatter("inh_pool_size", measure, data=df, c=color, label=style.latex_name(measure))
            ax.legend(loc="center left", bbox_to_anchor=(1, .5))
            ax.set_xlabel(style.latex_name("pool_size_inh"))
            ax.set_ylabel("$ " + ", ".join(map(lambda x: style.latex_name(x, dollar=False), measure_list)) + " $")
            ax.set_xticks([0, 20, 40])
            ax.set_xticks([10, 30], minor=True)
            ax.set_ylim(0, .6)
            ax.set_yticks(ax.get_ylim())
            ax.set_yticks([.1, .2, .3, .4, .5], minor=True)
            sns.despine(ax=ax)



@memory.cache
def sorn_sub_time_step():
    np.random.seed(42)
    from connectome.model import SORN, Network
    from connectome.analysis import standard_analysis
    import pandas as pd

    sorn = SORN(nr_neurons=2000, inh_ratio=.1, p_exc=.2, p_inh=.5,
                eta_stdp=.0012, eta_intrinsic=.75)
    sorn.NR_TIME_STEPS = 40

    class Recorder:
        def __init__(self):
            self.data = []

        def record(self, t, name, value):
            if name not in ["adjacency_matrix", "adjacency_matrix_before_add_synapses"]:
                return

            analysis = standard_analysis(Network(value))

            analysis["t"] = float(t) if name == "adjacency_matrix" else float(t) - .5
            self.data.append(analysis)

    sorn.recorder = Recorder()
    sorn()
    df = pd.DataFrame(sorn.recorder.data)
    return df


class SORNPOverTimeSubTimeStep(FolderPlot):
    def run(self):
        df = sorn_sub_time_step()[25:34]
        with self.plot("sorn_sub_time_step", scale=.6) as (fig, ax):
            ax.plot("t", "p_ee", ".-", data=df, label="Complete\nstep", clip_on=False)
            ax.plot("t", "p_ee", "^", data=df[1::2], label="Before\nstructural\nplasticity")
            ax.set_xlabel(style.latex_name("time"))
            ax.set_ylabel(style.latex_name("p_ee"))
            ax.set_xticks(df.t[::2])
            ax.get_yaxis().get_major_formatter().set_useOffset(False)
            plt.setp(ax.xaxis.get_majorticklabels(), rotation=style.vertical_label_angle)
            ax.locator_params(nbins=4, axis="y")
            ax.legend(loc="center left", bbox_to_anchor=(1, .5))
            sns.despine(ax=ax)



class APINBinomWeightDistribution(FolderPlot):
    def run(self):
        from itertools import product
        n_pow = [4, 6]
        f_dim = [5, 50]
        args = product(n_pow, f_dim)

        @memory.cache
        def get_weights(n_pow, f_dim):
            from connectome.model import API
            api = API(nr_neurons=2000, inh_ratio=.1, p_exc=.2, p_inh=.5)
            api(n_pow=n_pow, feature_space_dimension=f_dim)
            res = {"n_binom_exc": api.n_bin_exc_,
                   "n_binom_inh": api.n_bin_inh_,
                   "nr_exc": api.nr_exc,
                   "nr_inh": api.nr_inh,
                   "p": api.probabilities_.copy()}
            res.update(api.inputs)
            return res

        res = list(map(lambda x: get_weights(*x), args))

        def br(x):
            if x[0] == "$":
                return x
            return x + "\n"

        with self.plot("api_weight_distribution_as_function_of_n_bin") as (fig, ax):
            for val in res:
                p = val["p"]
                diag_to_nan(p)
                p = p[:val["nr_exc"], :val["nr_exc"]]
                p = p[~np.isnan(p)]
                hist, bin_edges = np.histogram(p.flatten(), bins="auto", normed=True)
                ax.plot(bin_edges[:-1], hist,
                        label=f'${val["n_pow"]}, {val["feature_space_dimension"]}$',
                        linewidth=style.linewidth)
                ax.set_yscale("log")
                ax.legend(loc="center left", bbox_to_anchor=(1, .5),
                          title=br(style.latex_name("n_pow") + ",")
                                + style.latex_name("feature_space_dimension"))
                ax.set_xlabel(style.latex_name("weight"), labelpad=style.default_label_pad_reduced)
                ax.set_ylabel(style.name["PDF"])
                #style.remove_middle_tick_labels(ax, "y")
                #style.ticks_to_major_minor(ax, "x")
                sns.despine(ax=ax)


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

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

    def run(self):
        from parallel import SGE, sge_available
        from itertools import product
        import scipy as sp
        mapper = SGE(num_threads=4, priority=0).map if sge_available() else map

        def make(lmbda, feature_space_dimension):
            from connectome.model.lasso_optimizer import DevelopmentalFeverLassoOptimizer
            from connectome.model.featurevector import uniform_spherical_feature_fectors
            from connectome.model import Network
            from connectome.analysis import standard_analysis
            nr_neurons = 2000
            nr_exc = int(nr_neurons * .9)
            optimizer = DevelopmentalFeverLassoOptimizer(nr_exc=nr_exc)
            feature_vectors = uniform_spherical_feature_fectors(nr_neurons, feature_space_dimension)
            lambda_exc, lambda_inh = lmbda, lmbda
            adjacency = optimizer.fit(feature_vectors, sp.zeros((nr_neurons, nr_neurons)), lambda_exc, lambda_inh)
            net = Network(adjacency, nr_exc=nr_exc, nr_inh=nr_neurons - nr_exc)
            analysis = standard_analysis(net)
            analysis.update({"lmbda": lmbda,
                             "feature_space_dimension": feature_space_dimension,
                             "nr_neurons": nr_neurons, "nr_exc": nr_exc})
            return analysis

        lmbdas = [1e-9, 1e-8, 1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2]
        feature_space_dimensions = [30, 80, 130, 180]
        pars = product(lmbdas, feature_space_dimensions)
        res = list(mapper(lambda x: make(*x), pars))
        self.output().save(res)


class FEVERLambdaFeatureSpaceDimensionConnRec(FolderPlot):
    def requires(self):
        return DevelopmentalFeverData()

    def run(self):
        import pandas as pd
        df_raw = pd.DataFrame(self.input().load())
        for measure in ["p_ee", "reciprocity_ee"]:
            df = df_raw.pivot("lmbda", "feature_space_dimension", measure)
            with self.plot(measure, scale=1) as (fig, ax):
                fig.set_size_inches(2, .5)
                sns.heatmap(df, ax=ax, annot=False, annot_kws={"fontsize": 8}, fmt=".2f",
                            cbar_kws={"label": style.latex_name(measure)}, vmin=0, vmax=1)
                cbar = fig.axes[1]
                cbar.set_yticklabels([0, "", "", "", "", 1])
                cbar.set_ylabel(style.latex_name(measure), labelpad=style.default_label_pad_reduced)

                plt.setp(ax.yaxis.get_majorticklabels(), rotation=0)
                ax.set_yticklabels([df.index[-1]] + [""] *(len(df.index)-2) +  [df.index[0]])

                ax.set_xlabel(style.latex_name("feature_space_dimension"))
                ax.set_ylabel("$\lambda$", labelpad=style.default_label_pad_reduced)






class TextureTraces(FolderPlot):
    def run(self):
        self.texture_traces()

    def texture_traces(self):
        from connectome.function.task.texture_data import TextureData
        textures = TextureData()
        textures, classes = textures.get_data(1, 7, 500, 50)
        offset_scale = 1.5

        with self.plot("texture_traces", scale=1.3, despine=no_despine) as (fig, ax):
            print("run")

            from itertools import cycle
            colors = cycle(sns.get_color_cycle())
            for data, cls, offset in zip(textures, classes, sp.arange(len(classes)) * offset_scale):
                ax.plot(data -offset, label="$T_{{{}}}$".format(cls), color="k")
                ax.text(-25, -offset, "$T_{{{}}}$".format(cls+1), ha="right", va="center")
            ax.spines["left"].set_visible(False)
            ax.spines["right"].set_visible(False)
            ax.spines["top"].set_visible(False)
            ax.set_yticks([])
            ax.tick_params("y", left="off", right="off")
            ax.set_xticks([0, 250, 500],)
            ax.set_xticks([125, 375], minor=True)
            ax.tick_params("x", top="off", which="both")
            ax.set_xlabel(TaskTexture.TIME_LABEL)
            ax.spines["bottom"].set_position(("data", -offset-1))
            style.annotate_group(ax, 250, 500, -offset-2.9, "Training", length=-.4, text_offset=1.7)


def model_name(comment):
    return comment.split("(")[0]

def texture_palette():
    return sns.get_color_cycle()


TEXTURE_TIME_ABBR = "$t \; (\Delta t)$"


class TaskTexture(FolderPlot):
    O = "o"
    TIME_LABEL = style.latex_name.long("texture_presentation_time") + " " + TEXTURE_TIME_ABBR
    INP_POOL_READOUT_WIDTH = 1
    READOUT_LABEL = style.latex_name("output_activity")
    #MODEL_IDs = [842, 843, 844]
    MODEL_IDs = [896, 897, 898]

    def requires(self):
        return [OutputFolder(), GrayScaleTextures()]

    def example_trace(self):
        file = os.path.join(self.input()[1].fn, "1-example_trace.png")
        import skimage.io
        img = skimage.io.imread(file)
        return img

    def run(self):
        data_dir = os.path.join(self.data_folder(), "model_checkpoints")
        db = os.path.join(data_dir, "keras_logs.db")

        data = partial(model_eval_data, db, data_dir)
        accuracies = []
        traces_models = []
        y_gt_models = []


        for model_id in self.MODEL_IDs:
            traces, y_gt, conf_mat, nr_layers, comment, y_pred = data(model_id)
            traces_models.append((model_name(comment), traces))
            y_gt_models.append((model_name(comment), y_gt))
            conf_mat = conf_mat / conf_mat.sum(axis=1, keepdims=True)
            accuracies.append({"Model": model_name(comment), "Accuracy": conf_mat.diagonal().mean()})
            self.plot_input_pool_readout(db, data_dir, model_id)
            for plot_trace_nr in [0, 1, 2, 3, 4]:
                self.plot_readout_trace(comment, model_id, plot_trace_nr, traces, y_gt)
            self.plot_conf_mat(conf_mat, model_id, nr_layers, model_name(comment))
        self.accuracy_bar_plot(accuracies, traces_models, y_gt_models)

    def accuracy_bar_plot(self, accuracies, traces_models, y_gt_models):
        df = pd.DataFrame(accuracies)


        # traces is text_examples x texture_length x readout_neuron
        y_preds = [(name, traces[:,-250:,:].sum(1).argmax(1)) for name, traces in traces_models]
        y_gts = [(name, y_gts.argmax(1)) for name, y_gts in y_gt_models]


        df_err = pd.DataFrame([{"Model": name_pred, "Correct": p == g}
                              for (name_pred, pred), (name_gt, gt) in zip(y_preds, y_gts)
                              for p, g in zip(pred, gt)])

        with open(os.path.join(self.output_folder(), "n_texture_exposures.txt"), "w") as f:
            f.write(f"N texture exposures={len(df_err)}")


        df_boot = pd.DataFrame()

        n_boot = 100
        for _ in range(n_boot):
            this_boot = df_err.groupby("Model").apply(
                lambda x: pd.DataFrame.sample(x, frac=1,
                                              replace=True).mean()).reset_index()
            df_boot = df_boot.append(this_boot)


        df["Model"] = df["Model"].map(style.name)
        with self.plot("accuracy_box", scale=(.55, .8)) as (fig, ax):
            sns.boxplot("Model", "Correct", data=df_boot, fliersize=.7,
                        linewidth=.2,
                        color="C2",
                        ax=ax)
            ax.set_ylabel(style.latex_name.long("texture_classification_accuracy"))
            ax.set_xlabel("")
            style.reduced_ticks_0_1(ax, axis="y", label_pad_reduce=0)
            plt.setp(ax.xaxis.get_majorticklabels(), **PlotPosteriorDistributionMatrix.LABEL_ROTATION_PARS)
            channce_level = 1/7
            chance_color = sns.utils.get_color_cycle()[1]
            ax.axhline(channce_level, dashes=style.linestyle.dashed, color=chance_color)
            ax.text(x=2.8, y=channce_level, s=style.latex_name.long("chance"), color=chance_color, ha="left")

        with self.plot("accuracy_bar_with_errorbar", scale=(.55, .3)) as (fig, ax):
            sns.barplot("Model", "Correct", data=df_err,
                        errcolor=sns.utils.get_color_cycle()[2],
                        ax=ax, edgecolor="none", color="k")
            ax.set_ylabel(style.latex_name.long("texture_classification_accuracy"))
            ax.set_xlabel("")
            style.reduced_ticks_0_1(ax, axis="y", label_pad_reduce=0)
            plt.setp(ax.xaxis.get_majorticklabels(), **PlotPosteriorDistributionMatrix.LABEL_ROTATION_PARS)
            channce_level = 1/7
            chance_color = sns.utils.get_color_cycle()[1]
            ax.axhline(channce_level, dashes=style.linestyle.dashed, color=chance_color)
            ax.text(x=2.8, y=channce_level, s=style.latex_name.long("chance"), color=chance_color, ha="left")

        with open(os.path.join(self.output_folder(), "bar.txt"), "w") as f:
            f.write(str(df))

    def plot_readout_trace(self, comment, model_id, plot_trace_nr, readout_traces, y_gt):
        name = comment.split("(")[0]
        traces_to_plot = readout_traces[plot_trace_nr]
        gt_class = y_gt[plot_trace_nr].argmax()
        for with_title in [True, False]:
            for others_dotted in [True, False]:
                title_str = ("" if with_title else "-no-title") + ("-dotted" if others_dotted else "")
                with self.plot(f"readout_traces_{model_id}-{comment}-trace_nr={plot_trace_nr}" + title_str,
                               scale=.35) as (fig, ax):
                    colors = cycle(texture_palette())
                    for k, trace in enumerate(traces_to_plot.T):
                        ax.plot(trace, label=f"${self.O}_{k+1}$", color=next(colors),
                                linestyle="dotted" if k != gt_class else "solid",
                                zorder=100 if k == gt_class else 0)
                    ax.set_xticks([0, 500])
                    ax.set_xticks([250], minor=True)
                    ax.set_xlim(ax.get_xticks())
                    ax.set_xlabel(style.latex_name("texture_presentation_time"), labelpad=0)
                    ax.set_ylabel(self.READOUT_LABEL, labelpad=5)

                    if with_title:
                        ax.set_title(style.name[name])
                    gt_label = "$\\mathdefault{{ T_{{ {} }} }}$".format(gt_class+1)
                    ax.text(700, .65, gt_label, va="center", ha="right", color=sns.utils.get_color_cycle()[gt_class])
                    ax.set_xlim(0, traces_to_plot.shape[0])
                    style.reduced_ticks_0_1(ax, "y", label_pad_reduce=1)
                    ax.set_ylim(0, 1.05)

    def plot_conf_mat(self, conf_mat, model_id, nr_layers, model_name):
        for with_cbar in [True, False]:
            suffix = "_cbar" if with_cbar else ""
            with self.plot(("conf_mat_{}_layers_{}".format(model_id, nr_layers)) + suffix, scale=.5) as (fig, ax):
                mappable = ax.matshow(conf_mat, vmin=0, vmax=1)
                if with_cbar:
                    cbar = fig.colorbar(mappable)
                    cbar.outline.remove()
                    cbar.set_label(style.latex_name("fractional_occurence"), labelpad=-5)
                    cbar.set_ticks([0, 1])
                ax.set_xlabel(style.latex_name.long("predicted"), labelpad=0)
                ax.set_ylabel("True", labelpad=0)
                labels = ["", "$T_1$"] + [""] * (conf_mat.shape[0]-2) + ["$T_{{{}}}$".format(conf_mat.shape[0])]
                ax.set_xticklabels(labels)
                ax.set_yticklabels(labels)
                ax.set_title(style.name[model_name])


    def plot_input_pool_readout(self, db, checkpoints, id):
        example_trace = self.example_trace()[0].astype(float)
        example_trace -= example_trace.min()
        example_trace /= example_trace.max()

        @memory.cache
        def make_textures_examples_inp_pool_readout(db, checkpoints, id):
            from connectome.function.task.texture import TextureModelAnalyser
            e = TextureModelAnalyser(db, checkpoints, id)

            tr = e.get_pool_traces_single_input(example_trace)
            #X, _ = e.get_test_data(1)
            y_pred = e.predict_single_trace(example_trace)
            return tr, y_pred, e.net_

        tr, y_pred, net = make_textures_examples_inp_pool_readout(db, checkpoints, id)

        with self.plot("complete_network-{}".format(id), 3, sharex=True, despine=lambda x: None) as (fig, (ax1, ax2, ax3)):
            fig.set_size_inches((self.INP_POOL_READOUT_WIDTH, 1.8))
            fig.tight_layout(h_pad=.1)

            sns.despine(fig)
            ax1.set_xlim(0, 500)


            # ax2.plot(tr[0], color="k");
            colors = cycle(texture_palette())
            for this_pred in y_pred[0].T:
                ax1.plot(this_pred, color=next(colors))
            ax1.set_yticks([0, 1])
            ax1.set_ylim(0, 1.1)
            ax1.set_ylabel(self.READOUT_LABEL, labelpad=style.default_label_pad_reduced)

            orig_lines = [l for l in ax1.lines]

            for n_line, line in enumerate(orig_lines):
                 color = line.get_color()
                 y_pos = .75 - n_line / len(orig_lines) * .5
                 ax1.add_line(Line2D([1.05, 1.2], [y_pos, y_pos],
                                     transform=ax1.transAxes,
                                     color=color,
                                     clip_on=False))

            xTEXT = 1.25
            ax1.text(xTEXT, .75, f"${self.O}_1$", ha="left", va="center",
                     transform=ax1.transAxes,)
            #ax1.text(xTEXT, .5, r"$\ldots$", ha="left", va="center",
            #transform = ax1.transAxes)
            ax1.text(xTEXT, .25, f"${self.O}_7$", ha="left", va="center",
            transform = ax1.transAxes)

            plot_subset = tr[0].T
            plot_subset[net.nr_exc:,:] *= -1
            plot_subset = tr[0].T[::50,:]
            absmax = sp.absolute(plot_subset).max()
            mappable = ax2.matshow(plot_subset, aspect="auto", vmin=-absmax, vmax=absmax, cmap="RdBu")

            ax2.set_ylabel(style.latex_name("neuron"), labelpad=style.default_label_pad_reduced-5.7)
            ax2.xaxis.set_ticks_position("bottom")
            ax2.set_yticks([0, plot_subset.shape[0]])

            cax = fig.add_axes([.85,  .5, .1, .1])
            cbar = fig.colorbar(mappable, cax=cax)
            cbar.outline.set_linewidth(0)
            cbar.set_ticks([-absmax, 0, absmax])
            cbar.set_ticklabels([-1, 0, 1])
            cax.set_title("e", y=.8)
            cbar.set_label(style.latex_name("pool_activity"), labelpad=1)
            cax.set_xlabel("i")


            ax3.plot(example_trace, color=config.maybe_to_black(texture_palette()[1]))
            ax3.set_xticks([0, 250, 500])
            ax3.set_xlabel(TaskTexture.TIME_LABEL)
            ax3.set_ylabel(style.latex_name("input_activity"), labelpad=style.default_label_pad_reduced)
            ax3.set_yticks([0, 1])

            style.reduced_ticks_0_1(ax1, "y")
            style.reduced_ticks_0_1(ax3, "y")






DsynResult = namedtuple("DsynResult", "membrane_v membrane_t spike_i spike_t pools")


@memory.cache
def make_dsyn():
    from connectome.model import SYN
    from connectome.function.task.propagation.task import PropagationTask
    sp.random.seed(22)
    pt = PropagationTask()
    syn = SYN(nr_neurons=1780, inh_ratio=.1, p_exc=.15, p_inh=.5,
              pool_size=290)
    dsyn = pt.dynamic_syn_factory(syn)
    dsyn.run(pt.run_time)
    return DsynResult(sp.asarray(dsyn.state_monitor_.v), sp.asarray(dsyn.state_monitor_.t),
                      sp.asarray(dsyn.spike_monitor_.i), sp.asarray(dsyn.spike_monitor_.t),
                      dsyn.pools_)


def syn_pool_label(this_subset, all_subsets):
    this_str = " \cap ".join("P_{{{}}}".format(subset) for subset in this_subset)
    this_str = "(" + this_str + ")" if len(this_subset) > 1 and len(this_subset) != len(all_subsets) else this_str

    if len(this_subset) == len(all_subsets):
        other_str = ""
    else:
        other_str = " \cup ".join("P_{{{}}}".format(subset) for subset in all_subsets if subset not in this_subset)
        if len(all_subsets) - len(this_subset) > 1:
            other_str = "(" + other_str + ")"
        other_str = setminus_unicode + other_str

    return "$" +  this_str  + other_str + "$"


class SynActivityRaster(FolderPlot):
    FIG_SIZE = (1, 2)
    #force_replot = True
    MAX_SUM_TRACE_PLOT_T = 10

    def run(self):
        dsyn = make_dsyn()
        xlim = self.plot_spikes_all_groups()
        self.plot_traces(dsyn, xlim)


    def plot_traces(self, dsyn, xlim):
        si = dsyn.spike_i
        st = (dsyn.spike_t * 1000000).astype(int)
        pools = dsyn.pools[1:]
        n_spikes_at_t_list = []
        with self.plot("sum_traces", len(pools), scale=1, sharex=True, despine=lambda ax: None) as (fig, axs):
            fig.set_size_inches(self.__class__.FIG_SIZE)
            colors = sns.get_color_cycle()[1:len(pools)+1]
            for k, ax, pool, color in zip(reversed(range(len(pools))), axs, reversed(pools), reversed(colors)):
                where = sp.in1d(si, pool)
                times = st[where]
                nr_spikes_at_t = sp.bincount(times) / len(dsyn.pools[1])
                n_spikes_at_t_list.append(nr_spikes_at_t)
                t_for_plot = sp.arange(len(nr_spikes_at_t))/1000

                t_in_range = t_for_plot <= self.MAX_SUM_TRACE_PLOT_T
                ax.plot(t_for_plot[t_in_range], nr_spikes_at_t[t_in_range], clip_on=False, color=color, zorder=100)

                ax.set_xlim(xlim)
                ax.set_xlim(-.5)

                ax.set_ylim(0, 1)
                #style.reduced_ticks_0_1(ax, "y", despine=False)
                #ax.set_ylabel("$P_{{{}}}$".format(k + 1), labelpad=.2)


                ax.spines["left"].set_position(("data", -.5))
                ax.spines["bottom"].set_position(("data", -.2))
                style.set_ticksmiddle_minor(ax, sp.arange(xlim[0], xlim[1] + 1e-5), "x")
                plt.setp(ax.get_xticklabels(), visible=False)
                ax.set_yticks([0, 1])
                ax.set_yticks([0.5], minor=True)
                ax.set_yticklabels([])
            ax.set_yticklabels([0, 1], size=8)

            style.set_ticksmiddle_minor(ax, sp.arange(xlim[0], xlim[1] + 1e-5), "x")

            def try_round(x):
                if int(x) == x:
                    return int(x)
                return x

            ax.set_xticklabels([try_round(xlim[0]), try_round(xlim[1])])
            plt.setp(ax.get_xticklabels(), visible=True)
            ax.set_xlabel(style.latex_name("time") + "\, (ms)$")
            fig.tight_layout(h_pad=.00005)
            fig.text(.15, .65, style.latex_name("fractional_pool_activation"), rotation=90, va="center", ha="center")
            sns.despine(fig)

    def plot_spikes_all_groups(self):
        import operator
        from functools import reduce
        import itertools

        dsyn = make_dsyn()
        ilist, t = dsyn.spike_i, dsyn.spike_t
        pools = list(map(set, dsyn.pools[1:]))
        selected_pools = list(range(len(pools)))
        pools = [pools[k] for k in selected_pools]
        from copy import deepcopy
        pools_complete = deepcopy(pools)

        subsample = 15
        for k in reversed(range(len(pools))):
            pools[k] = reduce(operator.sub, pools[:k], pools[k])


        concatenated = list(itertools.chain(*pools))[::subsample]
        pools_complete_subsampled = [[p for p in pool if p in concatenated] for pool in pools_complete]

        re_enumeration = {original: k for k, original in enumerate(concatenated)}
        max_pool_ind = max(map(lambda x: re_enumeration[x], pools_complete_subsampled[-1]))

        where = sp.array([i in concatenated for i in ilist])
        i_selected = ilist[where]
        i_selected = sp.array([re_enumeration[i] for i in i_selected])
        t_selected = t[where] * 1000

        with self.plot("spike_raster_groups", despine=lambda ax: None) as (fig, ax):
            fig.set_size_inches(self.__class__.FIG_SIZE[0]*.7, self.__class__.FIG_SIZE[1]*.7)
            colors = sns.get_color_cycle()[1:len(pools_complete_subsampled)+1]
            pool_group_block_offset = -1.5
            pool_membership_spacing = .5
            x_marker_pos = [(-len(pools_complete_subsampled) + k) * pool_membership_spacing + pool_group_block_offset
                            for k in range(len(pools_complete_subsampled))]
            for k, (color, pool) in enumerate(zip(colors, pools_complete_subsampled)):
                ax.scatter(sp.zeros(len(pool)) + x_marker_pos[k],
                           [re_enumeration[n] for n in pool], color=color, s=5, clip_on=False)


            ax.scatter(t_selected, i_selected, color="k", s=5, clip_on=False)
            xlim = 0, sp.ceil(t_selected.max())
            ax.set_xlim(xlim)
            ax.set_ylim(0, max_pool_ind)
            ymin, ymax = ax.get_ylim()
            ax.set_yticks(ax.get_ylim())
            ax.set_yticks(sp.arange(ymin, ymax+1, 5), minor=True)
            style.set_ticksmiddle_minor(ax, sp.arange(xlim[0], xlim[1] + 1e-5), "x")

            ax.set_xlabel(style.latex_name("time") + "\, (ms)$", labelpad=style.default_label_pad_reduced)

            y_marker_offset = -1
            ax.text(x_marker_pos[0], y_marker_offset, "$P_1$", ha="center", va="top", color=colors[0])
            ax.text(x_marker_pos[len(x_marker_pos)//2]-.1, -13, "Pool\nmembership", ha="center", va="top")
            ax.text(x_marker_pos[-1], y_marker_offset, "$P_{{{}}}$".format(len(x_marker_pos)), ha="center", va="top",
                    color=colors[len(x_marker_pos)-1])

            ax.spines["top"].set_visible(False)
            ax.spines["right"].set_visible(False)
            ax.spines["left"].set_position(("data", pool_group_block_offset
                                             - pool_membership_spacing * (len(x_marker_pos) + 1.5)))
            ax.spines["bottom"].set_position(("data", -1))
            ax.xaxis.tick_bottom()
            ax.yaxis.tick_left()
            ax.set_ylabel("Neuron", labelpad=style.default_label_pad_reduced)


            ax.legend()

        return ax.get_xlim()


    def plot_spikes_seven_groups(self, dsyn):
        import itertools
        import operator
        from functools import reduce

        dsyn = make_dsyn()
        i, t = dsyn.spike_i, dsyn.spike_t
        pools = list(map(set, dsyn.pools[1:]))
        s = len(pools[0])
        n_pools = len(pools)

        selected_pool_indices = [0, 1, len(pools)-1]
        combinations = list(itertools.chain(*[list(itertools.combinations(selected_pool_indices, k))
                                              for k in range(1, len(selected_pool_indices) + 1)]))
        intersections = [reduce(operator.and_, (pools[ind] for ind in indices))
                         for indices in combinations]

        unique_subsets = [reduce(operator.sub, (pool for k, pool in enumerate(pools)
                                                if k in selected_pool_indices and
                                                k not in combination), intersection)
                          for combination, intersection in zip(combinations, intersections)]
        doubled = any(itertools.starmap(operator.and_, itertools.permutations(unique_subsets, 2)))
        assert not doubled, "Reduction problem"

        re_enumeration = {neuron: k for k, neuron in enumerate(itertools.chain(*unique_subsets))}
        unique_subsets_relabeled = [list(map(lambda x: re_enumeration[x], us)) for us in unique_subsets]

        where = sp.array([i in re_enumeration for i in i])
        i_selected = i[where]
        t_selected = t[where]
        i_relabeled = sp.array(list(map(lambda x: re_enumeration[x], i_selected)))

        with self.plot("overlapping_activity") as (fig, ax):
            fig.set_size_inches((3, 4))
            color = itertools.cycle(sns.color_palette("hls", 7))
            marker = itertools.cycle([".", "o"])
            for subs, comb in zip(unique_subsets_relabeled, combinations):
                where = sp.array([i in subs for i in i_relabeled])
                this_marker = next(marker)
                this_color = next(color)
                ax.scatter(t_selected[where] * 1000, i_relabeled[where],
                           edgecolor=this_color if this_marker == "o" else "w",
                           facecolor="w" if this_marker == "o" else this_color,
                           label=syn_pool_label(comb, selected_pool_indices),
                           s=20 if this_marker == "o" else 80,
                           linewidth=1 if this_marker =="o" else .2,
                           marker=this_marker,
                           clip_on=False)
            ax.legend(loc="center left", bbox_to_anchor=(1, .5))
            ax.set_xlabel("$t\;(ms)$")
            ax.set_ylabel("Neuron")
            ax.set_xlim(-2)
            ax.set_ylim(0)
            sns.despine(fig)


@memory.cache
def api_data(feature_space_dimension=10):
    from connectome.model import API
    from connectome.function import TuningTask

    M = API
    feature_space_dimension = 10
    n_pow = 5
    model = M(nr_neurons=2000, inh_ratio=.1, p_exc=.2,
            feature_space_dimension=feature_space_dimension, p_inh=.5,
            n_pow=n_pow)

    res = TuningTask(return_arrays=True)(model)
    return res


@memory.cache
def fever_data(feature_space_dimension, seed):
    from connectome.model import ERFEVER
    from connectome.function import TuningTask
    import random
    random.seed(seed)
    sp.random.seed(seed)

    feature_space_dimension = feature_space_dimension
    feverization_ratio = .8
    model = ERFEVER(nr_neurons=2000, inh_ratio=.1, p_exc=.2,
            feature_space_dimension=feature_space_dimension, p_inh=.5,
                    feverization_ratio=feverization_ratio)

    res = TuningTask(return_arrays=True)(model)
    return res


def alpha_cmap(color):
    return mpl.colors.LinearSegmentedColormap.from_list("mycm",[color + (0,), color + (1,)])


class APITuning(FolderPlot):
    C_SIM_XLABEL = style.latex_name("cosine_similarity") + "$(FV, Stim)$"
    DYNAIMC_TUNING_STYLE = dict(facecolor=sns.get_color_cycle()[0], edgecolor="w", s=5)
    from .api_normalization import const as api_const
    name_prefix = "api"
    linewidth = .2

    @staticmethod
    def plot_dynamic_api_tuning(ax, pool_activity, cosine_similarity_to_stimulus, xlabel=None,
                                c=None, autonormalize=api_const):
        pool_activity = pool_activity.copy()
        if autonormalize is True:
            autonormalize = sp.absolute(pool_activity).max()
        pool_activity /= autonormalize

        ax.scatter(cosine_similarity_to_stimulus, pool_activity[-1],
                   label=r"$a(t_{\mathrm{end}})$",
                   clip_on=False,
                   linewidth=APITuning.linewidth,
                   **APITuning.DYNAIMC_TUNING_STYLE, c=c)
        APITuning.set_ax_lim(ax, autonormalize)
        ax.set_xlabel(xlabel if xlabel is not None else APITuning.C_SIM_XLABEL,
                      fontsize=config.maybe_modify_font_size(8))
        ax.set_yticks([-1, 0, 1])
        ax.set_yticklabels([style.latex_name.long("inhibition") + " 1", 0, style.latex_name.long("excitation") + " 1"],
                           fontsize=config.maybe_modify_font_size(8))


    @staticmethod
    def set_ax_lim(ax, minmax):
        ax.set_xlim(-1, 1)

        ax.set_ylim(-minmax, minmax)
        ax.set_xticks([-1, 0, 1])
        ax.set_xticklabels([-1, 0, 1], fontsize=config.maybe_modify_font_size(8))
        ax.set_yticks([-minmax, 0, minmax])

    def res(self):
        return api_data()

    def run(self):
        res = self.res()
        cosine_similarity_to_stimulus = res["cosine_similarity_to_stimulus"]
        thalamic_input = res["thalamic_input"]
        pool_activity = res["pool_activity"]
        feature_vectors = res["feature_vectors"]
        stimulus = res["stimulus"]
        # feature_vectors.shape == (feature_space_dimension, neurons)
        # pool_activity.shape == (timesteps, neurons)
        # cosine_similarity_to_stimulus.shape == (neurons,)
        # thalamic_input.shape == (neurons,)

        self.plot_tuning_scatter(cosine_similarity_to_stimulus, pool_activity, thalamic_input)
        self.plot_tuning_curve(cosine_similarity_to_stimulus, pool_activity, thalamic_input)
        self.plot_tuning_histogram(cosine_similarity_to_stimulus, pool_activity, thalamic_input)
        self.plot_tuning_examples(cosine_similarity_to_stimulus, pool_activity, thalamic_input, feature_vectors, stimulus)


    @staticmethod
    def format_y_axis(ax, min_max_f):
        ax.spines["bottom"].set_position("zero")
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.set_xticks([])
        ax.set_ylim(-min_max_f, min_max_f)
        ax.set_yticks([-min_max_f * .8, min_max_f * .8])
        ax.set_yticklabels(["−", "+"])

    def plot_tuning_examples(self, cosine_similarity_to_stimulus, pool_activity, thalamic_input, feature_vectors, stimulus):
        order = sp.argsort(cosine_similarity_to_stimulus)
        high_sim_idx = order[-2]
        mid_sim_idx = sp.argmin(sp.absolute(cosine_similarity_to_stimulus))
        low_sim_idx = order[0]
        with self.plot("tuning_examples", nrows=3, despine=lambda x: None,
                       sharex=True, sharey=True, scale=(.6, 1)) as (fig, axes):
            indices = [high_sim_idx, mid_sim_idx, low_sim_idx]
            colors = ["C1", "C2", "C3"]
            min_max_f = sp.absolute(feature_vectors[:,indices]).max()
            for ax, idx, color in zip(axes, indices, colors):
                ax.plot(stimulus, label="Stimulus", zorder=100)
                ax.plot(feature_vectors[:, idx],
                        label=f"{cosine_similarity_to_stimulus[idx]:.1f}",
                        color=color)
                ax.set_title(f"{sp.arccos(cosine_similarity_to_stimulus[idx]) / sp.pi * 180:.0f}°",
                             loc="right",
                             y=1,
                             x=1.2,
                             va="top",
                             size="small",
                             color=color)
                self.format_y_axis(ax, min_max_f)
            activation_ax = fig.add_axes([1, .15, .5, 1 - .15*2])
            last_activity = pool_activity[-1]
            activities = [last_activity[idx] for idx in indices]
            activation_ax.barh(list(reversed(sp.arange(len(activities)))),
                               activities,
                               .3,
                               color=colors)
            activation_ax.spines["top"].set_visible(False)
            activation_ax.spines["right"].set_visible(False)
            activation_ax.spines["left"].set_position("zero")
            min_max_a = .18
            activation_ax.set_xlim([-min_max_a, min_max_a])
            activation_ax.set_xticks([-min_max_a*.8, min_max_a*.8])
            activation_ax.set_xticklabels(["−", "+"])
            activation_ax.set_yticks([])
            fig.text(0, -.1, "Feature")
            fig.text(1, -.1, "Response")
            fig.text(-0.2, .5, "Strength", rotation=90, ha="center",
                     rotation_mode="anchor")
            #fig.legend(handles=[Line2D([],[])] + [Line2D([],[], color=c) for c in colors],
            #           labels=["Stimulus, 0°"] + [f"≈{d:.0f}°" for d in sp.arccos(cosine_similarity_to_stimulus[indices]) / sp.pi * 180],
            #           loc="center left", bbox_to_anchor=(2.2,.7))



    def plot_tuning_histogram(self, cosine_similarity_to_stimulus, pool_activity, thalamic_input):
        with self.plot(f"{self.name_prefix}i_tuning_hist",
                       scale=.6) as (fig, ax):
            sns.jointplot(cosine_similarity_to_stimulus, pool_activity[-1], ax=ax,kind ="kde")
            ax.set_xlabel(APITuning.C_SIM_XLABEL)
            ax.set_ylabel(r"$a(t_{\mathrm{end}})$")

    def plot_tuning_curve(self, cosine_similarity_to_stimulus, pool_activity,
                                 thalamic_input):
        from sklearn.gaussian_process import GaussianProcessRegressor
        gp = GaussianProcessRegressor()
        sort_indices = sp.argsort(cosine_similarity_to_stimulus)
        last_activity = pool_activity[-1][sort_indices]
        to_angle = lambda x: sp.arccos(x) / sp.pi * 180
        cosine_similarity_to_stimulus = cosine_similarity_to_stimulus[sort_indices]


        gp.fit(sp.atleast_2d(cosine_similarity_to_stimulus).T, sp.atleast_2d(last_activity).T)
        x_pred = cosine_similarity_to_stimulus
        y_pred = gp.predict(sp.atleast_2d(x_pred).T).flatten()

        df = pd.DataFrame({"x": cosine_similarity_to_stimulus, "res": y_pred - last_activity})
        df["bins"] = pd.cut(cosine_similarity_to_stimulus, 10)
        res_std = df.groupby("bins")["res"].std()
        res_std.index = [interval.mid for interval in res_std.index]
        error_y = gp.predict(res_std.index[:, None]).flatten()

        with self.plot(f"{self.name_prefix}_tuning_curve", despine=lambda x:None,
                       scale=.6) as (fig, ax):
            ax.plot(to_angle(x_pred), y_pred, zorder=-100)
            ax.errorbar(to_angle(res_std.index), error_y, yerr=res_std,
                        linewidth=0, elinewidth=1, c="gray")
            ax.set_xlabel("Relative orientation")
            ax.set_ylabel("Response")
            ax.set_xlim(180, 0)
            ax.set_ylim(-.2, .2)
            min_max = .2 #sp.absolute(y_pred).max()
            ax.set_ylim(-min_max, min_max)
            ax.yaxis.tick_right()
            ax.yaxis.set_ticks_position('right')
            ax.yaxis.set_label_position("right")
            ax.spines["left"].set_visible(False)
            ax.spines["top"].set_visible(False)
            ax.set_yticks([-min_max, 0, min_max])
            ax.set_yticklabels(["- Max", "0", "+Max"])
            ax.set_xticks([180, 90, 0])
            ax.set_xticklabels(["180°", "90°", "0°"])


    def plot_tuning_scatter(self, cosine_similarity_to_stimulus, pool_activity,
                            thalamic_input):
        for leg_top in [True, False]:
            suffix = "" if leg_top else "_leg_right"
            with self.plot(f"{self.name_prefix}i_tuning" + suffix,
                           scale=.6) as (fig, ax):
                dict(facecolor=sns.get_color_cycle()[0], edgecolor="w", s=5)
                APITuning.plot_dynamic_api_tuning(ax, pool_activity.copy(),
                                                  cosine_similarity_to_stimulus.copy())
                c2 = sns.get_color_cycle()[1]
                # ax2 = ax.twinx()
                ax.scatter(cosine_similarity_to_stimulus, thalamic_input,
                           s=self.__class__.DYNAIMC_TUNING_STYLE["s"],
                           linewidth=APITuning.linewidth,
                           edgecolor=c2, facecolor="w",
                           label=style.latex_name.full("input_activity"),
                           clip_on=False)
                if leg_top:
                    ax.legend(loc="lower left", bbox_to_anchor=(-.5, 1.1),
                              ncol=2,
                              columnspacing=.5, handletextpad=.2,
                              handlelength=0, borderaxespad=0, borderpad=0)
                else:
                    ax.legend(loc="center left", bbox_to_anchor=(1, .5))


class FEVERTuning(APITuning):
    name_prefix = "fever"
    feature_space_dimension = 10

    def res(self):
        return fever_data(self.feature_space_dimension, 1)


@memory.cache
def sorn_example(eta_stdp, eta_intrinsic):
    from connectome.model import SORN
    np.random.seed(42)

    sorn = SORN(nr_neurons=2000, inh_ratio=.1, p_exc=.2, p_inh=.5,
                eta_stdp=eta_stdp, eta_intrinsic=eta_intrinsic)

    class ActivityRecorder:
        def __init__(self):
            self.activity = []

        def record(self, t, name, value):
            if name == "current_active_neurons":
                self.activity.append(value.copy())
                if t % 10 == 0:
                    print("t={}: {:.2f}".format(t, t/sorn.NR_TIME_STEPS))

    sorn.recorder = ActivityRecorder()
    sorn()
    return sp.asarray(sorn.recorder.activity), sorn.nr_exc


class SORNExamples(FolderPlot):
    good = (.0006, .05)
    bad = (.01, .75)
    good_color = "C1"
    bad_color = "C2"

    def run(self):
        good_activity, nr_exc = sorn_example(*self.good)
        for colored_legend in [True, False]:
            suff = "_colored_legend" if colored_legend else ""
            for with_legend in [True, False]:
                with self.plot(f"good_activity_leg={with_legend}" + suff, scale=.5) as (fig, ax):
                    visualize_activity(fig, ax, good_activity, nr_exc, with_legend, self.good_color, colored_legend)

                bad_activity, nr_exc = sorn_example(*self.bad)
                with self.plot(f"bad_activity_leg={with_legend}" + suff, scale=.5) as (fig, ax):
                    visualize_activity(fig, ax, bad_activity, nr_exc, with_legend, self.bad_color, colored_legend)


class SORNSaneActivity(FolderPlot):
    POPULATION_CORRELATION_LABEL = style.latex_name("median_correlation")

    def run(self):
        df = pd.read_hdf(os.path.join(self.data_folder(), "sorn_corr", "sorn_corr.h5"), "sorn_corr")
        df.rename(columns={"eta": "eta_stdp", "corr": "pair_correlation"}, inplace=True)

        good_corr = pair_correlation(sorn_example(*SORNExamples.good)[0])
        bad_corr = pair_correlation(sorn_example(*SORNExamples.bad)[0])

        from connectome.model import default_definitions
        from itertools import accumulate
        limits_stdp = list(accumulate(
            default_definitions()["models"]["model_specific_parameters"]
            ["SORN"]["eta_stdp"]["args"]))
        limits_intrinsic = list(accumulate(
            default_definitions()["models"]["model_specific_parameters"]
            ["SORN"]["eta_intrinsic"]["args"]))

        s0, s1 = limits_stdp
        i0, i1 = limits_intrinsic
        df_prior_support = df.query(f"{s0} < eta_stdp < {s1} & {i0} < eta_intrinsic < {i1}")
        try:
            os.makedirs(self.output_folder())
        except FileExistsError:
            pass
        df_prior_support.describe().to_excel(os.path.join(self.output_folder(), "prior_filtered.xlsx"))

        eta_intrinsic_c_max = .4

        def format_cbar(cbar, name):
            cbar.set_label(name, labelpad=style.default_label_pad_reduced-4)
            cbar.outline.set_linewidth(0)
            #cbar.ax.minorticks_on()
            cbar.set_ticks([0, eta_intrinsic_c_max])
            #cbar.ax.set_yticks([.1, .2, .3], minor=True)
            x_intrinsic = 2.5
            for limit in [.25/2, .25]:
                cbar.ax.plot([0, x_intrinsic], [limit] * 2, color="k", clip_on=False,
                             solid_capstyle="butt", linewidth=1, dashes=(.5, .5))

        def format_ax(ax):
            for lim in limits_stdp:
                ax.axvline(lim, dashes=style.linestyle.dashed, color="k")

            ax.set_xlim(0, .01)
            ax.set_xticks([0, 0.01])
            ax.set_xticklabels(["0", "0.01"])
            #ax.set_xticks([1e-3, 2e-3, 3e-3], minor=True)
            #ax.xaxis.set_major_formatter(style.formatter.ExpFormatter())


        corr_scale = .8
        with self.plot("sorn_corr", scale=corr_scale) as (fig, ax):
            style.two_d_measure_scatter(df, ax, fig, "eta_stdp", "pair_correlation",  "eta_intrinsic",
                                        style_cbar=format_cbar, tight_ax_lims=True,
                                        vmin=0, vmax=eta_intrinsic_c_max, s=10)
            ax.scatter([SORNExamples.good[0]], [good_corr], clip_on=False, s=13,
                       edgecolor=SORNExamples.good_color, facecolor="w", zorder=20,
                       linewidth=1)
            ax.scatter([SORNExamples.bad[0]], [bad_corr], clip_on=False, s=13,
                       edgecolor=SORNExamples.bad_color, facecolor="w", zorder=20,
                       linewidth=1)
            ax.set_xlabel(style.latex_name("eta_stdp"), labelpad=0)
            ax.set_ylabel(self.POPULATION_CORRELATION_LABEL)
            ax.set_ylim(0, .4)
            ax.set_yticks([0, .4])
            ax.set_yticks([.1, .2, .3], minor=True)
            format_ax(ax)

        size = list(mpl.rcParams["figure.figsize"])
        size[1] /= 2
        size[0] /= 3

        for sse_with_corr_scale in [True, False]:
            suff = "_corr_scale" if sse_with_corr_scale else ""
            scale = corr_scale if sse_with_corr_scale else size
            with self.plot("sorn_sse" + suff, scale=scale) as (fig, ax):
                cbar = style.two_d_measure_scatter(df, ax, fig, "eta_stdp",
                                                   "sse",  "eta_intrinsic",
                                            style_cbar= format_cbar, tight_ax_lims=True,
                                            vmin=0, vmax=.5, s=10)
                fig.delaxes(cbar.ax)
                format_ax(ax)
                ax.set_yticks([0, 1])
                ax.set_ylabel(style.latex_name("sse"), labelpad=-8)
                style.remove_middle_tick_labels(ax, "y")


class FEVERMemory(FolderPlot):
    model_short_name = "fever"
    csim_label = "$" + style.latex_name("cosine_similarity", dollar=False) + "(FV_0, FV_{\\mathrm{end}})$"
    hdf5_file = "fever_memory.h5"
    hdf5_dataset = "fever_memory"
    hdf5_file_heatmap = "fever_memory_heatmap.h5"
    hdf5_dataset_heatmap = "fever_memory_heatmap"
    hdf5_folder = "fever_memory"
    free_par = "feverization_ratio"

    def run(self):
        df_heat = self.read_df()


        with self.plot(self.model_short_name + "_memory_heatmap", scale=(1, 1.4)) as (fig, ax):
            df_heat_plot = (df_heat
                            .groupby(
                [self.free_par, "feature_space_dimension"])
                       ["cosine_similarity"]
                       .mean()
                       .unstack(0)
                       .sort_index(ascending=False))
            mappable = ax.matshow(df_heat_plot.as_matrix(), cmap="viridis",
                                  aspect=1)
            ax.set_aspect(.8)
            cbar = fig.colorbar(mappable,
                                label=self.csim_label,
                                shrink=.8)
            cbar.outline.set_visible(False)
            cbar.set_ticks([0, .2, .4, .6, .8])
            cbar.set_ticklabels(["0", "0.2", "0.4", "0.6", "0.8"])
            ax.set_xlabel(style.latex_name(self.free_par),
                          labelpad=0)

            def safe_int(a):
                if int(a) == a:
                    return int(a)
                return a

            ax.set_xticks(
                [0, len(df_heat_plot.columns) // 2, len(df_heat_plot.columns)-1])
            ax.set_xticklabels([safe_int(df_heat[self.free_par].iloc[0]),
                                 f"{safe_int(df_heat[self.free_par].iloc[len(df_heat_plot.columns) // 2]):.1f}",
                                 safe_int(df_heat[self.free_par].iloc[-1])])

            ax.set_ylabel(style.latex_name("feature_space_dimension"),
                          labelpad=style.default_label_pad_reduced-8)
            ax.set_yticks([0, len(df_heat_plot)-1])
            ax.set_yticklabels([df_heat.feature_space_dimension.iloc[-1],
                                df_heat.feature_space_dimension.iloc[0]])

    def read_df(self):
        df_heat = pd.read_hdf(
            os.path.join(self.data_folder(), self.hdf5_folder,
                         self.hdf5_file_heatmap), self.hdf5_dataset_heatmap)
        return df_heat


class APIMemory(FEVERMemory):
    model_short_name = "api"
    hdf5_file = "api_memory.h5"
    hdf5_dataset = "api_memory"
    hdf5_file_heatmap = "api_memory_heatmap.h5"
    hdf5_dataset_heatmap = "api_memory_heatmap"
    hdf5_folder = "api_memory"
    free_par = "n_pow"


class FEVERRepresentation(FolderPlot):
    model = "ERFEVER"

    def run(self):
        sp.random.seed(40)
        with open(os.path.join(self.data_folder(), "fever_and_api_representation_memory", "api_fever_representation_memory.msgpack"), "rb") as f:
            df = pd.read_msgpack(f)
        self.plot_n(3, df)

    def plot_single(self, datum, ax):
        style = "-"
        ax.plot(datum.target_representation, style, c="C0", label="Stimulus", zorder=100)
        ax.plot(datum.final_representation, style, c="C1",
                label="Represented", zorder=50)
        if self.model == "API":
            ax.plot(datum.final_representation
                    / la.norm(datum.final_representation)
                    * la.norm(datum.target_representation)
                    , style,
                    c="C2", label="Represented,\nnormalized",
                    alpha=1)


        # ax.set_title(datum.model_name)
        #ax.set_ylim(-.3, .3)
        # ax.legend(loc="center left", bbox_to_anchor=(1,.5))
        APITuning.format_y_axis(ax, .3)
        return ax

    def plot_n(self, n, df):
        with self.plot(self.model + "_memory_example", sharex=True, sharey=True, nrows=n, despine = lambda x: None) as (fig, axes):
            samples = df[df.model_name == self.model].sample(n, replace=False)
            for ax, sample in zip(axes, samples.iterrows()):
                self.plot_single(sample[1], ax=ax)
                ax.spines["bottom"].set_position("center")
            if self.model == "API":
                axes[n// 2].legend(loc="center left", bbox_to_anchor=(1, .5))
            axes[-1].set_xticks([])
            axes[-1].set_xlabel("Feature", labelpad=15)
            fig.text(-.1, .5, "Strength", rotation=90, rotation_mode="anchor", ha="center")


class APIRepresentation(FEVERRepresentation):
    model = "API"


class APIFeatureSpaceDimension(FolderPlot):
    scale = style.size.xsq * 1
    scatter_size = 5

    def run(self):
        df = pd.read_pickle(os.path.join(self.data_folder(), "api_feature_space_dimension",
                                         "api_feature_space_dimension.pd.pickle"))
        df.rename(columns={"tuning_activity_correlation": "correlation"}, inplace=True)
        scale_fdim_scatter = (self.scale[0]*1.3, self.scale[1] * 1.3)

        with self.plot("api_feature_space_dim", scale=scale_fdim_scatter) as (fig, ax):
            ax.scatter("feature_space_dimension", "correlation", data=df, facecolor="k", edgecolor="w", s=self.scatter_size)
            ax.set_xlabel(style.latex_name("feature_space_dimension"))
            ax.set_ylabel(style.latex_name("correlation"))
            ax.set_xlim(0)
            ax.set_ylim(0, 1)
            ax.set_xticks([0, 100])
            style.reduced_ticks_0_1(ax, "y")
            style.set_ticksmiddle_minor(ax, range(0, 101, 10), "x")

        for inset_scale in [.5, 1]:
            suff = "" if inset_scale == .5 else "-large_inset"
            with self.plot("api_feature_space_dim_neg_cos_sim_corr",
                           scale=scale_fdim_scatter, despine = lambda x: None) as (fig, ax):
                with self.plot("api_feature_space_dim_neg_cos_sim_corr_examples_1" + suff,
                               scale=self.scale*inset_scale, despine=lambda x: None) as (fig2, ax2):
                    with self.plot("api_feature_space_dim_neg_cos_sim_corr_examples_2" + suff,
                                   scale=self.scale*inset_scale, despine=lambda x: None) as (fig3, ax3):
                        ax.scatter("feature_space_dimension", "neg_cos_sim_corr",
                                   data=df, facecolor="k",
                                   s=self.scatter_size,
                                   edgecolor="w", clip_on=False, zorder=101)
                        ax.set_xlabel(style.latex_name("feature_space_dimension"),
                                      labelpad=style.default_label_pad_reduced)
                        ax.set_xlim(0)
                        style.set_ticksmiddle_minor(ax, range(0, 101, 10), "x")

                        style.reduced_ticks_0_1(ax, "y")
                        #style.set_ticksmiddle_minor(ax, [0, 50, 100], "x")

                        ax.set_ylabel(style.latex_name("correlation") + "$|_{(c_{sim} < 0)}$", labelpad=0)

                        from connectome.model import default_definitions
                        min_max = (default_definitions()["models"]
                                   ["model_specific_parameters"]
                                   ["API"]
                                   ["feature_space_dimension"]
                                   ["args"])
                        for val in min_max:
                            ax.axvline(val, dashes=style.linestyle.dashed, color="k", zorder=-200)


                        ax2.set_aspect("equal")
                        ax3.set_aspect("equal")

                        info = []
                        for fdim, this_axis, c in zip([10, 70], (ax2, ax3), sns.get_color_cycle()[2:4]):
                            datum = self.plot_inset(this_axis, df.copy(), fdim, c)
                            plotted_neg_cos_sim = datum["neg_cos_sim_corr"]
                            ax.scatter(datum.feature_space_dimension, plotted_neg_cos_sim, c=c,
                                       s=self.scatter_size,
                                       zorder=102)
                            info.append({"feature_space_dimension": int(datum.feature_space_dimension),
                                         "plotted_neg_cos_sim": float(plotted_neg_cos_sim),
                                         "color": list(c)})
                        #ax2.set_xticklabels([])
                        #ax2.set_xlabel("")

                        sns.despine(fig)
                        sns.despine(fig2)
                        sns.despine(fig3)
        import yaml
        with open(os.path.join(self.output_folder(), "info.yaml"), "w") as f:
            yaml.dump(info, f)


    def plot_inset(self, ax_in, df, fdim, c):
        datum = df[df.feature_space_dimension > fdim].iloc[0]
        from .api_normalization import const as api_const
        APITuning.plot_dynamic_api_tuning(ax_in, datum.pool_activity, datum.cosine_similarity_to_stimulus,
                                          xlabel=None, c=c, autonormalize=api_const)
        return datum


def to_0_1_range(img):
    img -= img.min()
    img /= img.max()


class GrayScaleTextures(FolderPlot):
    def run(self):
        self.make_folder()
        import skimage.io as skio
        from connectome.function.task.texture import TextureData
        t = TextureData()
        example_traces = []
        colors = cycle(texture_palette())
        for k in range(len(t)):
            img_original = t.load_grayscale_image_normalized(k)
            center = list(map(lambda x: x//2, img_original.shape))
            size_x = 250
            size_y = 1000
            img_cropped = img_original[center[0]-size_x//2:center[0]+size_x//2, center[1]-size_y//2:center[1]+size_y//2]
            example_trace = img_original[center[0]:center[0] + 1, center[1] - 250:center[1] + 250].copy()
            example_traces.append(example_trace[0])
            if k == 1:
                width = 10
                img_original[center[0]-width:center[0]+1+width, center[1]-250:center[1]+250] = 1

            to_0_1_range(img_cropped)
            skio.imsave(os.path.join(self.output_folder(), "{}.png".format(k)), img_cropped)
            to_0_1_range(example_trace)
            skio.imsave(os.path.join(self.output_folder(), "{}-example_trace.png".format(k)), example_trace)

            with self.plot("example-trace{}-line".format(k)) as (fig, ax):
                ax.plot(example_trace[0], color=next(colors))
                ax.set_ylim(-2, 2)
                ax.axis("off")

        self.plot_example_traces_stacked(example_traces)

    def plot_example_traces_stacked(self, example_traces):
        offset_scale = 1.5

        with self.plot("example_traces_stacked", despine=no_despine,
                       scale=(TaskTexture.INP_POOL_READOUT_WIDTH, .8)) as (fig, ax):
            from itertools import cycle
            colors = cycle(texture_palette())
            half_texture_range = .5
            for data, cls, offset in zip(example_traces,
                                         range(len(example_traces)),
                                         sp.arange(len(example_traces)) * offset_scale):
                ax.plot(data - offset, label="$T_{{{}}}$".format(cls), color=next(colors))
                if cls in [0, len(example_traces)-1]:
                    ax.text(-25, -offset + half_texture_range,
                            "T$_{{{}}}$".format(cls + 1), ha="right", va="center")
            ax.text(-25, -offset//2 + half_texture_range,
                    "$\\vdots$", ha="right", va="center")
            ax.spines["left"].set_visible(False)
            ax.spines["right"].set_visible(False)
            ax.spines["top"].set_visible(False)
            ax.set_yticks([])
            ax.tick_params("y", left="off", right="off")
            ax.set_xticks([0, 250, 500], )
            ax.set_xticks([125, 375], minor=True)
            ax.tick_params("x", top="off", which="both")
            ax.set_xlabel(TaskTexture.TIME_LABEL)
            ax.spines["bottom"].set_position(("data", -offset - 1))
            style.annotate_group(ax, 250, 500, 2.1, "Training", length=0, text_offset=.6, linewidth=5)


class ABCIllustration(FolderPlot):
    def run(self):
        with self.plot("abc_illustration", 4, despine=lambda x: None) as (fig, axes):
            (ax_prior, ax_proposal, ax_data, ax_posterior) = axes.flatten()
            import scipy.stats as st

            sp.random.seed(7)
            pdf_linewidth = 2
            pdf_color = "grey"

            n_samples_prior = 7
            accepted_region = sp.array([-.1, .5])
            x = sp.linspace(-.7, .7, 10000)
            data = accepted_region.mean()


            prior = st.uniform(-.5, 1)

            def sample_data(x):
                return x + sp.rand() * .2

            prior_samples = prior.rvs(n_samples_prior)
            palette = sp.array(sns.husl_palette(len(prior_samples)))
            sample_data = sp.array(list(map(sample_data, prior_samples)))
            accepted = (sample_data < accepted_region[1]) & (
            (sample_data > accepted_region[0]))
            acc_samples = sample_data[accepted]
            kde = st.gaussian_kde(acc_samples)

            def rug(ax, data, mask, **style):
                ax.scatter(data[mask],
                           sp.zeros_like(data[mask]),
                           c=palette[mask],
                           clip_on=False,
                           zorder=10,
                           **style)

            ax_prior.fill(x, prior.pdf(x), color=pdf_color,
                          linewidth=pdf_linewidth)

            data_size = 2
            rug(ax_proposal, prior_samples,
                sp.ones_like(prior_samples).astype(bool), s=data_size)

            ax_data.plot(accepted_region, [0, 0], zorder=-1, color="k",
                         linewidth=2,
                         solid_capstyle="butt")

            rug(ax_data, sample_data, accepted, marker="|")
            rug(ax_data, sample_data, ~accepted, marker="x")
            ax_data.scatter([data], [0], marker="|", color="k")
            # ax_data.annotate("Observation", (data, .001), (.9, 1.7), textcoords="axes fraction",
            #                 arrowprops={"arrowstyle": "-|>", "color": "k"})

            ax_posterior.fill(x, kde.pdf(x) * prior.pdf(x), color=pdf_color,
                              zorder=-1, linewidth=pdf_linewidth,
                              clip_on=False)
            rug(ax_posterior, prior_samples, accepted, s=data_size)

            # for ax in ()

            for ax in (ax_prior, ax_proposal, ax_data, ax_posterior):
                ax.set_xlim(x[0], x[-1])
                ax.spines["bottom"].set_visible(True)
                ax.spines["top"].set_visible(False)
                ax.spines["left"].set_visible(False)
                ax.spines["right"].set_visible(False)
                ax.set_xticks([])
                ax.set_yticks([])

            ax_prior.set_title("Prior $p(\\theta)$")
            ax_proposal.set_title("Samples $\\theta$")
            ax_data.set_title("Data $D$")
            ax_posterior.set_title("Posterior $p(\\theta|D)$")

            xlabelpad = 4
            for ax in (ax_prior, ax_proposal, ax_posterior):
                ax.set_xlabel("$\\theta$", labelpad=xlabelpad)
            ax_data.set_xlabel("D", labelpad=xlabelpad)

            fig.set_size_inches(1.3, 2.3)
            fig.tight_layout(h_pad=0)



class FEVEROriginalVsModifiedFEVERRuleSatisfactionPLot(FolderPlot):
    def requires(self):
        return FEVEROriginalVsModifiedRuleSatisfaction()

    def run(self):
        df = self.input().load()
        df["kind"] = df["p_init"].apply(lambda x: "Orig." if x == 0 else "Adap.")
        os.makedirs(self.output_folder(), exist_ok=True)
        with open(os.path.join(self.output_folder(), "info.txt"), "w") as my_file:
            my_file.write(str(df.groupby("kind").mean()))

        with self.plot("fever_rule", scale=(1, .3)) as (fig, ax):
            sns.boxplot(x="kind", y="norm", ax=ax, data=df,
                        #facecolor="w",
                        #errcolor="k",
                        #edgecolor="k",
                        linewidth=1)
            ax.set_ylim(0, 5)
            ax.set_ylabel("Error")
            ax.set_xlabel("")
            fig.set_clip_on(False)


class FEVERNetworkRepresentationPolar(FolderPlot):
    def run(self):
        from connectome.model import ERFEVER
        from connectome.function.task.memory import FEVERDynamics
        sp.random.seed(42)

        @memory.cache
        def make_fever(seed):
            sp.random.seed(seed)
            erfever = ERFEVER(feverization_ratio=.7,
                              feature_space_dimension=150,
                              nr_neurons=2000,
                              inh_ratio=.1,
                              p_exc=.2, p_inh=.5)
            network = erfever()
            feature_vectors = erfever.get_feature_vectors()
            feverdynamics = FEVERDynamics(network, feature_vectors)
            feverdynamics.run(15)
            return feverdynamics.target_representation_, feverdynamics.end_representation_



        from matplotlib.patches import FancyArrowPatch
        from mpl_toolkits.mplot3d import proj3d

        class Arrow3D(FancyArrowPatch):
            def __init__(self, xs, ys, zs, *args, **kwargs):
                FancyArrowPatch.__init__(self, (0, 0), (0, 0), *args, **kwargs)
                self._verts3d = xs, ys, zs

            def draw(self, renderer):
                xs3d, ys3d, zs3d = self._verts3d
                xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
                self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
                FancyArrowPatch.draw(self, renderer)

        fig = plt.figure()
        scale = 2
        fig.set_size_inches((4 / 3 * scale, scale))
        ax = fig.add_axes([.2, .3, .4, .5], projection='3d')
        dimensions = sp.random.randint(0, 150, 3)

        ax.scatter([0], [0], [0], c="black")

        legend_proxy = []
        legend_labels = []
        import numpy.linalg as la

        for seed, ha, va in zip([41, 42, 43], ["right", "right", "right"], ["bottom", "bottom", "top"]):
            target_raw, end = make_fever(seed)
            target = target_raw[dimensions].flatten()
            end = end[dimensions].flatten()

            csim = sp.dot(target, end) / la.norm(target) / la.norm(end)

            ax.plot(*zip([0] * 3, target), linewidth=0)
            ax.plot(*zip([0] * 3, end), linewidth=0)
            c1, c2, *rest = sns.get_color_cycle()

            a1 = Arrow3D(*zip([0]*3, target), arrowstyle="-|>", lw=1, mutation_scale=20, color=c1)
            ax.add_artist(a1)

            a2 = Arrow3D(*zip([0]*3, end), arrowstyle="-|>", lw=1, mutation_scale=20, color=c2)
            ax.add_artist(a2)

            #ax.text(*target, "$c_{{sim}}={:.3f}$".format(csim), ha=ha, va=va)



        legend_labels += ["$FV_0$", "$FV_{\mathrm{end}}$"]
        legend_proxy += [mpl.lines.Line2D((0,0), (0,0), color=c) for c in (c1, c2)]


        ax.locator_params(nbins=2)
        ax.legend(legend_proxy,
                  legend_labels,
                  loc="lower left", bbox_to_anchor=(config.legend_x_transform(.8), .7),
                  labelspacing=0)

        fmt = "$d_{{{}}}$"
        labelpad = -16
        with config.force_pad():
            ax.set_xlabel(fmt.format(dimensions[0]), labelpad=labelpad)
            ax.set_ylabel(fmt.format(dimensions[1]), labelpad=labelpad)
            ax.set_zlabel(fmt.format(dimensions[2]), labelpad=labelpad)
        ax.tick_params(axis='both', which='major', pad=-2)

        ax.set_title(style.Tex(style.latex_name("feature_space_dimension", dollar=False) + "={}".format(len(target_raw))),
                     y=config.title_y_transoform(.95))

        ax.set_xticklabels([])
        ax.set_yticklabels([])
        ax.set_zticklabels([])

        self.make_folder()
        ax.figure.savefig(os.path.join(self.output_folder(), "representation_3d." + self.extension), transparent=True,
                          bbox_inches="tight")
        plt.close(ax.figure)


class Colors(FolderPlot):
    def run(self):
        with self.plot("colors") as (fig , ax):
            ax.legend(*zip(*[[mpl.patches.Patch(color=color), name] for name, color in style.color.items()]))
            ax.axis("off")

        for name, num in [("cycle", 7)]:
            with self.plot("{} {}".format(name, num)) as (fig, ax):
                ax.legend(*zip(*[[mpl.patches.Patch(color=color), name + " " + str(k)] for k, color in enumerate(sns.get_color_cycle())]))
                ax.axis("off")


class Names(FolderPlot):
    def run(self):
        import subprocess
        import os
        for name, short in style.latex_name.items():
            uncropped_name = name + "_uncropped"
            with self.plot(uncropped_name) as (fig, ax):
                center = .5, .5
                ax.text(*center, short, ha="center", va="center")
                ax.add_patch(mpl.patches.Circle(center, .25, facecolor="none", linewidth=style.linewidth))
                ax.axis("off")
                ax.set_aspect("equal")
            uncropped_file = os.path.join(self.output_folder(), uncropped_name + "." + self.extension)
            subprocess.run(["pdfcrop",
                            uncropped_file,
                            os.path.join(self.output_folder(), name + ".pdf")])
            os.remove(uncropped_file)


class SBMPlot(FolderPlot):
    def run(self):
        source_folder = os.path.join(self.data_folder(), "sbm",)
        nrows_orig = 1
        ncols_orig = 4


        def no_mod(x):
            return x

        def shuffle(x):
            n = x.shape[0]
            permutation = sp.random.choice(n, n)
            return x[permutation].T[permutation].T

        panels = {"main": [("LL", shuffle), ("LL", no_mod), ("ER", no_mod), ("SYN", no_mod)],
                  "supplements": [("EXP", no_mod), ("SORN", no_mod), ("ERFEVER", no_mod), ("API", no_mod)]}

        for panel_name, contents in panels.items():
            for scale in [1, .8, 0.75]:
                if scale != 1:
                    suffix = f"_scale={scale}"
                    nrows, ncols = 4, 1
                    row_extra_space = 1.3
                else:
                    nrows, ncols = nrows_orig, ncols_orig
                    suffix = ""
                    row_extra_space = 0
                with self.plot(panel_name + suffix,
                               nrows=nrows, ncols=ncols,
                               sharex=False, sharey=False,
                               despine=lambda x: None,
                               squeeze=False) as (fig, axes):
                    #if nrows == 1:
                    #    axes = axes[None,:]
                    fig.set_size_inches(scale * ncols + .4 * (ncols - 1), scale * nrows + row_extra_space)  # the + in width for the ticklabels
                    ax_iter = iter(axes.flatten())

                    for file, modifier in contents:
                        ax = next(ax_iter)
                        adj = sp.load(os.path.join(source_folder, file + ".npy"))
                        adj = modifier(adj)
                        ax.matshow(adj, cmap="Greys", aspect='equal')
                        name = os.path.splitext(os.path.basename(file))[0]
                        ax.set_title(f"$C_{{\\mathrm{{{style.name[name]}}}}}$", y=.98)

                    for ax in axes.flatten():
                        ax.xaxis.tick_bottom()
                        ax.yaxis.tick_left()
                        ax.set_adjustable('box-forced')
                        ax.set_xticks([])
                        ax.set_yticks([])

                    for ax in ax_iter:
                        ax.axis("off")

                    for ax in axes.flatten():
                        ax.set_aspect("equal")

                    for ax in axes[-1,:]:
                        ax.set_xlabel("Post", labelpad=2)
                        ax.set_xticks([0, adj.shape[0]-1])
                        ax.set_xticklabels([0, "$n_e$"])
                        ax.set_yticks([0, adj.shape[0]-1])
                        ax.set_yticklabels([])


                    for ax in axes[:,0]:
                        if scale != 1 and panel_name != "supplements":
                            ax.set_ylabel("Pre", labelpad=style.default_label_pad_reduced)
                            ax.set_yticks([0, adj.shape[0]-1])
                            ax.set_yticklabels([0, "$n_e$"])



                    n_blocks = 4
                    n_exc = 1800
                    sbm_ticks = sp.concatenate(([0], sp.linspace(n_exc/n_blocks/2, n_exc -n_exc/n_blocks/2, 4), [adj.shape[0]-1]), axis=0)
                    sbm_xtick_labels = ["0  "] + [f"$b_{n}$" for n in range(1,5)] + ["   $n_e$"]
                    sbm_ytick_labels = [""] + [f"$b_{n}$" for n in range(1, 5)] + [""]
                    if panel_name == "main":
                        axes.flatten()[1].set_xticks(sbm_ticks)
                        axes.flatten()[1].set_xticklabels(sbm_xtick_labels)
                        axes.flatten()[1].set_yticks(sbm_ticks)
                        axes.flatten()[1].set_yticklabels(sbm_ytick_labels)

                        ticklines = [axes.flatten()[1].xaxis.get_majorticklines(),
                                     axes.flatten()[1].yaxis.get_majorticklines()]
                        for ind in [2, 4, 6, 8]:
                            for tl in ticklines:
                                tl[ind].set_visible(False)


class ReconstructionTimes(FolderPlot):
    #force_replot = True

    def run(self):
        km = 1e3
        mm = 1e-3
        um = 1e-6
        h = 3600

        path_length_density = 10 * km / mm**3
        barrel_volume = (300 * um)**3
        total_barrel_path_length = barrel_volume * path_length_density
        annotation_speed = 1.5 * mm / h
        total_time = total_barrel_path_length / annotation_speed

        self.make_folder()
        with open(os.path.join(self.output_folder(), "barrel_reconstruction_time.txt"), "w") as f:
            f.write(f"{total_time/h}h at path lenth density {path_length_density / (km / mm**3)} km / mm**3")


class SYNPoolPropagationSummary(FolderPlot):
    #force_replot = True

    def run(self):
        df = pd.read_csv(os.path.join(self.data_folder(), "syn_propagation", "syn_propagation.csv"))
        df = df.groupby("pool_size")["fraction_pools_activated"].agg(["mean", "std"])
        with self.plot("syn_summary", scale=(style.size.s[0]*.85, style.size.s[1])) as (fig, ax):
            ax.plot(df["mean"], clip_on=False)
            ax.fill_between(df.index, df["mean"] + df["std"], df["mean"] - df["std"], alpha=.4)
            ax.set_ylim(0, 1)
            ax.set_xlim(80)
            ax.set_xticks([100, 200, 300])
            ax.set_yticks([0, 1])
            ax.set_yticks([0.2, .4, .6, .8], minor=True)
            ax.set_xlabel(style.latex_name("pool_size"))
            ax.set_ylabel(style.latex_name("fractional_chain_activation"), labelpad=style.default_label_pad_reduced)


class SynapseNumbers(FolderPlot):
    def run(self):
        # From feldmeyer 1999
        counts = sp.array([0, 0, 2, 5, 2, 2])
        probs = counts / counts.sum()
        avg_n_syn_per_connection = (probs * sp.arange(len(probs))).sum()
        n_connection = (.2 * 1800 + .6 * 200) * 2000
        n_synapses = n_connection * avg_n_syn_per_connection
        self.make_folder()
        numbers = f"Nr Synapses: {n_synapses}\naverage number of syn per conn: {avg_n_syn_per_connection}"
        with open(os.path.join(self.output_folder(),
                               "synapse_numbers.txt"), "w") as f:
            f.write(numbers)


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

    def output(self):
        return luigi.LocalTarget(self.input().path + ".mean_duration")

    def run(self):
        infile = "sqlite:///" + self.input().path
        abc = pd.read_sql_table("abc_smc", infile)
        mean_duration = (abc.end_time - abc.start_time).mean()

        with open(self.output().path, "w") as f:
            f.write(str(mean_duration))


class SYNGeGiInterpolation(FolderPlot):
    def run(self):
        from connectome.function.task.propagation.dynamicsyn import LUT
        with self.plot("syn_ge_gi_lut") as (fig, ax):
            ax.plot(LUT["pool_size"], LUT["log10_ge"], label="Log-excitatory\nefficacy $\log_{10} g_e$")
            ax.plot(LUT["pool_size"], LUT["log10_gi"], label="Log-inhibitory\nefficacy  $\log_{10}  g_i$")
            ax.set_xlabel(style.latex_name("pool_size"))
            ax.set_ylabel("Log-efficacy")
            ax.legend(loc="center left", bbox_to_anchor=(1, .5))


