import param
from immutabledict import immutabledict
import numpy as np
import pandas as pd
import holoviews as hv
import straxen

# Define function which plots points with radius defined in data space:
# Taken from:
# https://stackoverflow.com/questions/60361810/how-do-i-set-the-scatter-circle-radius-in-holoviews
from holoviews.element.chart import Chart
from holoviews.plotting.bokeh import PointPlot


class PlotPMTArrayTPC:
    """Class to plot TPC PMT array.

    :param gains: If a list of gain values are supplied PMTs which gain are set to zero are
        displayed in gray.
    :param top_pmt_array: Boolean whether top or bottom PMT array should be plotted.

    """

    def __init__(
        self,
        gains=None,
        top_pmt_array=True,
        opts=immutabledict(),
    ):
        self.is_top_pmt_array = top_pmt_array

        self._init_start_end_index()
        self._init_pmt_array(gains)
        self._init_common_options(opts)

    def _init_start_end_index(self):
        """Defines index-range of channels to be used for the plot.

        Depends on whether the top or bottom array should be plotted.

        """
        self.pmt_start = 0 if self.is_top_pmt_array else straxen.n_top_pmts + 1
        self.pmt_end = straxen.n_top_pmts if self.is_top_pmt_array else straxen.n_tpc_pmts

    def _init_pmt_array(self, gains):
        """Initalized a dataframe which will contain the Top PMT pattern array."""
        _pmts = pd.DataFrame()
        _pmts["channel"] = np.arange(0, straxen.n_tpc_pmts, 1)
        _pmts["x"] = straxen.pmt_positions().loc[:, "x"].values
        _pmts["y"] = straxen.pmt_positions().loc[:, "y"].values
        _pmts["gains"] = np.ones(straxen.n_tpc_pmts)
        if np.any(gains):
            if len(gains) != straxen.n_tpc_pmts:
                raise ValueError("Expected gains to have the same shape as number of PMTs in TPC!")
            _is_zero = gains == 0
            _pmts.loc[_is_zero, "gains"] = np.nan
        self.pmts = _pmts.loc[self.pmt_start : self.pmt_end - 1, :]

    def _init_common_options(self, opt_dict):
        """Creates dictionary with common options required by PMT array plot."""
        if not isinstance(opt_dict, (dict, immutabledict)):
            raise ValueError("opts_dict must be a dictionary!")

        if self.is_top_pmt_array:
            title = "Top array top view"
        else:
            title = "Bottom array top view"

        r = straxen.tpc_r
        r *= 1.1
        default_settings = {
            "xlim": (-r, r),
            "ylim": (-r, r),
            "xlabel": "x [cm]",
            "ylabel": "y [cm]",
            "title": title,
        }
        self.settings = immutabledict({**default_settings, **opt_dict})

    @staticmethod
    def _plot_tpc(**opts):
        """Function which plots the TPC as a simple circle."""
        tpc = straxen.holoviews_utils.plot_tpc_circle(straxen.tpc_r)
        tpc.opts(**opts)
        return tpc

    def _plot_top_array(self, peak_data, label, **opts):
        pmts = pd.DataFrame()
        pmts["area"] = peak_data["area_per_channel"][self.pmt_start : self.pmt_end]

        # Cops other static values:
        # (we cannot use self.pmts as plot input as new plots will also change prev.
        # plots due to same reference.)
        pmts["channel"] = self.pmts["channel"]
        pmts["x"] = self.pmts["x"]
        pmts["y"] = self.pmts["y"]
        pmts["gains"] = self.pmts["gains"]

        pmts["area"] /= pmts["gains"]
        pmt_array = Circle(
            pmts,
            kdims=[hv.Dimension("x", label="x [cm]")],
            vdims=[
                hv.Dimension("y", label="y [cm]"),
                hv.Dimension("area", label="Area [pe]"),
                hv.Dimension("channel", label="Channel"),
            ],
            label=label,
        )
        pmt_array.opts(
            color="area",
            cmap="viridis",
            line_color="black",
            radius=3 * 2.54 / 2,
            tools=["hover"],
            clim=(10, None),
            clipping_colors={"NaN": "gray"},
        )
        pmt_array.opts(**opts)
        return pmt_array

    def plot_pmt_array(self, peak, label="", **opts):
        """Function which plots the specified PMT array for the given peak.

        :param peak: Peak for which PMT array should be plotted. Must contain area per channel
            information.
        :param label: Plot abel to be used as title.
        :param opts: Option which can be supplied to a holovies plot. E.g. you can use logz=True for
            a logarithmic colorscale or colorbar=True for adding a colorbar.
        :return: Holoviews plot overlay.

        """
        # Allow to return empty plots required for event display if alt.
        if peak is None:
            return hv.Points(None)

        plot = self._plot_tpc()
        plot *= self._plot_top_array(peak, label=label).opts(**self.settings, **opts)
        return plot


class Circle(Chart):
    group = param.String(default="Circle", constant=True)
    size = param.Integer()


class CirclePlot(PointPlot):
    _plot_methods = dict(single="circle", batched="circle")

    style_opts = ["radius" if so == "size" else so for so in PointPlot.style_opts if so != "marker"]


hv.Store.register({Circle: CirclePlot}, "bokeh")
