""" """
# TODO: marker alpha

import functools
import itertools

import numpy as np

from ..manage import auto_xyz_ds
from .core import (
    PLOTTER_DEFAULTS,
    AbstractHeatMap,
    AbstractLinePlot,
    AbstractScatter,
    Plotter,
    calc_row_col_datasets,
    intercept_call_arg,
    prettify,
)


@functools.lru_cache(1)
def _init_bokeh_nb():
    """Cache this so it doesn't happen over and over again."""
    from bokeh.plotting import output_notebook
    from bokeh.resources import INLINE

    output_notebook(resources=INLINE)


def bshow(figs, nb=True, interactive=False, **kwargs):
    """ """
    from bokeh.plotting import show

    if nb:
        _init_bokeh_nb()
        show(figs, notebook_handle=interactive)
    else:
        show(figs)


# --------------------------------------------------------------------------- #
#                     Main lineplot interface for bokeh                       #
# --------------------------------------------------------------------------- #


class PlotterBokeh(Plotter):
    def __init__(self, ds, x, y, z=None, **kwargs):
        """ """
        # bokeh custom options / defaults
        kwargs["return_fig"] = kwargs.pop("return_fig", False)
        self._interactive = kwargs.pop("interactive", False)

        super().__init__(ds, x, y, z, **kwargs, backend="BOKEH")

    def prepare_axes(self):
        """Make the bokeh plot figure and set options."""
        from bokeh.plotting import figure

        if self.add_to_axes is not None:
            self._plot = self.add_to_axes

        else:
            # Currently axes scale type must be set at figure creation?
            self._plot = figure(
                # convert figsize to roughly matplotlib dimensions
                width=int(
                    self.figsize[0] * 80
                    + (100 if self._use_legend else 0)
                    + (20 if self._ytitle else 0)
                    + (20 if not self.yticklabels_hide else 0)
                ),
                height=int(
                    self.figsize[1] * 80
                    + (20 if self.title else 0)
                    + (20 if self._xtitle else 0)
                    + (20 if not self.xticklabels_hide else 0)
                ),
                x_axis_type=("log" if self.xlog else "linear"),
                y_axis_type=("log" if self.ylog else "linear"),
                y_axis_location=("right" if self.ytitle_right else "left"),
                title=self.title,
                toolbar_location="above",
                toolbar_sticky=False,
                active_scroll="wheel_zoom",
            )

    def set_axes_labels(self):
        """Set the labels on the axes."""
        if self._xtitle:
            self._plot.xaxis.axis_label = self._xtitle
        if self._ytitle:
            self._plot.yaxis.axis_label = self._ytitle

    def set_axes_range(self):
        """Set the plot ranges of the axes, and the panning limits."""
        from bokeh.models import DataRange1d

        self.calc_data_range()

        # plt_x_centre = (self._data_xmax + self._data_xmin) / 2
        # plt_x_range = self._data_xmax - self._data_xmin
        # xbounds = (plt_x_centre - plt_x_range, plt_x_centre + plt_x_range)
        xbounds = None
        self._plot.x_range = (
            DataRange1d(
                start=self._xlims[0], end=self._xlims[1], bounds=xbounds
            )
            if self._xlims
            else DataRange1d(bounds=xbounds)
        )

        # plt_y_centre = (self._data_ymax + self._data_ymin) / 2
        # plt_y_range = abs(self._data_ymax - self._data_ymin)
        # ybounds = (plt_y_centre - plt_y_range, plt_y_centre + plt_y_range)
        ybounds = None
        self._plot.y_range = (
            DataRange1d(
                start=self._ylims[0], end=self._ylims[1], bounds=ybounds
            )
            if self._ylims
            else DataRange1d(bounds=ybounds)
        )

    def set_spans(self):
        """Set custom horizontal and verical line spans."""
        from bokeh.models import Span

        span_opts = {
            "level": "glyph",
            "line_dash": "dashed",
            "line_color": (127, 127, 127),
            "line_width": self.span_width,
        }

        if self.hlines:
            for hl in self.hlines:
                self._plot.add_layout(
                    Span(location=hl, dimension="width", **span_opts)
                )
        if self.vlines:
            for vl in self.vlines:
                self._plot.add_layout(
                    Span(location=vl, dimension="height", **span_opts)
                )

    def set_gridlines(self):
        """Set whether to use gridlines or not."""
        if not self.gridlines:
            self._plot.xgrid.visible = False
            self._plot.ygrid.visible = False
        else:
            self._plot.xgrid.grid_line_dash = self.gridline_style
            self._plot.ygrid.grid_line_dash = self.gridline_style

    def set_tick_marks(self):
        """Set custom locations for the tick marks."""
        from bokeh.models import FixedTicker

        if self.xticks:
            self._plot.xaxis[0].ticker = FixedTicker(ticks=self.xticks)
        if self.yticks:
            self._plot.yaxis[0].ticker = FixedTicker(ticks=self.yticks)

        if self.xticklabels_hide:
            self._plot.xaxis.major_label_text_font_size = "0pt"
        if self.yticklabels_hide:
            self._plot.yaxis.major_label_text_font_size = "0pt"

    def set_sources_heatmap(self):
        from bokeh.plotting import ColumnDataSource

        # initialize empty source
        if not hasattr(self, "_source"):
            self._source = ColumnDataSource(data=dict())

        # remove mask from data -> not necessary soon? / convert to nan?
        var = np.ma.getdata(self._heatmap_var)

        self._source.add([var], "image")
        self._source.add([self._data_xmin], "x")
        self._source.add([self._data_ymin], "y")
        self._source.add([self._data_xmax - self._data_xmin], "dw")
        self._source.add([self._data_ymax - self._data_ymin], "dh")

    def set_sources(self):
        """Set the source dictionaries to be used by the plotter functions.
        This is seperate to allow interactive updates of the data only.
        """
        from bokeh.plotting import ColumnDataSource

        # check if heatmap
        if hasattr(self, "_heatmap_var"):
            return self.set_sources_heatmap()

        # 'copy' the zlabels iterator into src_zlbs
        self._zlbls, src_zlbs = itertools.tee(self._zlbls)

        # Initialise with empty dicts
        if not hasattr(self, "_sources"):
            self._sources = [
                ColumnDataSource(dict()) for _ in range(len(self._z_vals))
            ]

        # range through all data and update the sources
        for i, (zlabel, data) in enumerate(zip(src_zlbs, self._gen_xy())):
            self._sources[i].add(data["x"], "x")
            self._sources[i].add(data["y"], "y")
            self._sources[i].add([zlabel] * len(data["x"]), "z_coo")

            # check for color for scatter plot
            if "c" in data:
                self._sources[i].add(data["c"], "c")

            # check if should set y_err as well
            if "ye" in data:
                y_err_p = data["y"] + data["ye"]
                y_err_m = data["y"] - data["ye"]
                self._sources[i].add(
                    list(zip(data["x"], data["x"])), "y_err_xs"
                )
                self._sources[i].add(list(zip(y_err_p, y_err_m)), "y_err_ys")

            # check if should set x_err as well
            if "xe" in data:
                x_err_p = data["x"] + data["xe"]
                x_err_m = data["x"] - data["xe"]
                self._sources[i].add(
                    list(zip(data["y"], data["y"])), "x_err_ys"
                )
                self._sources[i].add(list(zip(x_err_p, x_err_m)), "x_err_xs")

    def plot_legend(self, legend_items=None):
        """Add a legend to the plot."""
        if self._use_legend:
            from bokeh.models import Legend

            loc = {"best": "top_left"}.get(self.legend_loc, self.legend_loc)
            where = {None: "right"}.get(self.legend_where, self.legend_where)

            # might be manually specified, e.g. from multiplot
            if legend_items is None:
                legend_items = self._lgnd_items

            lg = Legend(items=legend_items)
            lg.location = loc
            lg.click_policy = "hide"
            self._plot.add_layout(lg, where)

            # Don't repeatedly redraw legend
            self._use_legend = False

    def set_mappable(self):
        import matplotlib as plt
        from bokeh.models import LinearColorMapper, LogColorMapper

        mappr_fn = LogColorMapper if self.colormap_log else LinearColorMapper
        bokehpalette = [plt.colors.rgb2hex(m) for m in self.cmap(range(256))]

        self.mappable = mappr_fn(
            palette=bokehpalette, low=self._zmin, high=self._zmax
        )

    def plot_colorbar(self):
        if self._use_colorbar:
            where = {None: "right"}.get(self.legend_where, self.legend_where)

            from bokeh.models import BasicTicker, ColorBar, LogTicker

            ticker = LogTicker if self.colormap_log else BasicTicker
            color_bar = ColorBar(
                color_mapper=self.mappable,
                location=(0, 0),
                ticker=ticker(desired_num_ticks=6),
                title=self._ctitle,
            )
            self._plot.add_layout(color_bar, where)

    def set_tools(self):
        """Set which tools appear for the plot."""
        from bokeh.models import HoverTool

        tooltips = [
            (
                "({}, {})".format(
                    self.x_coo,
                    self.y_coo if isinstance(self.y_coo, str) else None,
                ),
                "(@x, @y)",
            ),
        ]
        if self.z_coo:
            tooltips.append((self.z_coo, "@z_coo"))

        self._plot.add_tools(HoverTool(tooltips=tooltips))

    def update(self):
        from bokeh.io import push_notebook

        self.set_sources()
        push_notebook()

    def show(self, **kwargs):
        """Show the produced figure."""
        if self.return_fig:
            return self._plot
        bshow(self._plot, **kwargs)
        return self

    def prepare_plot(self):
        self.prepare_axes()
        self.set_axes_labels()
        self.set_axes_range()
        self.set_spans()
        self.set_gridlines()
        self.set_tick_marks()
        self.set_sources()


def bokeh_multi_plot(fn):
    """Decorate a plotting function to plot a grid of values."""

    @functools.wraps(fn)
    def multi_plotter(ds, *args, row=None, col=None, link=False, **kwargs):

        if (row is None) and (col is None):
            return fn(ds, *args, **kwargs)

        # Set some global parameters
        p = fn(ds, *args, **kwargs, call=False)
        p.prepare_data_multi_grid()

        kwargs["xlims"] = kwargs.get("xlims", [p._data_xmin, p._data_xmax])
        kwargs["ylims"] = kwargs.get("ylims", [p._data_ymin, p._data_ymax])

        kwargs["vmin"] = kwargs.pop("vmin", p.vmin)
        kwargs["vmax"] = kwargs.pop("vmax", p.vmax)

        # split the dataset into its respective rows and columns
        ds_r_c, nrows, ncols = calc_row_col_datasets(ds, row=row, col=col)

        # intercept figsize as meaning *total* size for whole grid
        figsize = kwargs.pop("figsize", None)
        if figsize is None:
            av_n = (ncols + nrows) / 2
            figsize = (2 * (4 / av_n) ** 0.5, 2 * (4 / av_n) ** 0.5)
        else:
            figsize = (figsize[0] / ncols, figsize[1] / nrows)
        kwargs["figsize"] = figsize

        # intercept return_fig for the full grid and other options
        return_fig = kwargs.pop("return_fig", False)

        subplots = {}

        # range through rows and do subplots
        for i, ds_r in enumerate(ds_r_c):
            # range through columns
            for j, sub_ds in enumerate(ds_r):
                skws = {"legend": False, "colorbar": False}

                # if not last row
                if i != nrows - 1:
                    skws["xticklabels_hide"] = True
                    skws["xtitle"] = ""

                # if not first column
                if j != 0:
                    skws["yticklabels_hide"] = True
                    skws["ytitle"] = ""

                # label each column
                if (i == 0) and (col is not None):
                    col_val = prettify(ds[col].values[j])
                    skws["title"] = "{} = {}".format(col, col_val)
                    fx = "fontsize_xtitle"
                    skws["fontsize_title"] = kwargs.get(
                        fx, PLOTTER_DEFAULTS[fx]
                    )

                # label each row
                if (j == ncols - 1) and (row is not None):
                    skws["ytitle_right"] = True
                    row_val = prettify(ds[row].values[i])
                    skws["ytitle"] = "{} = {}".format(row, row_val)

                subplots[i, j] = fn(
                    sub_ds,
                    *args,
                    return_fig=True,
                    call=False,
                    **{**kwargs, **skws},
                )

        from bokeh.layouts import gridplot

        plts = [[subplots[i, j]() for j in range(ncols)] for i in range(nrows)]

        # link zooming and panning between all plots
        if link:
            x_range, y_range = plts[0][0].x_range, plts[0][0].y_range
            for i in range(nrows):
                for j in range(ncols):
                    plts[i][j].x_range = x_range
                    plts[i][j].y_range = y_range

        # the main grid
        p._plot = gridplot(plts)

        if p._use_legend or p._use_colorbar:
            from bokeh.layouts import row
            from bokeh.models import ColorBar, GlyphRenderer, Legend, Range1d

            # plot dummy using last sub_ds
            skws = {
                "title": "",
                "legend_loc": "center_left",
                "legend_where": "left",
            }
            lgren = fn(sub_ds, *args, return_fig=True, **{**kwargs, **skws})

            # remove all but legend, colorbar and glyph renderers
            lgren.renderers = [
                r
                for r in lgren.renderers
                if isinstance(r, (Legend, GlyphRenderer, ColorBar))
            ]
            lgren.toolbar_location = None
            lgren.outline_line_color = None

            # size it - this is pretty hacky at the moment
            lgren.width = 120
            lgren.height = int(80 * figsize[1] * nrows + 100)
            lgren.x_range = Range1d(0, 0)
            lgren.y_range = Range1d(0, 0)

            # append to the right of the gridplot
            p._plot = row([p._plot, lgren])

        if return_fig:
            return p._plot
        bshow(p._plot)

    return multi_plotter


class ILinePlot(PlotterBokeh, AbstractLinePlot):
    def __init__(self, ds, x, y, z=None, y_err=None, x_err=None, **kwargs):
        super().__init__(ds, x, y, z=z, y_err=y_err, x_err=x_err, **kwargs)

    def plot_lines(self):
        """Plot the data and a corresponding legend."""
        self._lgnd_items = []

        for src in self._sources:
            col = next(self._cols)
            zlabel = next(self._zlbls)
            legend_pics = []

            if self.lines:
                line = self._plot.line(
                    "x",
                    "y",
                    source=src,
                    color=col,
                    line_dash=next(self._lines),
                    line_width=next(self._lws) * 1.5,
                )
                legend_pics.append(line)

            if self.markers:
                marker = next(self._mrkrs)
                m = getattr(self._plot, marker)(
                    "x",
                    "y",
                    source=src,
                    name=zlabel,
                    color=col,
                    fill_alpha=0.5,
                    line_width=0.5,
                    size=self._marker_size,
                )
                legend_pics.append(m)

            # Check if errors specified as well
            if self.y_err:
                err = self._plot.multi_line(
                    xs="y_err_xs",
                    ys="y_err_ys",
                    source=src,
                    color=col,
                    line_width=self.errorbar_linewidth,
                )
                legend_pics.append(err)
            if self.x_err:
                err = self._plot.multi_line(
                    xs="x_err_xs",
                    ys="x_err_ys",
                    source=src,
                    color=col,
                    line_width=self.errorbar_linewidth,
                )
                legend_pics.append(err)

            # Add the names and styles of drawn lines for the legend
            self._lgnd_items.append((zlabel, legend_pics))

    def __call__(self):
        self.prepare_data_single()
        # Bokeh preparation
        self.prepare_plot()
        self.plot_lines()
        self.plot_legend()
        self.plot_colorbar()
        self.set_tools()
        return self.show(interactive=self._interactive)


@bokeh_multi_plot
@intercept_call_arg
def ilineplot(ds, x, y, z=None, y_err=None, x_err=None, **kwargs):
    """From  ``ds`` plot lines of ``y`` as a function of ``x``, optionally for
    varying ``z``. Interactive,

    Parameters
    ----------
    ds : xarray.Dataset
        Dataset to plot from.
    x : str
        Dimension to plot along the x-axis.
    y : str or tuple[str]
        Variable(s) to plot along the y-axis. If tuple, plot each of the
        variables - instead of ``z``.
    z : str, optional
        Dimension to plot into the page.
    y_err : str, optional
        Variable to plot as y-error.
    x_err : str, optional
        Variable to plot as x-error.
    row : str, optional
        Dimension to vary over as a function of rows.
    col : str, optional
        Dimension to vary over as a function of columns.
    plot_opts
        See ``xyzpy.plot.core.PLOTTER_DEFAULTS``.
    """
    return ILinePlot(ds, x, y, z, y_err=y_err, x_err=x_err, **kwargs)


class AutoILinePlot(ILinePlot):
    """Interactive raw data multi-line plot."""

    def __init__(self, x, y_z, **lineplot_opts):
        ds = auto_xyz_ds(x, y_z)
        super().__init__(ds, "x", "y", z="z", **lineplot_opts)


def auto_ilineplot(x, y_z, **lineplot_opts):
    """Auto version of :func:`~xyzpy.ilineplot` that accepts array arguments
    by converting them to a ``Dataset`` first.
    """
    return AutoILinePlot(x, y_z, **lineplot_opts)()


# --------------------------------------------------------------------------- #


class IScatter(PlotterBokeh, AbstractScatter):
    def __init__(self, ds, x, y, z=None, **kwargs):
        super().__init__(ds, x, y, z, **kwargs, markers=True)

    def plot_scatter(self):
        self._lgnd_items = []

        for src in self._sources:
            if "c" in src.column_names:
                col = {"field": "c", "transform": self.mappable}
            else:
                col = next(self._cols)
            marker = next(self._mrkrs)
            zlabel = next(self._zlbls)
            legend_pics = []

            m = getattr(self._plot, marker)(
                "x",
                "y",
                source=src,
                name=zlabel,
                color=col,
                fill_alpha=0.5,
                line_width=0.5,
                size=self._marker_size,
            )
            legend_pics.append(m)

            # Add the names and styles of drawn markers for the legend
            self._lgnd_items.append((zlabel, legend_pics))

    def __call__(self):
        self.prepare_data_single()
        # Bokeh preparation
        self.prepare_plot()
        self.plot_scatter()
        self.plot_legend()
        self.plot_colorbar()
        self.set_tools()
        return self.show(interactive=self._interactive)


@bokeh_multi_plot
@intercept_call_arg
def iscatter(ds, x, y, z=None, y_err=None, x_err=None, **kwargs):
    """From  ``ds`` plot a scatter of ``y`` against ``x``, optionally for
    varying ``z``. Interactive.

    Parameters
    ----------
    ds : xarray.Dataset
        Dataset to plot from.
    x : str
        Quantity to plot along the x-axis.
    y : str or tuple[str]
        Quantity(s) to plot along the y-axis. If tuple, plot each of the
        variables - instead of ``z``.
    z : str, optional
        Dimension to plot into the page.
    y_err : str, optional
        Variable to plot as y-error.
    x_err : str, optional
        Variable to plot as x-error.
    row : str, optional
        Dimension to vary over as a function of rows.
    col : str, optional
        Dimension to vary over as a function of columns.
    plot_opts
        See ``xyzpy.plot.core.PLOTTER_DEFAULTS``.
    """
    return IScatter(ds, x, y, z, y_err=y_err, x_err=x_err, **kwargs)


class AutoIScatter(IScatter):
    def __init__(self, x, y_z, **iscatter_opts):
        ds = auto_xyz_ds(x, y_z)
        super().__init__(ds, "x", "y", z="z", **iscatter_opts)


def auto_iscatter(x, y_z, **iscatter_opts):
    """Auto version of :func:`~xyzpy.iscatter` that accepts array arguments
    by converting them to a ``Dataset`` first.
    """
    return AutoIScatter(x, y_z, **iscatter_opts)()


# --------------------------------------------------------------------------- #


_HEATMAP_ALT_DEFAULTS = (
    ("legend", False),
    ("colorbar", True),
    ("colormap", "inferno"),
    ("gridlines", False),
    ("padding", 0),
    ("figsize", (5, 5)),  # try to be square, maybe use aspect_ratio??
)


class IHeatMap(PlotterBokeh, AbstractHeatMap):
    def __init__(self, ds, x, y, z, **kwargs):
        # set some heatmap specific options
        for k, default in _HEATMAP_ALT_DEFAULTS:
            if k not in kwargs:
                kwargs[k] = default
        super().__init__(ds, x, y, z, **kwargs)

    def plot_heatmap(self):
        self.calc_color_norm()
        self._plot.image(
            image="image",
            x="x",
            y="y",
            dw="dw",
            dh="dh",
            source=self._source,
            color_mapper=self.mappable,
        )

    def __call__(self):
        # Core preparation
        self.prepare_data_single()
        # matplotlib preparation
        self.prepare_plot()
        self.plot_heatmap()
        self.plot_colorbar()
        self.set_tools()
        return self.show(interactive=self._interactive)


@bokeh_multi_plot
@intercept_call_arg
def iheatmap(ds, x, y, z, **kwargs):
    """From  ``ds`` plot variable ``z`` as a function of ``x`` and ``y`` using
    a 2D heatmap. Interactive,

    Parameters
    ----------
    ds : xarray.Dataset
        Dataset to plot from.
    x : str
        Dimension to plot along the x-axis.
    y : str
        Dimension to plot along the y-axis.
    z : str, optional
        Variable to plot as colormap.
    row : str, optional
        Dimension to vary over as a function of rows.
    col : str, optional
        Dimension to vary over as a function of columns.
    plot_opts
        See ``xyzpy.plot.core.PLOTTER_DEFAULTS``.
    """
    return IHeatMap(ds, x, y, z, **kwargs)


class AutoIHeatMap(IHeatMap):
    def __init__(self, x, **iheatmap_opts):
        ds = auto_xyz_ds(x)
        super().__init__(ds, "y", "z", "x", **iheatmap_opts)


def auto_iheatmap(x, **iheatmap_opts):
    """Auto version of :func:`~xyzpy.iheatmap` that accepts array arguments
    by converting them to a ``Dataset`` first.
    """
    return AutoIHeatMap(x, **iheatmap_opts)()
