from __future__ import annotations
import logging
from typing import List, Optional
import panel as pn
import holoviews as hv
from param import Parameter
from functools import partial
import hvplot.xarray  # noqa pylint: disable=duplicate-code,unused-import
import xarray as xr

from bencher.utils import hmap_canonical_input, get_nearest_coords, get_nearest_coords1D
from bencher.results.panel_result import PanelResult
from bencher.results.bench_result_base import ReduceType

from bencher.plotting.plot_filter import PlotFilter, VarRange
from bencher.variables.results import ResultVar


hv.extension("bokeh", "plotly")


class HoloviewResult(PanelResult):
    @staticmethod
    def set_default_opts(width=600, height=600):
        width_heigh = {"width": width, "height": height, "tools": ["hover"]}
        hv.opts.defaults(
            hv.opts.Curve(**width_heigh),
            hv.opts.Points(**width_heigh),
            hv.opts.Bars(**width_heigh),
            hv.opts.Scatter(**width_heigh),
            hv.opts.HeatMap(cmap="plasma", **width_heigh, colorbar=True),
            # hv.opts.Surface(**width_heigh),
            hv.opts.GridSpace(plot_size=400),
        )
        return width_heigh

    def to(self, hv_type: hv.Chart, reduce: ReduceType = ReduceType.AUTO, **kwargs) -> hv.Chart:
        return self.to_hv_dataset(reduce).to(hv_type, **kwargs)

    def overlay_plots(self, plot_callback: callable) -> Optional[hv.Overlay]:
        results = []
        markdown_results = pn.Row()
        for rv in self.bench_cfg.result_vars:
            res = plot_callback(rv)
            if res is not None:
                if isinstance(res, pn.pane.Markdown):
                    markdown_results.append(res)
                else:
                    results.append(res)
        if len(results) > 0:
            overlay = hv.Overlay(results).collate()
            if len(markdown_results) == 0:
                return overlay
            return pn.Row(overlay, markdown_results)
        if len(markdown_results) > 0:
            return markdown_results
        return None

    def layout_plots(self, plot_callback: callable):
        if len(self.bench_cfg.result_vars) > 0:
            pt = hv.Layout()
            got_results = False
            for rv in self.bench_cfg.result_vars:
                res = plot_callback(rv)
                if res is not None:
                    got_results = True
                    pt += plot_callback(rv)
            return pt if got_results else None
        return plot_callback(self.bench_cfg.result_vars[0])

    def time_widget(self, title):
        return {"title": title}
        # if self.bench_cfg.over_time:
        #     time_widget_args = {"widget_type": "scrubber", "widget_location": "bottom"}
        #     time_widget_args["title"] = None  # use the title generated by the widget instead
        # else:
        #     time_widget_args = {"widget_type": "individual"}
        #     time_widget_args["title"] = title

        # return time_widget_args

    def to_bar(self, result_var: Parameter = None, **kwargs) -> Optional[pn.panel]:
        return self.filter(
            self.to_bar_ds,
            float_range=VarRange(0, 0),
            cat_range=VarRange(0, None),
            repeats_range=VarRange(1, 1),
            reduce=ReduceType.SQUEEZE,
            target_dimension=2,
            result_var=result_var,
            result_types=(ResultVar),
            **kwargs,
        )

    def to_bar_ds(self, dataset: xr.Dataset, result_var: Parameter = None, **kwargs):
        by = None
        if self.plt_cnt_cfg.cat_cnt >= 2:
            by = self.plt_cnt_cfg.cat_vars[1].name
        da_plot = dataset[result_var.name]
        title = self.title_from_ds(da_plot, result_var, **kwargs)
        time_widget_args = self.time_widget(title)
        return da_plot.hvplot.bar(by=by, **time_widget_args, **kwargs)

    def hv_container_ds(
        self, dataset: xr.Dataset, result_var: Parameter, container: hv.Chart = None, **kwargs
    ):
        return hv.Dataset(dataset[result_var.name]).to(container).opts(**kwargs)

    def to_hv_container(
        self,
        container: pn.pane.panel,
        reduce_type=ReduceType.AUTO,
        target_dimension: int = 2,
        result_var: Parameter = None,
        result_types=(ResultVar),
        **kwargs,
    ) -> Optional[pn.pane.panel]:
        return self.map_plot_panes(
            partial(self.hv_container_ds, container=container),
            hv_dataset=self.to_hv_dataset(reduce_type),
            target_dimension=target_dimension,
            result_var=result_var,
            result_types=result_types,
            **kwargs,
        )

    def to_line(self, result_var: Parameter = None, **kwargs) -> Optional[pn.panel]:
        return self.filter(
            self.to_line_ds,
            float_range=VarRange(1, 1),
            cat_range=VarRange(0, None),
            repeats_range=VarRange(1, 1),
            panel_range=VarRange(0, None),
            reduce=ReduceType.SQUEEZE,
            target_dimension=2,
            result_var=result_var,
            result_types=(ResultVar),
            **kwargs,
        )

    def to_line_ds(self, dataset: xr.Dataset, result_var: Parameter, **kwargs):
        x = self.plt_cnt_cfg.float_vars[0].name
        # y = self.plt_cnt_cfg.result_vars[0].name
        by = None
        if self.plt_cnt_cfg.cat_cnt >= 1:
            by = self.plt_cnt_cfg.cat_vars[0].name
        da_plot = dataset[result_var.name]
        title = self.title_from_ds(da_plot, result_var, **kwargs)
        time_widget_args = self.time_widget(title)
        return da_plot.hvplot.line(x=x, by=by, **time_widget_args, **kwargs)

    def to_curve(self, result_var: Parameter = None, **kwargs):
        return self.filter(
            self.to_curve_ds,
            float_range=VarRange(1, 1),
            cat_range=VarRange(0, None),
            repeats_range=VarRange(2, None),
            reduce=ReduceType.REDUCE,
            target_dimension=2,
            result_var=result_var,
            result_types=(ResultVar),
            **kwargs,
        )

    def to_curve_ds(
        self, dataset: xr.Dataset, result_var: Parameter, **kwargs
    ) -> Optional[hv.Curve]:
        hvds = hv.Dataset(dataset)
        # result_var = self.get_results_var_list(result_var)[0]
        title = self.title_from_ds(dataset, result_var, **kwargs)
        pt = hvds.to(hv.Curve).opts(title=title, **kwargs)
        pt *= hvds.to(hv.Spread).opts(alpha=0.2)
        if len(dataset.sizes) > 1:
            return pt.opts(legend_position="right").overlay()
        return pt.opts(legend_position="right")

    def to_heatmap(
        self,
        result_var: Parameter = None,
        tap_var=None,
        tap_container=None,
        target_dimension=2,
        **kwargs,
    ) -> Optional[pn.panel]:
        if tap_var is None:
            heatmap_cb = self.to_heatmap_ds
        else:
            heatmap_cb = partial(
                self.to_heatmap_container_tap_ds, result_var_plot=tap_var, container=tap_container
            )

        return self.filter(
            heatmap_cb,
            float_range=VarRange(2, None),
            cat_range=VarRange(0, None),
            input_range=VarRange(1, None),
            target_dimension=target_dimension,
            result_var=result_var,
            result_types=(ResultVar),
            **kwargs,
        )

    def to_heatmap_ds(
        self, dataset: xr.Dataset, result_var: Parameter, **kwargs
    ) -> Optional[hv.HeatMap]:
        if len(dataset.dims) >= 2:
            x = self.plt_cnt_cfg.float_vars[0].name
            y = self.plt_cnt_cfg.float_vars[1].name
            C = result_var.name
            title = f"Heatmap of {result_var.name}"
            time_args = self.time_widget(title)
            return dataset.hvplot.heatmap(x=x, y=y, C=C, cmap="plasma", **time_args, **kwargs)
        return None

    def to_heatmap_container_tap_ds(
        self,
        dataset: xr.Dataset,
        result_var: Parameter,
        result_var_plot: Parameter,
        container: pn.pane.panel = pn.pane.panel,
        **kwargs,
    ) -> pn.Row:
        htmap = self.to_heatmap_ds(dataset, result_var).opts(tools=["hover", "tap"], **kwargs)
        htmap_posxy = hv.streams.Tap(source=htmap, x=0, y=0)

        container_instance = container(**kwargs)
        title = pn.pane.Markdown("Selected: None")

        def tap_plot(x, y):  # pragma: no cover
            x_nearest = get_nearest_coords1D(
                x, dataset.coords[self.bench_cfg.input_vars[0].name].data
            )
            y_nearest = get_nearest_coords1D(
                y, dataset.coords[self.bench_cfg.input_vars[1].name].data
            )
            kdims = {}
            kdims[self.bench_cfg.input_vars[0].name] = x_nearest
            kdims[self.bench_cfg.input_vars[1].name] = y_nearest

            if hasattr(htmap, "current_key"):
                for d, k in zip(htmap.kdims, htmap.current_key):
                    kdims[d.name] = k

            ds = dataset[result_var_plot.name]
            val = ds.sel(**kdims)
            item = self.zero_dim_da_to_val(val)
            title.object = "Selected: " + ", ".join([f"{k}:{v}" for k, v in kdims.items()])
            container_instance.object = item
            if hasattr(container, "autoplay"):  # container is a video, set to autoplay
                container_instance.paused = False
                container_instance.time = 0
                container_instance.loop = True
                container_instance.autoplay = True

        htmap_posxy.add_subscriber(tap_plot)
        bound_plot = pn.Column(title, container_instance)
        return pn.Row(htmap, bound_plot)

    def to_error_bar(self) -> hv.Bars:
        return self.to_hv_dataset(ReduceType.REDUCE).to(hv.ErrorBars)

    def to_points(self, reduce: ReduceType = ReduceType.AUTO) -> hv.Points:
        ds = self.to_hv_dataset(reduce)
        pt = ds.to(hv.Points)
        if reduce:
            pt *= ds.to(hv.ErrorBars)
        return pt

    def to_scatter(self, **kwargs) -> Optional[pn.panel]:
        match_res = PlotFilter(
            float_range=VarRange(0, 0), cat_range=VarRange(0, None), repeats_range=VarRange(1, 1)
        ).matches_result(self.plt_cnt_cfg, "to_hvplot_scatter")
        if match_res.overall:
            hv_ds = self.to_hv_dataset(ReduceType.SQUEEZE)
            by = None
            subplots = False
            if self.plt_cnt_cfg.cat_cnt > 1:
                by = [v.name for v in self.bench_cfg.input_vars[1:]]
                subplots = False
            return hv_ds.data.hvplot.scatter(by=by, subplots=subplots, **kwargs).opts(
                title=self.to_plot_title()
            )
        return match_res.to_panel(**kwargs)

    # def to_scatter_jitter(self, **kwargs) -> Optional[hv.Scatter]:
    #     matches = PlotFilter(
    #         float_range=VarRange(0, 0),
    #         cat_range=VarRange(0, None),
    #         repeats_range=VarRange(2, None),
    #         input_range=VarRange(1, None),
    #     ).matches_result(self.plt_cnt_cfg, "to_scatter_jitter")
    #     if matches.overall:
    #         hv_ds = self.to_hv_dataset(ReduceType.NONE)

    #         by = None
    #         groupby = None
    #         subplots=False
    #         if self.plt_cnt_cfg.cat_cnt > 1:
    #             by = [v.name for v in self.bench_cfg.all_vars[1:]]
    #             subplots=False
    #         return hv_ds.data.hvplot.scatter(by=by,subplots=subplots, **kwargs).opts(title=self.to_plot_title())

    #         # pt = (
    #         #     hv_ds.to(hv.Scatter)
    #         #     .opts(jitter=0.1)
    #         #     .overlay("repeat")
    #         #     .opts(show_legend=False, title=self.to_plot_title(), **kwargs)
    #         # )
    #         # return pt
    #     return matches.to_panel()

    # def to_scatter_jitter_multi(self, **kwargs) -> Optional[hv.Scatter]:
    #     matches = PlotFilter(
    #         float_range=VarRange(0, 0),
    #         cat_range=VarRange(0, None),
    #         repeats_range=VarRange(2, None),
    #         input_range=VarRange(1, None),
    #     ).matches_result(self.plt_cnt_cfg, "to_scatter_jitter")
    #     if matches.overall:
    #         hv_dataset = self.to_hv_dataset(ReduceType.NONE)

    #         print("KDIMS",hv_dataset.kdims)
    #         # hv_dataset.kdims =[hv_dataset.kdims[2],hv_dataset.kdims[1],hv_dataset.kdims[0]]
    #         # print("KDIMS",hv_dataset.kdims)

    #         # exit()
    #         cb = partial(self.to_scatter_jitter_da, **kwargs)
    #         return self.to_panes_multi_panel(hv_dataset, None, plot_callback=cb, target_dimension=3)
    #     return matches.to_panel()

    # def to_scatter_jitter_da(self, ds: xr.Dataset, **kwargs) -> Optional[hv.Scatter]:
    #     matches = PlotFilter(
    #         float_range=VarRange(0, 0),
    #         cat_range=VarRange(0, None),
    #         repeats_range=VarRange(2, None),
    #         input_range=VarRange(1, None),
    #     ).matches_result(self.plt_cnt_cfg, "to_scatter_jitter")
    #     if matches.overall:

    #         print("DA IN",da)
    #         da = self.to_hv_dataset(ReduceType.NONE)
    #         hvds = hv.Dataset(da)
    #         # return None

    #         # print("DA FRESH",da)
    #         result_var = self.bench_cfg.result_vars[0]

    #         print(hvds.data.sizes)
    #         print("repeat" in hvds.data.sizes)
    #         # if "repeat" in hvds.data.sizes:
    #         # try:
    #         #     pt = (
    #         #         hvds.to(hv.Scatter)
    #         #         .opts(jitter=0.1)
    #         #         # .overlay()
    #         #         .overlay("repeat")
    #         #         .opts(show_legend=False, title=self.to_plot_title(), clabel=result_var.name, **kwargs)
    #         #     )
    #         # except:
    #         pt = (
    #             hvds.to(hv.Scatter)
    #             .opts(jitter=0.1)
    #             # .overlay()
    #             # .overlay("repeat")
    #             .opts(show_legend=False, title=self.to_plot_title(), clabel=result_var.name, **kwargs)
    #         )
    #         return pt
    #     return matches.to_panel()

    def to_scatter_jitter(
        self,
        result_var: Parameter = None,
        **kwargs,  # pylint: disable=unused-argument
    ) -> List[hv.Scatter]:
        return self.overlay_plots(partial(self.to_scatter_jitter_single, **kwargs))

    def to_scatter_jitter_single(self, result_var: Parameter, **kwargs) -> Optional[hv.Scatter]:
        matches = PlotFilter(
            float_range=VarRange(0, 0),
            cat_range=VarRange(0, None),
            repeats_range=VarRange(2, None),
            input_range=VarRange(1, None),
        ).matches_result(self.plt_cnt_cfg, "to_scatter_jitter")
        if matches.overall:
            ds = self.to_hv_dataset(ReduceType.NONE)
            pt = (
                ds.to(hv.Scatter, vdims=[result_var.name], label=result_var.name)
                .opts(jitter=0.1, show_legend=False, title=self.to_plot_title(), **kwargs)
                .overlay("repeat")
            )
            return pt
        return matches.to_panel()

    def to_heatmap_single(
        self, result_var: Parameter, reduce: ReduceType = ReduceType.AUTO, **kwargs
    ) -> hv.HeatMap:
        matches_res = PlotFilter(
            float_range=VarRange(2, None),
            cat_range=VarRange(0, None),
            input_range=VarRange(1, None),
        ).matches_result(self.plt_cnt_cfg, "to_heatmap")
        if matches_res.overall:
            z = result_var
            title = f"{z.name} vs ("

            for iv in self.bench_cfg.input_vars:
                title += f" vs {iv.name}"
            title += ")"

            color_label = f"{z.name} [{z.units}]"

            return self.to(hv.HeatMap, reduce).opts(clabel=color_label, **kwargs)
        return matches_res.to_panel()

    def to_heatmap_tap(
        self,
        result_var: Parameter,
        reduce: ReduceType = ReduceType.AUTO,
        width=800,
        height=800,
        **kwargs,
    ):
        htmap = self.to_heatmap_single(result_var, reduce).opts(
            tools=["hover", "tap"], width=width, height=height
        )
        htmap_posxy = hv.streams.Tap(source=htmap, x=0, y=0)

        def tap_plot(x, y):
            kwargs[self.bench_cfg.input_vars[0].name] = x
            kwargs[self.bench_cfg.input_vars[1].name] = y
            return self.get_nearest_holomap(**kwargs).opts(width=width, height=height)

        tap_htmap = hv.DynamicMap(tap_plot, streams=[htmap_posxy])
        return htmap + tap_htmap

    def to_nd_layout(self, hmap_name: str) -> hv.NdLayout:
        return hv.NdLayout(self.get_hmap(hmap_name), kdims=self.bench_cfg.hmap_kdims).opts(
            shared_axes=False, shared_datasource=False
        )

    def to_holomap(self, name: str = None) -> hv.HoloMap:
        return hv.HoloMap(self.to_nd_layout(name)).opts(shared_axes=False)

    def to_holomap_list(self, hmap_names: List[str] = None) -> hv.HoloMap:
        if hmap_names is None:
            hmap_names = [i.name for i in self.result_hmaps]
        col = pn.Column()
        for name in hmap_names:
            self.to_holomap(name)
        return col

    def get_nearest_holomap(self, name: str = None, **kwargs):
        canonical_inp = hmap_canonical_input(
            get_nearest_coords(self.ds, collapse_list=True, **kwargs)
        )
        return self.get_hmap(name)[canonical_inp].opts(framewise=True)

    def to_dynamic_map(self, name: str = None) -> hv.DynamicMap:
        """use the values stored in the holomap dictionary to populate a dynamic map. Note that this is much faster than passing the holomap to a holomap object as the values are calculated on the fly"""

        def cb(**kwargs):
            return self.get_hmap(name)[hmap_canonical_input(kwargs)].opts(
                framewise=True, shared_axes=False
            )

        kdims = []
        for i in self.bench_cfg.input_vars + [self.bench_cfg.iv_repeat]:
            kdims.append(i.as_dim(compute_values=True))

        return hv.DynamicMap(cb, kdims=kdims)

    def to_grid(self, inputs=None):
        if inputs is None:
            inputs = self.bench_cfg.inputs_as_str()
        if len(inputs) > 2:
            inputs = inputs[:2]
        return self.to_holomap().grid(inputs)

    def to_table(self):
        return self.to(hv.Table, ReduceType.SQUEEZE)

    def to_surface(self, result_var: Parameter = None, **kwargs) -> Optional[pn.panel]:
        return self.filter(
            self.to_surface_ds,
            float_range=VarRange(2, None),
            cat_range=VarRange(0, None),
            input_range=VarRange(1, None),
            reduce=ReduceType.REDUCE,
            target_dimension=2,
            result_var=result_var,
            result_types=(ResultVar),
            **kwargs,
        )

    def to_surface_ds(
        self, dataset: xr.Dataset, result_var: Parameter, alpha: float = 0.3, **kwargs
    ) -> Optional[pn.panel]:
        """Given a benchCfg generate a 2D surface plot

        Args:
            result_var (Parameter): result variable to plot

        Returns:
            pn.pane.holoview: A 2d surface plot as a holoview in a pane
        """
        matches_res = PlotFilter(
            float_range=VarRange(2, 2),
            cat_range=VarRange(0, None),
            vector_len=VarRange(1, 1),
            result_vars=VarRange(1, 1),
        ).matches_result(self.plt_cnt_cfg, "to_surface_hv")
        if matches_res.overall:
            # xr_cfg = plot_float_cnt_2(self.plt_cnt_cfg, result_var)

            # TODO a warning suggests setting this parameter, but it does not seem to help as expected, leaving here to fix in the future
            # hv.config.image_rtol = 1.0

            mean = dataset[result_var.name]

            hvds = hv.Dataset(dataset[result_var.name])

            x = self.plt_cnt_cfg.float_vars[0]
            y = self.plt_cnt_cfg.float_vars[1]

            try:
                surface = hvds.to(hv.Surface, vdims=[result_var.name])
                surface = surface.opts(colorbar=True)
            except Exception as e:
                logging.warning(e)

            if self.bench_cfg.repeats > 1:
                std_dev = dataset[f"{result_var.name}_std"]

                upper = mean + std_dev
                upper.name = result_var.name

                lower = mean - std_dev
                lower.name = result_var.name

                surface *= (
                    hv.Dataset(upper)
                    .to(hv.Surface)
                    .opts(alpha=alpha, colorbar=False, backend="plotly")
                )
                surface *= (
                    hv.Dataset(lower)
                    .to(hv.Surface)
                    .opts(alpha=alpha, colorbar=False, backend="plotly")
                )

            surface = surface.opts(
                zlabel=f"{result_var.name} [{result_var.units}]",
                title=f"{result_var.name} vs ({x.name} and {y.name})",
                backend="plotly",
                **kwargs,
            )

            if self.bench_cfg.render_plotly:
                hv.extension("plotly")
                out = surface
            else:
                # using render disabled the holoviews sliders :(
                out = hv.render(surface, backend="plotly")
            return pn.Column(out, name="surface_hv")

        return matches_res.to_panel()

    # def plot_scatter2D_hv(self, rv: ParametrizedSweep) -> pn.pane.Plotly:
    # import plotly.express as px

    #     """Given a benchCfg generate a 2D scatter plot

    #     Args:
    #         bench_cfg (BenchCfg): description of benchmark
    #         rv (ParametrizedSweep): result variable to plot

    #     Returns:
    #         pn.pane.Plotly: A 3d volume plot as a holoview in a pane
    #     """

    #     # bench_cfg = wrap_long_time_labels(bench_cfg)
    #     self.ds.drop_vars("repeat")

    #     df = self.to_pandas()

    #     names = rv.index_names()

    #     return px.scatter(
    #         df, x=names[0], y=names[1], marginal_x="histogram", marginal_y="histogram"
    #     )


HoloviewResult.set_default_opts()
