import contextlib
import matplotlib.pyplot as plt
import os
import matplotlib as mpl
import luigi
import seaborn as sns
#import datashader as ds
#import datashader.transfer_functions as tf
#from datashader.colors import viridis
import pandas as pd
import scipy as sp
from collections import OrderedDict
#import xarray as xr
from matplotlib.figure import Figure
from matplotlib.axes import Axes
from figures import style
from numbers import Number
from figures.analysispipeline.config import config

def despine(ax):
    sns.despine(ax=ax)


def no_despine(ax):
    pass


def modify_0_1(lim):
    """
    Check whether the axis limit is close enough to (0, 1)

    Parameters
    ----------
    lim: tuple of axis limits

    Returns
    -------

    modif: bool

    """
    return lim[0] == 0 and .95 < lim[1] <= 1.15


@contextlib.contextmanager
def plot(output_file: str, *args, show=False, scale=1, despine=despine, **kwargs):
    base_figsize = mpl.rcParams["figure.figsize"]
    if isinstance(scale, Number):
        new_figsize = tuple(map(lambda x: x * scale, base_figsize))
    else:
        new_figsize = scale
    fig, ax = plt.subplots(*args, **kwargs)
    fig.set_size_inches(new_figsize)
    yield (fig, ax)
    if show:
        ext = os.path.splitext(output_file)[-1]
        output_file = os.path.join(os.path.expanduser("~"), "tmp", "show." + ext)
    try:
        os.makedirs(os.path.dirname(output_file))
    except FileExistsError:
        pass
    despine(ax)

    fig.savefig(config.sanitize_filename(output_file), bbox_inches="tight", transparent=True)
    plt.close(fig)


class UnsatisfiedTarget(luigi.target.Target):
    def __init__(self, file):
        self.fn = file

    def exists(self):
        return False


class ReplotMixin:
    force_replot = False

    def complete(self):
        if self.force_replot:
            return False
        return super().complete()


class FolderPlot(ReplotMixin, luigi.Task):
    extension = config.extension
    superfolder = config.plot_superfolder

    def make_folder(self):
        try:
            os.makedirs(self.output_folder())
        except FileExistsError:
            pass

    def data_folder(self):
        return os.path.join(os.path.dirname(os.environ["PLOT_PATH"]), "data")

    def output_folder(self):
        return os.path.join(os.environ["PLOT_PATH"], self.superfolder, self.__class__.__name__)

    def output(self):
        return luigi.LocalTarget(self.output_folder())

    def plot(self, file_name, *args, **kwargs) -> (Figure, Axes):
        self.force_replot = False
        return plot(os.path.join(self.output().path, file_name + "." + self.extension), *args, **kwargs)


def traces_to_png(traces, path, width=900, height=300, how="eq_hist"):
    """
    Paramters
    ---------

    texture_traces : 2d array of shape (nr_time_steps, nr_neurons)

    width: int
        Width of the ong

    height: int
        Height of the png

    how: str
        Normalization. Can be "eq_hist", "linear", "log

    path: str
        Where to store the png.
        Has to end with "png"
    """
    assert path.endswith("png"), "Has to be a png file"

    df = pd.DataFrame(traces)
    df.rename(columns={col: str(col) for col in df.columns}, inplace=True)
    str_cols = df.columns
    df["t"] = sp.arange(len(df)).astype(float)

    cvs = ds.Canvas(x_range=[df.t.min(), df.t.max()],
                    y_range=[df[str_cols].min().min(), df[str_cols].max().max()],
                    plot_height=height, plot_width=width)
    aggs = OrderedDict((c, cvs.line(df, 't', c)) for c in str_cols)
    merged = xr.concat(aggs.values(), dim=pd.Index(str_cols, name='cols'))

    total = tf.interpolate(merged.sum(dim='cols'), cmap=reversed(viridis), how=how)
    pil_img = total.to_pil()
    pil_img.save(path)


def plot_stacked(ax, arr):
    arr = arr.T
    arr = sp.vstack((sp.zeros(arr.shape[1]),arr))
    cumsum = sp.cumsum(arr, axis=0)
    from itertools import cycle
    color = cycle(sns.get_color_cycle())
    for k, (lower, upper) in enumerate(zip(cumsum[:-1], cumsum[1:])):
        ax.fill_between(sp.arange(len(lower)), upper, lower, color=next(color), label=k)
    ax.legend(loc="center left", bbox_to_anchor=(1, .5))


def contour_polygon(X, Y, Z, level, **polygon_kwargs):
    from matplotlib import _cntr as cntr
    ct = cntr.Cntr(X, Y, Z)
    res = ct.trace(level)
    nseg = len(res) // 2
    segments, codes = res[:nseg], res[nseg:]
    p = plt.Polygon(segments[0], **polygon_kwargs)
    return p