# $$_ Lines starting with # $$_* autogenerated by jup_mini. Do not modify these
# $$_markdown
# # Interactive Plot
# $$_end_markdown
# $$_markdown
# # Description
# 
# Allows interactive plotting of multidimensional data
# $$_end_markdown
# $$_code
# $$_ %%checkall
from __future__ import annotations
import os
import sys
import math
import colorsys
from dataclasses import dataclass
import unittest
import doctest
import pandas as pd
import numpy as np
from IPython.display import display, clear_output
from ipywidgets import widgets
import plotly
import plotly.callbacks
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from typing import Callable, Any
from collections.abc import Sequence
import traitlets
from pyqstrat.pq_utils import bootstrap_ci, get_child_logger

ROOT_DIR = os.path.join(sys.path[0])
sys.path.insert(1, ROOT_DIR)

_logger = get_child_logger(__name__)

DEFAULT_PLOTLY_COLORS = ['rgb(31, 119, 180)', 'rgb(255, 127, 14)',
                         'rgb(44, 160, 44)', 'rgb(214, 39, 40)',
                         'rgb(148, 103, 189)', 'rgb(140, 86, 75)',
                         'rgb(227, 119, 194)', 'rgb(127, 127, 127)',
                         'rgb(188, 189, 34)', 'rgb(23, 190, 207)']

LineDataType = tuple[str, pd.DataFrame, dict[Any, pd.DataFrame]]
    
DimensionFilterType = Callable[[
    pd.DataFrame,
    str,
    list[tuple[str, Any]]],
    np.ndarray]

DataFilterType = Callable[[
    pd.DataFrame,
    list[tuple[str, Any]]],
    pd.DataFrame]

StatFuncType = Callable[[pd.DataFrame, str, str, str], list[LineDataType]]

DetailDisplayType = Callable[[
    widgets.Widget, 
    pd.DataFrame,
    bool],
    None]


PlotFuncType = Callable[[str, str, list[LineDataType]], list[widgets.Widget]]

DataFrameTransformFuncType = Callable[[pd.DataFrame], pd.DataFrame]

SeriesTransformFuncType = Callable[[pd.Series], pd.Series]

DisplayFormFuncType = Callable[[Sequence[widgets.Widget], bool], None]

UpdateFormFuncType = Callable[[int], None]

CreateSelectionWidgetsFunctype = Callable[[dict[str, str], dict[str, str], UpdateFormFuncType], dict[str, Any]]


def percentile_buckets(a: np.ndarray, n=10) -> np.ndarray:
    '''
    >>> np.random.seed(0)
    >>> a = np.random.uniform(size=10000)
    >>> assert np.allclose(np.unique(percentile_buckets(a)), np.arange(0.05, 1, 0.1), atol=0.01)
    '''
    if not len(a): return np.empty(0)
    pctiles = np.arange(0, 100, int(round(100 / n))) 
    buckets = np.nanpercentile(a, pctiles)
    conditions: list[Any] = []
    for i, bucket in enumerate(buckets[:-1]):
        if buckets[i] == buckets[i + 1]:
            conditions.append(a == buckets[i])
        else:
            conditions.append((a >= buckets[i]) & (a < buckets[i + 1]))
    conditions.append((a >= buckets[-1]))
    b = [np.mean(a[cond]) for cond in conditions]
    ret = np.select(conditions, b)
    return ret


def display_form(form_widgets: Sequence[widgets.Widget], debug=False) -> None:
    if not debug: clear_output()
    box_layout = widgets.Layout(
        display='flex',
        flex_flow='column',
        align_items='stretch',
        border='solid',
        width='100%')
    box = widgets.Box(children=list(form_widgets), layout=box_layout)
    display(box)
    
    
class SimpleTransform:
    '''
    Initial transformation of data.  For example, you might add columns that are quantiles of other columns
    '''
    def __init__(self, transforms: list[tuple[str, str, SeriesTransformFuncType]] | None = None) -> None:
        self.transforms = [] if transforms is None else transforms
        
    def __call__(self, data: pd.DataFrame) -> pd.DataFrame:
        for (colname, new_colname, func) in self.transforms:
            data[new_colname] = func(data[colname])
        return data
    

def simple_dimension_filter(data: pd.DataFrame, dim_name: str, selected_values: list[tuple[str, Any]]) -> np.ndarray:
    '''
    Produces a list to put into a dropdown for selecting a dimension value
    '''
    mask = np.full(len(data), True)
    for name, value in selected_values:
        if value == 'All': continue
        mask &= (data[name] == value)
    values = np.unique(data[mask][dim_name].values)  # will sort values before returning them
    return ['All'] + values.tolist()
    
    
def simple_data_filter(data: pd.DataFrame, selected_values: list[tuple[str, Any]]) -> pd.DataFrame:
    '''
    Filters a dataframe based on the selected values
    '''
    mask = np.full(len(data), True)
    for name, value in selected_values:
        if value == 'All': continue
        mask &= (data[name] == value)
    return data[mask]
    
    
class MeanWithCI:
    '''
    Computes mean (or median) and optionally confidence intervals for plotting
    '''
    def __init__(self, mean_func: Callable[[np.ndarray], np.ndarray] = np.nanmean, ci_level: int = 0) -> None:  # type: ignore
        '''
        Args:
            mean: The function to compute ci for
            ci_level: Set to 0 for no confidence intervals, or the level you want.
                For example, set to 95 to compute a 95% confidence interval. Default 0
        '''
        self.mean_func = mean_func
        self.ci_level = ci_level
        
    def __call__(self, filtered_data: pd.DataFrame, xcol: str, ycol: str, zcol: str) -> list[LineDataType]:
        '''
        For each unique value of x and z, compute mean (and optionally ci) of y.
        Return:
            x, y data for plotting lines of the mean of y versus x for each z and the data used to compute the mean
        '''
        zvals = np.unique(filtered_data[zcol])
        cols = [col for col in filtered_data.columns if col not in [xcol, ycol, zcol]]
        df = filtered_data[[xcol, ycol, zcol] + cols]
        ret = []
        columns = [xcol, ycol] if not self.ci_level else [xcol, ycol, f'ci_d_{self.ci_level}', f'ci_u_{self.ci_level}']
        for zvalue in zvals:
            df = filtered_data[filtered_data[zcol] == zvalue]
            plt_data: list[Any] = []
            for x, yseries in df.groupby(xcol)[ycol]:
                y = yseries.values
                if not len(y): raise Exception(y)
                mean = self.mean_func(y)
                if self.ci_level:
                    ci_up, ci_down = bootstrap_ci(y, ci_level=self.ci_level / 100)
                    plt_data.append((x, mean, ci_down, ci_up))
                else:
                    plt_data.append((x, mean))
            line = pd.DataFrame.from_records(plt_data, columns=columns)
            ret.append((zvalue, line, df))
        return ret
    
    
class SimpleDetailTable:
    '''
    Displays a pandas DataFrame under a plot that contains the data used to compute a statistic of y for each x, y pair
    '''
    
    def __init__(self, 
                 colnames: list[str] | None = None, 
                 float_format: str = '{:.4g}', 
                 min_rows: int = 100, 
                 copy_to_clipboard: bool = True) -> None:
        '''
        Args:
            colnames: list of column names to display. If None we display all columns. Default None
            float_format: Format for each floating point column. Default {:.4g}
            min_rows: Do not truncate the display of the table before this many rows. Default 100
            copy_to_clipboard: If set, we copy the dataframe to the clipboard. On linux, you must install xclip for this to work
       '''
        self.colnames = colnames
        self.float_format = float_format
        self.min_rows = min_rows
        self.copy_to_clipboard = True
        
    def __call__(self, detail_widget: widgets.Widget, data: pd.DataFrame, debug=False) -> None:
        '''
        Args:
            detail_widget: The widget to display the data in
            data: The dataframe to display
        '''
        if self.float_format:
            orig_float_format = pd.options.display.float_format
            pd.options.display.float_format = (self.float_format).format
        
        if self.min_rows:
            orig_min_rows = pd.options.display.min_rows
            pd.options.display.min_rows = self.min_rows
            
        with detail_widget:
            if not debug: clear_output()
            if self.colnames: data = data[self.colnames]
            data = data.reset_index(drop=True)
            display(data)
            if self.copy_to_clipboard: data.to_clipboard(index=False)
                
        if self.float_format: pd.options.display.float_format = orig_float_format
        if self.min_rows: pd.options.display.min_rows = orig_min_rows
            
            
def create_selection_dropdowns(dims: dict[str, str], labels: dict[str, str], update_form_func: UpdateFormFuncType) -> dict[str, Any]:
    '''
    Create a list of selection widgets
    '''
    selection_widgets: dict[str, widgets.Widget] = {}
    for name in dims.keys():
        label = labels[name] if name in labels else name
        widget = widgets.Dropdown(description=label, style={'description_width': 'initial'})
        selection_widgets[name] = widget
        
    for widget in selection_widgets.values():
        widget.observe(lambda x: on_widgets_updated(x, update_form_func, selection_widgets), names='value')

    return selection_widgets


def on_widgets_updated(change: traitlets.utils.bunch.Bunch, update_form_func, selection_widgets: dict[str, widgets.Widget]) -> None:
    '''
    Callback called by plotly when widgets are updated by the user.
    '''
    owner = change['owner']
    widgets = list(selection_widgets.values())
    owner_idx = widgets.index(owner)
    update_form_func(owner_idx)
            
            
@dataclass
class LineConfig:
    color: str | None = None
    thickness: float = math.nan
    secondary_y: bool = False
    marker_mode: str = 'lines+markers'
    show_detail: bool = True
        
        
def _plotly_color_to_rgb(plotly_color: str) -> tuple[int, int, int]:
    '''
    Convert plotly color which is a string into r, g, b values
    >>> assert _plotly_color_to_rgb('rgb(31, 119, 180)') == (31, 119, 180)
    '''
    plotly_color = plotly_color.replace('rgb(', '').replace(')', '')
    s = plotly_color.split(',')
    r, g, b = int(s[0]), int(s[1]), int(s[2])
    return r, g, b

        
def _lighten_color(r: int, g: int, b: int) -> tuple[int, int, int]:
    '''
    Lighten color so we can show confidence intervals in a lighter shade than the line itself
    We convert to hls and increase lightness and decrease saturation
    >>> assert _lighten_color(31, 119, 180) == (102, 168, 214)
    '''
    hls = colorsys.rgb_to_hls(r, g, b)
    light_hls = (hls[0], hls[1] * 1.5, hls[2] * 0.5)
    rgb = colorsys.hls_to_rgb(*light_hls)
    rgb = (int(round(rgb[0])), int(round(rgb[1])), int(round(rgb[2])))
    return rgb


def foo(name, old, new):
    import datetime
    print(f'hello: {datetime.datetime.now()} {name} {old} {new}')
    
        
class LineGraphWithDetailDisplay:
    '''
    Draws line graphs and also includes a detail pane.
    When you click on a point on the line graph, the detail pane shows the data used to compute that point.
    '''
    
    def __init__(self, 
                 display_detail_func: DetailDisplayType = SimpleDetailTable(), 
                 line_configs: dict[str, LineConfig] = {}, 
                 title: str | None = None, 
                 hovertemplate: str | None = None, 
                 debug=False) -> None:
        '''
        Args:
            display_detail_func: A function that displays the data on the detail pane. Default SimpleDetailTable
            line_configs: Configuration of each line. The key in this dict is the zvalue for that line.  Default {}
            title: Title of the graph. Default None
            hovertemplate: What to display when we hover over a point on the graph.  See plotly hovertemplate
        '''
        self.display_detail_func = display_detail_func
        self.line_configs = line_configs
        self.title = title
        self.hovertemplate = hovertemplate
        self.debug = debug
        self.default_line_config = LineConfig()
        self.detail_data: dict[Any, pd.DataFrame] = {}
        self.xcol = ''
        self.zvalues: dict[int, Any] = {}  # trace index by zvalue
            
    def __call__(self, xaxis_title: str, yaxis_title: str, line_data: list[LineDataType]) -> list[widgets.Widget]:
        '''
        Draw the plot and also set it up so if you click on a point, we display the data used to compute that point.
        Args:
            line_data: The zvalue, plot data, and detail data for each line to draw. The plot data must have 
            x as the first column and y as the second column
         Return:
            A list of widgets to draw.  In this case, a figure widget and a output widget which contains the detail display
        '''
        if not len(line_data): return []
        self.detail_data.clear()
        secondary_y = any([lc.secondary_y for lc in self.line_configs.values()])
        
        fig_widget = go.FigureWidget(make_subplots(specs=[[{"secondary_y": secondary_y}]]))
        # fig_widget.on_trait_change(foo, '_js2py_restyle')
        # fig_widget.on_trait_change()
        detail_widget = widgets.Output()
        
        trace_num = 0

        for line_num, (zvalue, line_df, _detail_data) in enumerate(line_data):
            x = line_df.iloc[:, 0].values
            self.xcol = line_df.columns[0]
            y = line_df.iloc[:, 1].values
            self.detail_data[zvalue] = _detail_data
            line_config = self.line_configs[zvalue] if zvalue in self.line_configs else self.default_line_config
            marker_mode = line_config.marker_mode
            color = line_config.color if line_config.color else DEFAULT_PLOTLY_COLORS[line_num]
            
            hovertemplate = self.hovertemplate
            customdata = None
            
            if hovertemplate is None:
                unique, counts = np.unique(_detail_data[self.xcol].values, return_counts=True)
                customdata = counts[np.searchsorted(unique, x)]
                hovertemplate = 'N: %{customdata}'  # number of entries used to compute each x
                hovertemplate += f' Series: {zvalue} {xaxis_title}: ' + '%{x:.4g} ' + f'{yaxis_title}: ' + '%{y:.4g}'
                
            trace = go.Scatter(
                x=x,
                y=y,
                customdata=customdata,
                mode=marker_mode,
                name=str(zvalue),
                line=dict(color=color),
                hovertemplate=hovertemplate              
            )
            
            self.zvalues[trace_num] = zvalue
            
            fig_widget.add_trace(trace, secondary_y=line_config.secondary_y)
            if line_config.show_detail:
                fig_widget.data[trace_num].on_click(self._on_graph_click, append=True)
            trace_num += 1

            if len(line_df.columns) > 2:  # x, y, ci up and ci down 
                fill_color = _plotly_color_to_rgb(color)
                fill_color = _lighten_color(*fill_color)
                # we set transparency to 0.5 so we can see lines under the ci fill
                fill_color_str = f'rgba({fill_color[0]},{fill_color[1]},{fill_color[2]},0.5)'
                ci_down = line_df.iloc[:, 2].values
                ci_up = line_df.iloc[:, 3].values
                ci_trace = go.Scatter(
                    x=np.concatenate([x, x[::-1]]),  # x, then x reversed
                    y=np.concatenate([ci_up, ci_down[::-1]]),  # upper, then lower reversed
                    fill='toself',
                    fillcolor=fill_color_str,
                    line=dict(color='rgba(255,255,255,0)'),
                    hoverinfo="skip",
                    showlegend=False)
            
                fig_widget.add_trace(ci_trace, secondary_y=line_config.secondary_y)
                trace_num += 1
            
        fig_widget.update_layout(title=self.title, xaxis_title=xaxis_title)
        fig_widget.update_layout(yaxis_title=yaxis_title)
        
        if secondary_y:
            fig_widget.update_yaxes(title_text=yaxis_title, secondary_y=True)
            
        self.fig_widget = fig_widget
        self.detail_widget = detail_widget
        self.line_data = line_data
                        
        return [self.fig_widget, self.detail_widget]
    
    def _on_graph_click(self, 
                        trace: go.Trace, 
                        points: plotly.callbacks.Points, 
                        selector: plotly.callbacks.InputDeviceState) -> None:
        '''
        Callback called by plotly when you click on a point on the graph.
        When you click on a point, we display the dataframe with the data we used to compute that point.
        '''
        if not len(points.xs): return
        trace_idx = points.trace_index
        zvalue = self.zvalues[trace_idx]
        _detail_data = self.detail_data[zvalue]
        _detail_data = _detail_data[_detail_data[self.xcol].values == points.xs[0]]
        self.display_detail_func(self.detail_widget, _detail_data, self.debug)


class InteractivePlot:
    '''
    Creates a multidimensional interactive plot off a dataframe.
    '''
    def __init__(self,
                 data: pd.DataFrame,
                 labels: dict[str, str] | None = None,
                 transform_func: DataFrameTransformFuncType = SimpleTransform(),
                 create_selection_widgets_func: CreateSelectionWidgetsFunctype = create_selection_dropdowns,
                 dim_filter_func: DimensionFilterType = simple_dimension_filter,
                 data_filter_func: DataFilterType = simple_data_filter,
                 stat_func: StatFuncType = MeanWithCI(),
                 plot_func: PlotFuncType = LineGraphWithDetailDisplay(),
                 display_form_func: DisplayFormFuncType = display_form,
                 debug=False) -> None:
        '''
        Args:
            data: The pandas dataframe to use for plotting
            labels: A dict where column names from the dataframe are mapped to user friendly labels. For any column names
                not found as keys in this dict, we use the column name as the label. Default None
            dim_filter_func: A function that generates the values of a dimension based on other dimensions. For example, if 
                the user chooses "Put Option" in a put/call dropdown, the valid strikes could change in a Strike 
                dropdown that follows. Default simple_dimension_filter
            data_filter_func: A function that filters the data to plot. For example, if the user chooses "Put Option" in a put/call dropdown,
                we could filter the dataframe to only include put options. Default simple_data_filter
            stat_func: Once we have filtered the data, we may need to plot some statistics, such as mean and confidence intervals.
                In this function, we compute these statistics. Default MeanWithCI()
            plot_func: A function that plots the data.  This could also display detail data used to compute the statistics associated
                with each data point.
            display_form_func: A function that displays the form given a list of plotly widgets (including the graph widget)
            debug: Dont clear forms if this is true so we can see print output
        '''
        self.data = data
        self.transform_func = transform_func
        self.create_selection_widgets_func = create_selection_widgets_func
        if labels is None: labels = {}
        self.labels = labels
        self.dim_filter_func = dim_filter_func
        self.data_filter_func = data_filter_func
        self.stat_func = stat_func
        self.plot_func = plot_func
        self.display_form_func = display_form_func
        self.selection_widgets: dict[str, Any] = {}
        self.debug = debug
        
    def create_pivot(self, xcol: str, ycol: str, zcol: str, dimensions: dict[str, Any]) -> None:
        '''
        Create the initial pivot
        Args:
            xcol: Column name to use as the x axis in the DataFrame
            ycol: Column name to use as the y axis in the DataFrame
            zcol: Column name to use for z-values. Each zvalue can be used for a different trace within this plot. For example, a column
            called "option_type" could contain the values "American", "European", "Bermudan" and we could plot the data for each type
            in a separate trace
            dimensions: The column names used for filter dimensions. For example, we may want to filter by days to expiration and put/call
            The key the column name and the value is the initial value for that column. For example, in a 
            dropdown for Put/Call we may want "Put" to be the initial value set in the dropdown.  Set to None if you 
            don't care what initial value is chosen.
        '''
        self.xlabel = xcol if xcol not in self.labels else self.labels[xcol]
        self.ylabel = ycol if ycol not in self.labels else self.labels[ycol]
        self.zcol = zcol
        self.xcol = xcol
        self.ycol = ycol
        self.selection_widgets = self.create_selection_widgets_func(dimensions, self.labels, self.update)
        self.update()
        
    def update(self, owner_idx: int = -1) -> None:
        '''
        Redraw the form using the values of all widgets above and including the one with index owner_idx.
        If owner_idx is -1, we redraw everything.
        '''
        select_conditions = [(name, widget.value) for name, widget in self.selection_widgets.items()]
        if owner_idx == -1:
            dim_select_conditions = []
        else:
            dim_select_conditions = select_conditions[:owner_idx + 1]  # for selecting lower widget options, use value of widgets above 
        
        for name in list(self.selection_widgets.keys())[owner_idx + 1:]:
            widget = self.selection_widgets[name]
            widget_options = self.dim_filter_func(self.data, name, dim_select_conditions)
            _logger.debug(f'setting values: {widget_options} on widget: {name}')
            widget.options = self.dim_filter_func(self.data, name, dim_select_conditions)
            
        if owner_idx == -1: return
                         
        filtered_data = self.data_filter_func(self.data, select_conditions)
        transformed_data = self.transform_func(filtered_data)
        lines = self.stat_func(transformed_data, self.xcol, self.ycol, self.zcol)
        plot_widgets = self.plot_func(self.xlabel, self.ylabel, lines)
        self.display_form_func(list(self.selection_widgets.values()) + plot_widgets, self.debug)
                
        
# unit tests
class TestInteractivePlot(unittest.TestCase):
    def test_interactive_plot(self):
        np.random.seed(0)
        size = 1000
        
        dte = np.random.randint(5, 10, size)
        put_call = np.random.choice(['put', 'call'], size)
        year = np.random.choice([2018, 2019, 2020, 2021], size)
        delta = np.random.uniform(0, 0.5, size)
        delta = np.where(put_call == 'call', delta, -delta)
        premium = np.abs(delta * 10) * dte + np.random.normal(size=size) * dte / 10
        data = pd.DataFrame({'dte': dte, 'put_call': put_call, 'year': year, 'delta': delta, 'premium': premium})
        
        labels = {'premium': 'Premium $', 'year': 'Year', 'dte': 'Days to Expiry', 'delta_rnd': 'Delta'}
        secy_line_config = LineConfig(secondary_y=True)
        
        ip = InteractivePlot(data, 
                             labels, 
                             transform_func=self.transform,
                             stat_func=MeanWithCI(ci_level=95),
                             plot_func=LineGraphWithDetailDisplay(line_configs={'put': secy_line_config}), debug=True)
        
        ip.create_pivot('delta_rnd', 'premium', 'put_call', dimensions={'year': 2018, 'dte': None})
    
    def transform(self, data: pd.DataFrame) -> pd.DataFrame:
        np.seterr('raise')
        data['delta_rnd'] = percentile_buckets(np.abs(data.delta.values), 10)
        return data
    
        
if __name__ == '__main__':
    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE | doctest.ELLIPSIS)
    unittest.main(argv=['first-arg-is-ignored'], exit=False)
    print('done')
# $$_end_code
