import copy
import pkg_resources
import warnings

import dclab
from dclab import kde_contours
import numpy as np
from PyQt5 import uic, QtCore, QtGui, QtWidgets
import pyqtgraph as pg
from pyqtgraph import exporters
from pyqtgraph.graphicsItems.GradientEditorItem import Gradients


from .. import plot_cache
from .. import util

from .widgets import ColorBarWidget, SimplePlotItem


# Register custom colormaps
Gradients["grayblue"] = {'ticks': [(0.0, (100, 100, 100, 255)),
                                   (1.0, (0, 0, 255, 255))],
                         'mode': 'rgb'}

Gradients["graygreen"] = {'ticks': [(0.0, (100, 100, 100, 255)),
                                    (1.0, (0, 180, 0, 255))],
                          'mode': 'rgb'}

Gradients["grayorange"] = {'ticks': [(0.0, (100, 100, 100, 255)),
                                     (1.0, (210, 110, 0, 255))],
                           'mode': 'rgb'}

Gradients["grayred"] = {'ticks': [(0.0, (100, 100, 100, 255)),
                                  (1.0, (200, 0, 0, 255))],
                        'mode': 'rgb'}


class ContourSpacingTooLarge(UserWarning):
    pass


class PipelinePlot(QtWidgets.QWidget):
    """Implements the plotting pipeline using pyqtgraph"""
    instances = {}

    def __init__(self, parent, pipeline, plot_id, *args, **kwargs):
        QtWidgets.QWidget.__init__(self, parent=parent, *args, **kwargs)
        path_ui = pkg_resources.resource_filename(
            "shapeout2.gui", "pipeline_plot.ui")
        uic.loadUi(path_ui, self)
        # used to avoid unnecessary plotting
        self._plot_data_hash = "unset"

        self.pipeline = pipeline
        self.identifier = plot_id
        self.update_content()
        PipelinePlot.instances[plot_id] = self

    def update_content(self):
        """Update the current plot"""
        parent = self.parent()
        dslist, slot_states = self.pipeline.get_plot_datasets(self.identifier)
        plot = self.pipeline.get_plot(self.identifier)
        plot_state = plot.__getstate__()
        # check whether anything changed
        # 1. plot state and all relevant slot states
        tohash = [slot_states, plot_state]
        # 2. all relevant filter states
        for slot_state in slot_states:
            slot_id = slot_state["identifier"]
            for filt_id in self.pipeline.filter_ids:
                if self.pipeline.is_element_active(slot_id, filt_id):
                    filt = self.pipeline.get_filter(filt_id)
                    filt_state = filt.__getstate__()
                    tohash.append([slot_id, filt_id, filt_state])
                    # also check whether the polygon filters changed (#26)
                    for pid in filt_state["polygon filters"]:
                        pf = dclab.PolygonFilter.get_instance_from_id(pid)
                        tohash.append(pf.__getstate__())
        plot_data_hash = util.hashobj(tohash)
        if plot_data_hash == self._plot_data_hash:
            # do nothing
            return
        else:
            self._plot_data_hash = plot_data_hash

        # abbreviations
        gen = plot_state["general"]
        lay = plot_state["layout"]
        sca = plot_state["scatter"]

        # auto range (overrides stored ranges)
        if gen["auto range"]:
            # default range is limits + 5% margin
            gen["range x"] = self.pipeline.get_min_max(feat=gen["axis x"],
                                                       plot_id=self.identifier,
                                                       margin=.05)
            gen["range y"] = self.pipeline.get_min_max(feat=gen["axis y"],
                                                       plot_id=self.identifier,
                                                       margin=0.05)

        # title
        self.setWindowTitle(lay["name"])

        # clear widget
        self.plot_layout.clear()

        if not slot_states:
            return

        labelx, labely = get_axes_labels(plot_state, slot_states)

        # font size for plot title (default size + 2)
        size = "{}pt".format(QtGui.QFont().pointSize() + 2)
        self.plot_layout.addLabel(lay["name"], colspan=3, size=size)
        self.plot_layout.nextRow()

        self.plot_layout.addLabel(labely, angle=-90)
        linner = self.plot_layout.addLayout()

        # limits in case of scatter plot and feature hue
        if lay["division"] == "merge":
            pp = PipelinePlotItem(parent=linner)
            linner.addItem(item=pp,
                           row=None,
                           col=None,
                           rowspan=1,
                           colspan=1)
            pp.redraw(dslist, slot_states, plot_state)
        elif lay["division"] == "each":
            colcount = 0
            for ds, sl in zip(dslist, slot_states):
                pp = PipelinePlotItem(parent=linner)
                linner.addItem(item=pp,
                               row=None,
                               col=None,
                               rowspan=1,
                               colspan=1)
                pp.redraw([ds], [sl], plot_state)
                colcount += 1
                if colcount % lay["column count"] == 0:
                    linner.nextRow()
        elif lay["division"] == "multiscatter+contour":
            colcount = 0
            # scatter plots
            plot_state_scatter = copy.deepcopy(plot_state)
            plot_state_scatter["contour"]["enabled"] = False
            for ds, sl in zip(dslist, slot_states):
                pp = PipelinePlotItem(parent=linner)
                linner.addItem(item=pp,
                               row=None,
                               col=None,
                               rowspan=1,
                               colspan=1)
                pp.redraw([ds], [sl], plot_state_scatter)
                colcount += 1
                if colcount % lay["column count"] == 0:
                    linner.nextRow()
            # contour plot
            plot_state_contour = copy.deepcopy(plot_state)
            plot_state_contour["scatter"]["enabled"] = False
            pp = PipelinePlotItem(parent=linner)
            linner.addItem(item=pp,
                           row=None,
                           col=None,
                           rowspan=1,
                           colspan=1)
            pp.redraw(dslist, slot_states, plot_state_contour)

        # colorbar
        colorbar_kwds = {}

        if sca["marker hue"] == "kde":
            colorbar_kwds["vmin"] = 0
            colorbar_kwds["vmax"] = 1
            colorbar_kwds["label"] = "density [a.u.]"
        elif sca["marker hue"] == "feature":
            feat = sca["hue feature"]
            label = dclab.dfn.get_feature_label(feat)
            fl_names = slot_states[0]["fl names"]
            if label.count("FL"):
                for key in fl_names:
                    if key in label:
                        label = label.replace(key, fl_names[key])
                        break
            colorbar_kwds["label"] = label
            if label.endswith("[a.u.]"):
                colorbar_kwds["vmin"] = 0
                colorbar_kwds["vmax"] = 1
            else:
                colorbar_kwds["vmin"] = sca["hue min"]
                colorbar_kwds["vmax"] = sca["hue max"]

        if colorbar_kwds:
            # add colorbar
            colorbar = ColorBarWidget(
                cmap=sca["colormap"],
                width=15,
                height=min(300, lay["size y"]//2),
                **colorbar_kwds
            )
            self.plot_layout.addItem(colorbar)

        # x-axis label
        self.plot_layout.nextRow()
        self.plot_layout.addLabel(labelx, col=1)

        # Set size in the end (after layout is populated)
        self.setMinimumSize(lay["size x"], lay["size y"])
        self.setMaximumSize(lay["size x"], lay["size y"])
        size_hint = self.parent().sizeHint()
        parent.setMinimumSize(size_hint)
        parent.setMaximumSize(size_hint)


class PipelinePlotItem(SimplePlotItem):
    def __init__(self, *args, **kwargs):
        super(PipelinePlotItem, self).__init__(*args, **kwargs)
        # circumvent problems with removed plots
        self.setAcceptHoverEvents(False)
        # Disable user interaction
        self.setMouseEnabled(x=False, y=False)
        # bring axes to front
        self.axes_to_front()
        # Keep track of all elements (for redraw)
        self._plot_elements = []

    def perform_export(self, file):
        """Performs export in new layout with axes labels set

        Overrides the basic functionality of SimplePlotItem.
        See https://github.com/ZELLMECHANIK-DRESDEN/ShapeOut2/issues/7
        """
        # Create a plot window
        win = pg.GraphicsLayoutWidget(
            size=(self.width() + 100, self.height() + 100),
            show=True)
        # fill layout
        labelx, labely = get_axes_labels(self.plot_state, self.slot_states)
        win.addLabel(labely, angle=-90)
        explot = PipelinePlotItem()
        explot.redraw(self.dslist, self.slot_states, self.plot_state)
        win.addItem(explot)
        win.addLabel("")  # spacer to avoid cut tick labels on the right(#7)
        win.nextRow()
        win.addLabel(labelx, col=1)
        # Update the UI (do it twice, otherwise the tick labels overlap)
        QtWidgets.QApplication.processEvents()
        QtWidgets.QApplication.processEvents()
        win.hide()
        # perform actual export
        suffix = file[-3:]
        if suffix == "png":
            exp = exporters.ImageExporter(win.scene())
            # translate from screen resolution (80dpi) to 300dpi
            exp.params["width"] = int(exp.params["width"] / 72 * 300)
        elif suffix == "svg":
            exp = exporters.SVGExporter(win.scene())
        exp.export(file)

    def redraw(self, dslist, slot_states, plot_state):
        # Remove everything
        for el in self._plot_elements:
            self.removeItem(el)

        if not dslist:
            return

        self.dslist = dslist
        self.slot_states = slot_states
        self.plot_state = plot_state

        # General
        gen = plot_state["general"]
        # TODO:
        # - test whether all datasets have same channel width / pixel size
        # Isoelastics
        if gen["isoelastics"]:
            cfg = dslist[0].config
            els = add_isoelastics(plot_item=self,
                                  axis_x=gen["axis x"],
                                  axis_y=gen["axis y"],
                                  channel_width=cfg["setup"]["channel width"],
                                  pixel_size=cfg["imaging"]["pixel size"])
            self._plot_elements += els
        # Modifications in log mode
        set_viewbox(self,
                    range_x=gen["range x"],
                    range_y=gen["range y"],
                    scale_x=gen["scale x"],
                    scale_y=gen["scale y"])
        # Scatter data
        sca = plot_state["scatter"]
        if sca["enabled"]:
            for rtdc_ds, ss in zip(dslist, slot_states):
                sct = add_scatter(plot_item=self,
                                  rtdc_ds=rtdc_ds,
                                  plot_state=plot_state,
                                  slot_state=ss
                                  )
                self._plot_elements += sct
        # Contour data
        if plot_state["contour"]["enabled"]:
            # show legend
            if plot_state["contour"]["legend"]:
                legend = self.addLegend(offset=(-.01, +.01))
            else:
                legend = None
            for rtdc_ds, ss in zip(dslist, slot_states):
                con = add_contour(plot_item=self,
                                  rtdc_ds=rtdc_ds,
                                  plot_state=plot_state,
                                  slot_state=ss,
                                  legend=legend,
                                  )
                self._plot_elements += con

        # Set subplot title and number of events
        if plot_state["layout"]["label plots"]:
            if len(dslist) == 1 and plot_state["scatter"]["enabled"]:
                # only one scatter plot
                # set title
                ss = slot_states[0]
                thtmls = "<span style='color:{}'>{}</span>"
                title = thtmls.format(ss["color"], ss["name"])
                self.setTitle(title)
                if plot_state["scatter"]["show event count"]:
                    # set event count
                    chtml = "<span style='font-size:{}pt'>".format(
                        # default font size - 1
                        QtGui.QFont().pointSize() - 1) + "{} events</span>"
                    label = QtWidgets.QGraphicsTextItem(
                        "",
                        # This is kind of hackish: set the parent to the right
                        # axis so that it is always drawn there.
                        parent=self.axes["right"]["item"])
                    label.setHtml(chtml.format(len(sct[0].data)))
                    # move the label to the left by its width
                    label.setPos(-label.boundingRect().width()+2, -5)

            elif (plot_state["contour"]["enabled"]
                    and not plot_state["scatter"]["enabled"]):
                # only a contour plot
                # set title
                self.setTitle("Contours")


def add_contour(plot_item, plot_state, rtdc_ds, slot_state, legend=None):
    gen = plot_state["general"]
    con = plot_state["contour"]
    try:
        x, y, density = plot_cache.get_contour_data(
            rtdc_ds=rtdc_ds,
            xax=gen["axis x"],
            yax=gen["axis y"],
            xacc=con["spacing x"],
            yacc=con["spacing y"],
            xscale=gen["scale x"],
            yscale=gen["scale y"],
            kde_type=gen["kde"],
        )
    except ValueError:
        # most-likely there is nothing to compute a contour for
        return []
    if density.shape[0] < 3 or density.shape[1] < 3:
        warnings.warn("Contour not possible; spacing may be too large!",
                      ContourSpacingTooLarge)
        return []
    plev = kde_contours.get_quantile_levels(
        density=density,
        x=x,
        y=y,
        xp=rtdc_ds[gen["axis x"]][rtdc_ds.filter.all],
        yp=rtdc_ds[gen["axis y"]][rtdc_ds.filter.all],
        q=np.array(con["percentiles"])/100,
        normalize=True)
    contours = []
    for level in plev:
        # make sure that the contour levels are not at the boundaries
        if not (np.allclose(level, 0, atol=1e-12, rtol=0)
                or np.allclose(level, 1, atol=1e-12, rtol=0)):
            cc = kde_contours.find_contours_level(
                density, x=x, y=y, level=level)
            contours.append(cc)

    elements = []
    for ii in range(len(contours)):
        style = linestyles[con["line styles"][ii]]
        width = con["line widths"][ii]
        for cci in contours[ii]:
            cline = pg.PlotDataItem(x=cci[:, 0],
                                    y=cci[:, 1],
                                    pen=pg.mkPen(color=slot_state["color"],
                                                 width=width,
                                                 style=style,
                                                 ),
                                    )
            elements.append(cline)
            plot_item.addItem(cline)
            if ii == 0 and legend is not None:
                legend.addItem(cline, slot_state["name"])
            # Always plot higher percentiles above lower percentiles
            # (useful if there are multiple contour plots overlapping)
            cline.setZValue(con["percentiles"][ii])
    return elements


def add_isoelastics(plot_item, axis_x, axis_y, channel_width, pixel_size):
    elements = []
    isodef = dclab.isoelastics.get_default()
    # We do not use isodef.get_with_rtdcbase, because then the
    # isoelastics would be shifted according to flow rate and.
    # viscosity. We could do it, but for visualization there is
    # really no need and also, the plots then look the same as
    # in Shape-Out 1.
    try:
        iso = isodef.get(
            method="numerical",
            channel_width=channel_width,
            flow_rate=None,
            viscosity=None,
            col1=axis_x,
            col2=axis_y,
            add_px_err=True,
            px_um=pixel_size)
    except KeyError:
        pass
    else:
        for ss in iso:
            iline = pg.PlotDataItem(x=ss[:, 0], y=ss[:, 1])
            plot_item.addItem(iline)
            elements.append(iline)
            # send them to the back
            iline.setZValue(-100)
    return elements


def add_scatter(plot_item, plot_state, rtdc_ds, slot_state):
    gen = plot_state["general"]
    sca = plot_state["scatter"]
    scatter = pg.ScatterPlotItem(size=sca["marker size"],
                                 pen=pg.mkPen(color=(0, 0, 0, 0)),
                                 brush=pg.mkBrush("k"),
                                 symbol="s")
    scatter.setAcceptHoverEvents(False)
    plot_item.addItem(scatter)

    if sca["marker hue"] == "kde":
        kde_type = gen["kde"]
    else:
        kde_type = "none"

    x, y, kde, idx = plot_cache.get_scatter_data(
        rtdc_ds=rtdc_ds,
        downsample=sca["downsample"] * sca["downsampling value"],
        xax=gen["axis x"],
        yax=gen["axis y"],
        xscale=gen["scale x"],
        yscale=gen["scale y"],
        kde_type=kde_type,
    )
    # define colormap
    # TODO:
    # - common code base with QuickView
    cmap = pg.ColorMap(*zip(*Gradients[sca["colormap"]]["ticks"]))
    if sca["marker hue"] == "kde":
        brush = []
        # Note: we don't expand the density to [0, 1], because the
        # colorbar will show "density" and because we don want to
        # compute the density in this function and not someplace else.
        for k in kde:
            brush.append(cmap.mapToQColor(k))
        # Note, colors could also be digitized (does not seem to be faster):
        # cbin = np.linspace(0, 1, 1000)
        # dig = np.digitize(kde, cbin)
        # for idx in dig:
        #     brush.append(cmap.mapToQColor(cbin[idx]))
    elif sca["marker hue"] == "feature":
        brush = []
        feat = np.asarray(rtdc_ds[sca["hue feature"]][idx], dtype=float)
        feat -= sca["hue min"]
        feat /= sca["hue max"] - sca["hue min"]
        for f in feat:
            if np.isnan(f):
                brush.append(pg.mkColor("#FF0000"))
            else:
                brush.append(cmap.mapToQColor(f))
    elif sca["marker hue"] == "dataset":
        alpha = int(sca["marker alpha"] * 255)
        colord = pg.mkColor(slot_state["color"])
        colord.setAlpha(alpha)
        brush = pg.mkBrush(colord)
    else:
        alpha = int(sca["marker alpha"] * 255)
        colork = pg.mkColor("#000000")
        colork.setAlpha(alpha)
        brush = pg.mkBrush(colork)

    # convert to log-scale if applicable
    if gen["scale x"] == "log":
        x = np.log10(x)
    if gen["scale y"] == "log":
        y = np.log10(y)

    scatter.setData(x=x, y=y, brush=brush)
    scatter.setZValue(-50)
    return [scatter]


def get_axes_labels(plot_state, slot_states):
    gen = plot_state["general"]
    labelx = dclab.dfn.get_feature_label(gen["axis x"])
    labely = dclab.dfn.get_feature_label(gen["axis y"])
    # replace FL-? with user-defined names
    fl_names = slot_states[0]["fl names"]
    if labelx.count("FL"):
        for key in fl_names:
            if key in labelx:
                labelx = labelx.replace(key, fl_names[key])
                break
    if labely.count("FL"):
        for key in fl_names:
            if key in labely:
                labely = labely.replace(key, fl_names[key])
                break
    return labelx, labely


def set_viewbox(plot, range_x, range_y, scale_x="linear", scale_y="linear",
                padding=0):
    # Set Log scale
    plot.setLogMode(x=scale_x == "log",
                    y=scale_y == "log")
    range_x = np.array(range_x)
    range_y = np.array(range_y)
    if scale_x == "log":
        if range_x[0] <= 0:
            if range_x[1] > 10:
                range_x[0] = 1e-1
            else:
                range_x[0] = 1e-3
        range_x = np.log10(range_x)
    if scale_y == "log":
        if range_y[0] <= 0:
            if range_y[1] > 10:
                range_y[0] = 1e-1
            else:
                range_y[0] = 1e-3
        range_y = np.log10(range_y)
    # Set Range
    plot.setRange(xRange=range_x,
                  yRange=range_y,
                  padding=padding,
                  )


linestyles = {
    "solid": QtCore.Qt.SolidLine,
    "dashed": QtCore.Qt.DashLine,
    "dotted": QtCore.Qt.DotLine,
}
