import param
import panel as pn
import holoviews as hv
import numpy as np
from immutabledict import immutabledict


class EventStats(param.Parameterized):
    """Class which displays all events in a user defined data space. Marks position of single
    selected event depending on the passed index or index widget.

    :param events: Events for which histogram should be plotted.
    :param default_config: Dictionary containing the default settings
        for the plot.
    :param exclude_fields: A list of fields to be excluded from the
        x/y-selection widget. Can be uselful as event_info has a lot of
        fields. Can be either a list/tuple or a function. If passed as
        function must take as an argument a list with field names and
        return a list with field names.
    :param _index_slider: Parameter to pass an external integer slider
        to class. E.g. required in event display.

    Warning:
        In case the initally selected value is outside of the specified
        plotting range the plot will break.

    """

    _x_edges = None
    _y_edges = None
    _entries = None
    _x_field_name = None
    _y_field_name = None
    _x_bins = None
    _y_bins = None
    _x_log = False
    _y_log = False
    _z_log = False

    # Static widgets used for setting plot layout.
    x_field = "cs1"
    x_bins = param.Tuple(default=(200, 0, 50), label="x binning (#bins, min max)")
    x_log = param.Boolean(default=False, label="Logarithmic x binning (np.logspace)")
    y_field = "cs2"
    y_bins = param.Tuple(default=(200, 1000, 2500), label="y binning (#bins, min max)")
    y_log = param.Boolean(default=False, label="Logarithmic y binning (np.logspace)")
    clim = param.Tuple(default=(1, None), label="Colorbarlimit (min, max)")
    z_log = param.Boolean(default=True, label="Logarithmic colorbar")
    VALID_CONFIGS = ["x_field", "x_bins", "x_log", "y_field", "y_bins", "y_log", "clim", "z_log"]

    def __init__(
        self,
        events,
        default_config=immutabledict(),
        exclude_fields=None,
        _index_slider=None,
        **params,
    ):
        for key, value in default_config.items():
            if not hasattr(self, key):
                raise ValueError(
                    f'"{key}" is not a valid config. '
                    f"Only the following attributes are recognized: {self.VALID_CONFIGS}."
                )
            setattr(self, key, value)

        self.exclude_fields = exclude_fields
        self.events = events
        self._create_widgets(
            self.x_field,
            self.y_field,
            _index_slider,
        )
        super().__init__(**params)

    def interactive_event_stats(self):
        """An interactive plot which displays a histogram of the supplied data.

        The user can manually change the axis and binning via the provided widgets. This plot is
        meant to be used as a stand alone outside of the event display.

        """
        stream_hist, stream_marker = self._create_streams()
        hist = hv.DynamicMap(self._plot_hist_data, streams=stream_hist)
        selected_points = hv.DynamicMap(self._plot_selected_event, streams=stream_marker)
        return pn.Column(
            (hist * selected_points).opts(width=400),
            self._make_widget_panel(include_index_slider=True),
        )

    def _create_streams(self):
        """Creates stream objects for the histogram and marker plot."""
        stream_hist = dict(
            x_field=self.x_field_widget,
            y_field=self.y_field_widget,
            x_bins=self.param.x_bins,
            y_bins=self.param.y_bins,
            clim=self.param.clim,
            x_log=self.param.x_log,
            y_log=self.param.y_log,
            z_log=self.param.z_log,
        )

        stream_marker = dict(
            x_field=self.x_field_widget, y_field=self.y_field_widget, index=self.index_slider
        )

        return stream_hist, stream_marker

    def _create_widgets(
        self,
        x_field="cs1",
        y_field="cs2",
        _index_slider=None,
    ):
        """Creates widgets for the x and y field as well as an slider for the event index if not
        supplied externally."""
        fields = list(self.events.dtype.names)
        if self.exclude_fields:
            if hasattr(self.exclude_fields, "__call__"):
                fields = self.exclude_fields(fields)
            elif isinstance(self.exclude_fields, (list, tuple)):
                fields = self._exclude_fields(fields, self.exclude_fields)
            else:
                raise ValueError('"exclude_fields" must be either a list or a function!')

        fields.sort()
        self.x_field_widget = pn.widgets.Select(value=x_field, options=fields, name="X-field")

        self.y_field_widget = pn.widgets.Select(value=y_field, options=fields, name="Y-field")

        if not _index_slider:
            self.index_slider = pn.widgets.IntSlider(
                start=0, end=len(self.events), name="Event Index:"
            )
        else:
            self.index_slider = _index_slider

    @staticmethod
    def _exclude_fields(fields, excluded_fields):
        """Function which drops certain fields from."""
        dtype = []
        for f in fields:
            if f not in excluded_fields:
                dtype.append(f)
        return dtype

    def _make_widget_panel(self, include_index_slider=False):
        """Put widgets used in the selection plot into a single panel which can be displayed.

        Adds index slider for stand alone plot.

        """
        widget_panel = pn.Row(
            pn.Column(
                self.x_field_widget,
                self.param.x_bins,
                self.param.x_log,
            ),
            pn.Column(
                self.y_field_widget,
                self.param.y_bins,
                self.param.y_log,
            ),
            pn.Column(
                self.param.clim,
                self.param.z_log,
            ),
            scroll=True,
        )

        if include_index_slider:
            widget_panel = pn.Column(self.index_slider, widget_panel)

        return widget_panel

    def _hist_data(self, x_field, y_field, xbins, ybins, x_log, y_log):
        """Function which checks if underlying data has to be rebinned."""
        _hist_has_not_changed = True
        _hist_has_not_changed &= x_field == self._x_field_name
        _hist_has_not_changed &= y_field == self._y_field_name
        _hist_has_not_changed &= xbins == self._x_bins
        _hist_has_not_changed &= x_log == self._x_log
        _hist_has_not_changed &= ybins == self._y_bins
        _hist_has_not_changed &= y_log == self._y_log

        if not _hist_has_not_changed:
            x_binning = self._get_binning(*xbins, x_log)
            y_binning = self._get_binning(*ybins, y_log)

            entries, x_edges, y_edges = np.histogram2d(
                self.events[x_field],
                self.events[y_field],
                bins=(
                    x_binning,
                    y_binning,
                ),
            )

            self._x_edges = x_edges
            self._y_edges = y_edges
            self._entries = entries.T

            self._x_field_name = x_field
            self._y_field_name = y_field
            self._x_bins = xbins
            self._y_bins = ybins
            self._x_log = x_log
            self._y_log = y_log

    @staticmethod
    def _get_binning(n_bins, bin_min, bin_max, log):
        """Function which will return an array of bin edges."""
        if log:
            bin_min = max(10**-3, bin_min)
            bins = np.logspace(np.log10(bin_min), np.log10(bin_max), n_bins)
        else:
            bins = np.linspace(bin_min, bin_max, n_bins)
        return bins

    def _plot_hist_data(
        self,
        x_field: str,
        y_field: str,
        x_bins: tuple,
        y_bins: tuple,
        clim: tuple = (1, None),
        x_log: bool = False,
        y_log: bool = False,
        z_log: bool = True,
    ):
        """Creates histogram based on user specified settings."""
        self._hist_data(
            x_field,
            y_field,
            x_bins,
            y_bins,
            x_log,
            y_log,
        )
        hist = hv.QuadMesh(
            (self._x_edges, self._y_edges, self._entries),
            [x_field, y_field],
        )
        hist = hist.apply.opts(
            logz=z_log,
            logx=x_log,
            logy=y_log,
            clim=clim,
            cmap="viridis",
            colorbar=True,
            framewise=True,
            clipping_colors={"min": "white"},
        )
        return hist

    def _plot_selected_event(self, x_field: str, y_field: str, index: int = 0):
        """Function which highlights event with selected index in plot."""
        point_x = self.events[x_field][index]
        point_y = self.events[y_field][index]
        _is_nan = np.isnan(point_x) or np.isnan(point_y)
        if not _is_nan:
            selected_points = hv.Points(
                (point_x, point_y), [self.x_field_widget.value, self.y_field_widget.value]
            ).opts(color="red", size=7, alpha=1)

            hline = hv.HLine(point_y).opts(color="red", alpha=0.4)
            vline = hv.VLine(point_x).opts(color="red", alpha=0.4)
        else:
            # In case one of the parameters is not defined
            # plot invisilbe point
            selected_points = hv.Points(
                (0, 0), [self.x_field_widget.value, self.y_field_widget.value]
            ).opts(color="k", alpha=0, size=7)

            hline = hv.HLine(0).opts(color="k", alpha=0)
            vline = hv.VLine(0).opts(color="k", alpha=0)

        return selected_points * hline * vline
