# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/plotting.controls.ipynb (unless otherwise specified).

__all__ = ['Filter', 'RangeFilter', 'CategoricalFilter', 'TimeFilter', 'ScatterFilter', 'GenericMulitScatterFilter']

# Cell
import datetime
from typing import Any, Iterable, Union, Tuple, List, Optional
from abc import ABCMeta, abstractmethod
from inspect import signature

import numpy as np
import pandas as pd
from bokeh.plotting import figure, show
from bokeh.io import push_notebook
from bokeh.models import ColumnDataSource, LassoSelectTool
import panel as pn
import panel.widgets as pnw


from .core import *
from .utils import *

# Cell
class Filter(metaclass=ABCMeta):
    """Abstract base class for filters."""
    def __init__(self, data, width=500, height=500) -> None:
        self.data = data
        self.width = width
        self.height = height
        self.build_gui()

    def mask_callback(self, callback):
        def new_callback(event):
            return callback(self.get_selection())
        return new_callback

    @abstractmethod
    def register_callback(self, callback):
        pass

    @abstractmethod
    def build_gui(self):
        pass

    @abstractmethod
    def get_selection(self, inverted):
        pass

    @abstractmethod
    def show(self):
        pass

# Cell
class RangeFilter(Filter):
    def __init__(self, data: np.ndarray, name: str, bins: int = 20, steps: int = 50, with_hist: bool = True, width: int = 500, height: int = 500) -> None:
        self.steps = steps
        self.step_width = (data.max()-data.min())/steps
        self.bins = bins
        self.name = name
        self.with_hist = with_hist
        super().__init__(data=data, width=width, height=height)

    def build_gui(self) -> None:
        val_min = min(self.data)
        val_max = max(self.data)
        # subtract and add a bit to the min and max value to ensure the whole range is captured
        dist = val_max-val_min if val_max != val_min else 1
        val_min = val_min-0.01*dist
        val_max = val_max+0.01*dist
        self.slider = pnw.RangeSlider(name=self.name, start=val_min, end=val_max, step=round(((val_max-val_min)/self.steps), 1), width=int(0.97*self.width))
        self.slider.param.watch(self.update_self, "value_throttled")
        if self.with_hist:
            self.hist = histogram(self.data, bins=self.bins, range=(self.data.min(), self.data.max()), height=100, width=self.width, remove_tools=True)
        else:
            self.hist = None
        self.gui = pn.Column(self.slider, self.hist)

    def update_self(self, event):
        mask = self.get_selection()
        self.update_with_mask(mask)

    def get_selection(self, inverted: bool = False) -> np.ndarray:
        selection = (self.data >= self.gui[0].value[0]) & (self.data <= self.gui[0].value[1])
        if inverted:
            selection = ~selection
        return selection

    def update_with_mask(self, mask: np.ndarray) -> None:
        if self.hist is None:
            return
        else:
            selection = self.data[np.logical_and(self.get_selection(), mask)]
            self.hist = histogram(selection, bins=self.bins, range=(self.data.min(), self.data.max()), height=100, width=self.width, remove_tools=True)
            self.gui[1] = self.hist

    def register_callback(self, callback):
        callback = self.mask_callback(callback)
        self.slider.param.watch(callback, "value_throttled")

    def show(self) -> pn.layout.base.Column:
        return self.gui

# Cell
class CategoricalFilter(Filter):
    """A multi select based filter for categorical data"""
    def __init__(self, data: np.ndarray, name: str, width: int = 500, height: int = 500) -> None:
        self.name = name
        super().__init__(data=data, width=width, height=height)

    def build_gui(self) -> None:
        self.selector = pnw.MultiSelect(name=self.name, options=np.unique(self.data).tolist(), value=np.unique(self.data).tolist(), width=self.width, height=self.height)
        self.gui = pn.Column(self.selector)

    def get_selection(self, inverted: bool = False) -> np.ndarray:
        selection = np.in1d(self.data, self.selector.value)
        if inverted:
            selection = ~selection
        return selection

    def show(self) -> pn.layout.base.Column:
        return self.gui

    def register_callback(self, callback):
        callback = self.mask_callback(callback)
        self.selector.param.watch(callback, "value")

    def update_with_mask(self, mask: np.ndarray, disabel_callbacks=True):
        """This does not work yet, maybe the whole structure needs to be changed to be based on the param lib."""
        self.selector.value = list(set([element for element, mask_value in zip(self.data, mask) if mask_value]))

# Cell
class TimeFilter(Filter):
    """A filter for time data with a start and an end date"""
    def __init__(self, start_times: Iterable[datetime.datetime], end_times: Iterable[datetime.datetime], start_time_label: str = "Start", end_time_label: str = "End", width: int = 500, height: int = 500) -> None:
        self.start_times = start_times
        self.end_times = end_times
        self.start_time_label = start_time_label
        self.end_time_label = end_time_label
        super().__init__((start_times, end_times), width=width, height=height)

    def build_gui(self):
        self.plot = time_arc_plot(self.start_times, self.end_times, width=self.width, height=self.height-100)
        min_start_time, max_start_time = get_min_and_max_dates(self.start_times)
        min_end_time, max_end_time = get_min_and_max_dates(self.end_times)
        min_date = min(min_start_time, min_end_time)
        max_date = max(max_start_time, max_end_time)
        self.start_time_slider = pnw.DateRangeSlider(name=self.start_time_label, start=min_date, end=max_date, width=self.width, height=50)
        self.end_time_slider = pnw.DateRangeSlider(name=self.end_time_label, start=min_date, end=max_date, width=self.width, height=50)
        self.gui = pn.Column(self.plot, self.start_time_slider, self.end_time_slider)

    def get_selection(self, inverted=False):
        start_selection = (self.start_times > self.start_time_slider.value[0]) & (self.start_times < self.start_time_slider.value[1])
        end_selection = (self.end_times > self.end_time_slider.value[0]) & (self.end_times < self.end_time_slider.value[1])
        return start_selection & end_selection

    def register_callback(self, callback):
        callback = self.mask_callback(callback)
        self.start_time_slider.param.watch(callback, "value_throttled")
        self.end_time_slider.param.watch(callback, "value_throttled")

    def show(self):
        return self.gui

# Cell
class ScatterFilter(Filter):
    """A filter based on a scatter plot with a lasso selection."""
    def __init__(self, x, y, x_label: str = "", y_label: str = "", width: int = 500, height: int = 500):
        self.x_label = x_label
        self.y_label = y_label
        super().__init__((x,y), width=width, height=height)

    def build_gui(self):
        self.source = ColumnDataSource({"x": self.data[0], "y": self.data[1]})
        p = figure(x_axis_label=self.x_label, y_axis_label=self.y_label, width=self.width, height=self.height, tools="lasso_select")
        p.scatter("x", "y", source=self.source)
        p.select_one(LassoSelectTool).select_every_mousemove = False
        self.figure = p
        self.gui = pn.Row(self.figure)

    def get_selection(self):
        selected_indices = self.source.selected.indices
        selection = [False if i not in selected_indices else True for i in range(len(self.data[0]))]
        return selection

    def mask_callback(self, callback):
        def new_callback(attr, new, old):
            callback(self.get_selection())
        return new_callback

    def register_callback(self, callback):
        callback = self.mask_callback(callback)
        self.source.selected.on_change("indices", callback)

    def show(self):
        return self.gui

# Cell
class GenericMulitScatterFilter(Filter):
    """A generic filter base on the scatter plot filter, that provides additional inputs for column selection and how the selections over the different columns should be combined."""
    def __init__(self, data, columns: Optional[List[str]] = None, mode: str = "symmetric", width: int = 500, height: int = 500):
        """mode: symmetric (default) or singular. If symmetric for a specific x-y combination the corrosponding y-x combination will be updated with the same values. Else they have different selections."""
        self.mode = mode
        self.columns = columns if columns is not None else data.columns
        self.source = ColumnDataSource(data)
        self.source.selected.on_change("indices", self.update_selection)
        self.selections = {x_key: {y_key: [False]*data.shape[0] for y_key in self.columns} for x_key in self.columns}
        super().__init__(data, width, height)

    def update_selection(self, attr, old, new):
        self.selections[self.x_select.value][self.y_select.value] = [True if index in new else False for index in range(len(self.selections[self.x_select.value][self.y_select.value]))]
        if self.mode == "symmetric":
            self.selections[self.y_select.value][self.x_select.value] = [True if index in new else False for index in range(len(self.selections[self.x_select.value][self.y_select.value]))]
        self.source.data["colors"] = self.get_colors()

    def update_plot(self, event):
        self.gui[3] = self.scatter_plot()

    def build_gui(self):
        self.x_select = pnw.Select(name="x-Axis", options=list(self.data.columns), value=self.data.columns[0], width=self.width//2, height=50)
        self.x_select.param.watch(self.update_plot, "value")
        second_select_index = 1 if len(self.data.columns) > 1 else 0
        self.y_select = pnw.Select(name="y-Axis", options=list(self.data.columns), value=self.data.columns[second_select_index], width=self.width//2, height=50)
        self.y_select.param.watch(self.update_plot, "value")
        self.combine_selections = pnw.RadioBoxGroup(name="Combine selections", options=["None", "Or", "And"], value="None", width=self.width//3, inline=True)
        self.combine_selections.param.watch(self.update_plot, "value")
        self.ignore_empty_selections = pnw.CheckBoxGroup(name="Ignore empty selections", value=["Ignore empty selections"], options=["Ignore empty selections"])
        self.ignore_empty_selections.param.watch(self.update_plot, "value")
        p = self.scatter_plot()
        legend = pn.Row("<p style='color:#B22222';>Selected in current</p>", "<p style='color:#FF00FF';>Selected in another</p>", "<p style='color:#808080';>Not selected</p>")
        self.gui = pn.Column(
            pn.Row(self.x_select, self.y_select),
            pn.Row("Combine selections with: ", self.combine_selections, self.ignore_empty_selections),
#             pn.Row("Ignore empty selections", ),
            legend, p
        )

    def scatter_plot(self):
        p = figure(x_axis_label=self.x_select.value, y_axis_label=self.y_select.value, width=self.width, height=self.height-50, tools="lasso_select")
        # get selections
        if self.combine_selections.value == "None":
            combined_selection_selected = self.selections[self.x_select.value][self.y_select.value]
        else:
            # put all singel selections into one list so they can be iterated over with zip
            if "Ignore empty selections" in self.ignore_empty_selections.value:
                single_selections = [selection for layer_1 in self.selections.values() for selection in layer_1.values() if any(selection)]
                if len(single_selections) == 0:
                    single_selections = [selection for selection in single_selections if any(selection)]
            else:
                single_selections = [selection for layer_1 in self.selections.values() for selection in layer_1.values()]
            if self.combine_selections.value == "And":
                combined_selection_selected = [all(point_selections) for point_selections in zip(*single_selections)]
            else:
                combined_selection_selected = [any(point_selections) for point_selections in zip(*single_selections)]
        self.source.selected.indices = [index for index, value in enumerate(combined_selection_selected) if value]
        self.source.data["colors"] = self.get_colors()
        radius = (self.data[self.x_select.value].max()-self.data[self.x_select.value].min() + self.data[self.y_select.value].max()-self.data[self.y_select.value].min())* 0.01
        p.scatter(self.x_select.value, self.y_select.value, source=self.source, fill_color="colors", line_color=None, radius=radius)
        p.select_one(LassoSelectTool).select_every_mousemove = False
        return p

    def get_colors(self):
        # create colors from selections
        colors = ["gray"]*self.data.shape[0]
        current_selections =self.selections[self.x_select.value][self.y_select.value]
        single_selections = [selection for layer_1 in self.selections.values() for selection in layer_1.values()]
        any_selections = [any(point_selections) for point_selections in zip(*single_selections)]
        for index, (any_selection, current_selection) in enumerate(zip(any_selections, current_selections)):
            if any_selection is True and current_selection is False:
                colors[index] = "fuchsia"
            elif current_selection is True:
                colors[index] = "firebrick"
        return colors

    def get_selection(self):
        if self.combine_selections.value == "None":
            selection = [value for value in self.selections[self.x_select.value][self.y_select.value]]
        else:
            if "Ignore empty selections" in self.ignore_empty_selections.value:
                single_selections = [selection for layer_1 in self.selections.values() for selection in layer_1.values() if any(selection)]
                if len(single_selections) == 0:
                    single_selections = [selection for selection in single_selections if any(selection)]
            else:
                single_selections = [selection for layer_1 in self.selections.values() for selection in layer_1.values()]
            if self.combine_selections.value == "And":
                selection = [all(point_selections) for point_selections in zip(*single_selections)]
            else:
                selection = [any(point_selections) for point_selections in zip(*single_selections)]
        return selection

    def mask_callback(self, callback):
        def new_callback(attr, old, new):
            return callback(self.get_selection())
        return new_callback

    def register_callback(self, callback):
        callback = self.mask_callback(callback)
        self.source.selected.on_change("indices", callback)

    def update_with_mask(self, mask):
        pass

    def show(self):
        return self.gui