from collections import UserDict

import matplotlib as mpl
import matplotlib.pyplot as plt
import scipy as sp
import seaborn as sns
from matplotlib.colors import LinearSegmentedColormap
from seaborn.utils import get_color_cycle
from figures.analysispipeline.config import config, rgb_to_rgb

sns.get_color_cycle = get_color_cycle   # API has chanbeg apparently not in the main seaborn module anymore

bar_line_width = 3


_color  = {"ER":   (  0,   0,   0),         # Black
          "SORN": (  0, 158, 115),     # Bluish green
          "ERF":  (240, 228, 66),  # Yellow
          "API":  (  0, 114, 178),     # Blue
          "EXP":  ( 86, 180, 233),    # Sky blue
          "EXPF": (204, 121, 167),   # Reddish purple
          "FF":   (213,  94,   0),      # Vermillon
          "LL":   (230, 159,   0),     # Orange
          "SF":   (213,  94,   0),      # Vermillon
          "gray": (100, 100, 100),    # Gray
          "black": (  0,   0,   0),
          "light_gray": (220, 220, 220),
          "delete_matrix_row_column": (0, 0, 0)
          }

_color["original_network"] = _color["gray"]
_color["AP"] = _color["API"]
_color["LSM"] = _color["EXP"]
_color["SYN"] = _color["SF"]
_color["FEVER"] = _color["ERF"]
_color["ERFEVER"] = _color["ERF"]
_color["RemoveAddNoiseAndSubsample"] = _color["gray"]


name = {
    "CDF": "CDF",
    "PDF": "PDF",
    "ER": "ER-ESN",
    "ERFEVER": "FEVER",
    "API": "API",
    "SYN": "SYNFIRE",
    "EXP": "EXP-LSM",
    "LL": "LAYERED",
    "noise": "Noise",
    "SYH": "SYNFIRE",
    "joint": "Shared",
    "SORN": "STDP-SORN",
    "model": "Model",
    "adapted": "Adapted"
}
name.update({"FEVER": name["ERFEVER"],
             "RemoveAddNoiseAndSubsample": name["noise"]})


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

_col_update = {}
for key, value in _color.items():
    try:
        _col_update[name[key]] = value
    except KeyError:
        pass

_color.update(_col_update)


class Tex:
    def __init__(self, *args):
        if len(args) == 0:
            self.text = ""
        elif len(args) == 1:
            self.text = str(args[0]).strip("$")
        else:
            self.text = sum(map(type(self), args), type(self)()).text

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return "$" + self.text + "$"

    def __add__(self, other):
        try:
            text = other.text
        except AttributeError:
            text = other
        return Tex(self.text + text)

    def expectation(self):
        return expectation(self)

    def __getitem__(self, item):
        if isinstance(item, slice):
            if item.start == 1 and item.stop == -1:
                return self.text[1:-1]
        raise Exception("Slicing only defined for 1:-1")



def expectation(tex: Tex):
    tex = Tex("\mathbb{E} \left [", tex, "\\right]")
    return tex





mpl.rcParams.update(config.mpl_style)

cmap = mpl.cm.viridis

internal_to_print_name_mapping = {"EXP": "LSM",
                                  "AP": "API"}
print_name_to_internal_papping = {value: key for key, value in internal_to_print_name_mapping.items()}


features_for_classification = ['rec_ee_shuffled', 'in_out_deg_corr', 'rec_ei_shuffled', 'cycles_exc_5']
features_for_display = features_for_classification + ["p_e", "p_i", "rec_ee_unshuffled", "nr_measured_vertices"]
feature_short_name_latex = {-2: 'leaf',
                             0: '$r_{ee}$',
                             1: '$r_{i/o}$',
                             2: '$r_{ei}$',
                             3: '$r^{(5)}$',
                             4: '$p_e$',
                             5: '$p_i$',
                             6: "$r_{ee, unsh}$",
                             7: "$n$"}

names = config.names


names.update({"reciprocity_exc": names["reciprocity_ee"],
                   "subsampling_fraction": names["subsampling"],
                   "fraction_remove_and_add": names["noise"]})

NameDict = config.name_dict()
latex_name = NameDict(names)



for k, name_ in enumerate(features_for_display):
    feature_short_name_latex.update({name_: feature_short_name_latex[k]})


def internal_name_to_print_name(name):
    if name in internal_to_print_name_mapping:
        return internal_to_print_name_mapping[name]
    else:
        return name


def strip_characters_for_save(string):
    return (string.replace("$", "").replace("^","").replace("{","").replace("}","").replace("/", "")
           .replace("(","") .replace(")","").replace(">",""))


class ColorDictionary(UserDict):
    def __missing__(self, key):
        return self["gray"]







def rgb_to_hex(rgb):
    colors = [str(hex(int(channel*255)))[2:] for channel in rgb]
    colors = [col if len(col) == 2 else "0" + col for col in colors ]
    return "#" + "".join(colors)

color = ColorDictionary({key: rgb_to_rgb(*value) for key, value in _color.items()})


# colors from tango icon theme: chocolate and plum
matrix_colors = [(0xe9/255, 0xb9/255, 0x6e/255),
                 (0, 0, 0),
                 (0xad/255, 0x7f/255, 0xa8/255)]
adjacency_matrix_cmap = LinearSegmentedColormap.from_list("ExcInh", matrix_colors, N=3)
adjacency_matrix_subsampling_cmap = LinearSegmentedColormap.from_list("ExcInh", matrix_colors, N=4)


internal_names = ["ER", "SORN", "ERF", "AP", "EXP", "EXPF", "FF", "LL", "SF"]
print_names = list(map(internal_name_to_print_name, internal_names))


plausible_regime = {
    "p": [0.15, 0.25],  # connectivity
    "r": [0.15, 0.35]  # reciprocity
}


vertical_label_angle = -90
linewidth = config.linewidth

heat_map_cmap = None

despine = sns.despine


def mimic_alpha(color: tuple, alpha: float, background_color=(1,1,1)):
    return tuple(c * alpha + (1-alpha) * bc for c, bc in zip(color, background_color))


def make_plausible_regime_r_over_p_rectangle(ax, color="black"):
    ax.plot([plausible_regime["p"][0], plausible_regime["p"][1],
              plausible_regime["p"][1], plausible_regime["p"][0],plausible_regime["p"][0]],
             [plausible_regime["r"][0], plausible_regime["r"][0],
              plausible_regime["r"][1], plausible_regime["r"][1], plausible_regime["r"][0]], color=color)


def make_plausible_regime_r_over_p_2(ax):
    x = plausible_regime["p"]
    upper = [plausible_regime["r"][1]] * 2
    lower = [plausible_regime["r"][0]] * 2
    ax.fill_between(x, upper, lower, color=color["light_gray"], edgecolors=None, label="Biological regime")


def do_nothing(ax):
    pass


make_plausible_regime_r_over_p = do_nothing  # disable due to connectivity discussion



default_label_pad_reduced = -5


def reduced_ticks_0_1(ax, axis="both", despine=True, label_pad_reduce=default_label_pad_reduced):
    major_ticks = [0, 1]
    minor_ticks = [.2, .4, .6, .8]
    major_labels = [0, 1]

    if despine:
        sns.despine(ax=ax)

    if axis == "x" or axis == "both":
        ax.set_xticks(major_ticks)
        ax.set_xticklabels(major_labels)
        ax.set_xticks(minor_ticks, minor=True)
        config.modify_labelpad(ax.xaxis, label_pad_reduce)
        ax.set_xlim(0, 1)

    if axis == "y" or axis == "both":
        ax.set_yticks(major_ticks)
        ax.set_yticklabels(major_labels)
        ax.set_yticks(minor_ticks, minor=True)
        config.modify_labelpad(ax.yaxis, label_pad_reduce)
        ax.set_ylim(0, 1)


def style_r_over_p(ax, view="all"):
    plt.xlabel(latex_name("p_ee"), labelpad=default_label_pad_reduced)
    plt.ylabel(latex_name("reciprocity_ee"), labelpad=default_label_pad_reduced)
    sns.despine(ax=ax)
    ax.set_aspect('equal')

    for which in "xy":
        if view == "all":
            set_ticksmiddle_minor(ax, [0, .2, .4, .6, .8, 1], which)
        if view == "half":
            set_ticksmiddle_minor(ax, [0, .1, .2, .3, .4, .5], which)
        if view == "dot6":
            set_ticksmiddle_minor(ax, [0, .1, .2, .3, .4, .5, .6], which)
    make_plausible_regime_r_over_p_rectangle(ax)


def rotate_pair_grid(g, rotation=45):
    for ax in g.axes[-1]:
        plt.setp( ax.xaxis.get_majorticklabels(), rotation=rotation)


def style_cbar_0_1(cbar, label=None, labelpad=default_label_pad_reduced):
    if label is not None:
        cbar.set_label(label, labelpad=labelpad)
    cbar.set_ticks([0, .2, .4, .6, .8, 1])
    cbar.set_ticklabels([0, "", "", "", "", 1])
    cbar.outline.set_linewidth(0)


def make_cbar(ax, color_list, tick_labels, label=""):
    bounds = sp.arange(len(color_list)+1)

    cmap = mpl.colors.ListedColormap(color_list)
    norm = mpl.colors.BoundaryNorm(bounds, cmap.N)

    cbar = mpl.colorbar.ColorbarBase(ax, cmap=cmap, norm=norm, boundaries=bounds, ticks=bounds + .5)
    cbar.set_ticklabels(tick_labels)
    cbar.outline.set_linewidth(0)
    cbar.set_label(label)
    return cbar


def two_d_measure_scatter(df, ax, fig, x, y, hue, style_cbar=style_cbar_0_1,
                          vmin=0, vmax=1, s=20, tight_ax_lims=True, labelpad=None, clip_on=True,
                          with_cbar=True):

    mappable = ax.scatter(df[x], df[y], c=df[hue], s=s, edgecolor="none", vmin=vmin, vmax=vmax,clip_on=clip_on)

    if tight_ax_lims:
        ax.set_ylim(df[y].min(), df[y].max())
        ax.set_xlim(df[x].min(), df[x].max())
    if labelpad is not None:
        ax.set_xlabel(latex_name(x), labelpad=labelpad)
        ax.set_ylabel(latex_name(y), labelpad=labelpad)
    else:
        ax.set_xlabel(latex_name(x))
        ax.set_ylabel(latex_name(y))

    ax.locator_params(nbins=4)

    if with_cbar:
        cbar = fig.colorbar(mappable, ticks=mpl.ticker.MaxNLocator(4))
        style_cbar(cbar, latex_name(hue))
        cbar.outline.set_linewidth(0)
        return cbar


def vertical_xlabel(ax):
    plt.setp(ax.xaxis.get_majorticklabels(), rotation=vertical_label_angle)


def colorbar(ax, mappable, label, nticks=4, labelpad=None):
    ticker = mpl.ticker.MaxNLocator(nticks)
    cb = ax.figure.colorbar(mappable, ticks=ticker)
    cb.outline.set_linewidth(0)
    cb.set_label(label, labelpad=labelpad)
    return cb


def annotate_group(ax, xmin, xmax, y, text, length=.1, text_offset=0, linewidth=linewidth):
    ax.plot([xmin, xmin, xmax, xmax],[y-length, y, y, y-length], color="k", clip_on=False, linewidth=linewidth, solid_capstyle="butt")
    ax.text((xmin+xmax)/2, y+text_offset, text,  va="bottom" if length >= 0 else "top", ha="center")


def ticks_to_major_minor(ax, axis="both"):
    if axis == "x" or axis == "both":
        xticks = ax.get_xticks()
        ax.set_xticks([xticks[0], xticks[-1]])
        ax.set_xticks(xticks[1:-1], minor=True)
        ax.set_xticklabels([xticks[0], xticks[-1]])

    if axis == "y" or axis == "both":
        yticks = ax.get_yticks()
        ax.set_yticks([yticks[0], yticks[-1]])
        ax.set_yticks(yticks[1:-1], minor=True)
        ax.set_yticklabels(ax.get_yticks())


def remove_middle_tick_labels(ax, axis="both"):
    if axis == "x" or axis == "both":
        xticks = ax.get_xticklabels()
        ax.set_xticklabels(xticks[:1] + [""] * (len(xticks)-2) + xticks[-1:])

    if axis == "y" or axis == "both":
        yticks = ax.get_yticks()
        ax.set_yticklabels([str(yticks[0])] + [""] * (len(yticks) - 2) + [str(yticks[-1])])


def set_ticksmiddle_minor(ax, ticks, axis):
    ticks = list(ticks)
    if axis == "x":
        set_lim = ax.set_xlim
        set_ticks = ax.set_xticks
    if axis == "y":
        set_lim = ax.set_ylim
        set_ticks = ax.set_yticks

    set_lim(ticks[0], ticks[-1])
    set_ticks([ticks[0], ticks[-1]])
    set_ticks(ticks[1:-1], minor=True)



annotarrowprops = dict(arrowstyle="-|>", facecolor="black")
