import param
import pandas as pd
import functools
import inspect
from typing import List
import panel as pn


class DataCombiner(param.Parameterized):

    datas: List[pd.DataFrame] = param.List(
        default=[], item_type=pd.DataFrame, allow_refs=True, nested_refs=True
    )
    index_columns = param.List(default=[], item_type=str)
    transformers = param.List(default=[])
    merged_data: pd.DataFrame = param.DataFrame()

    highlight_series = param.Series(default=None)
    highlight_name = param.ObjectSelector(default=None, objects=[], allow_None=True)
    highlight_index = param.ObjectSelector(default=None, objects=[], allow_None=True)

    highlight_ranges = param.Dict(default={})

    def __init__(self, **params):
        super().__init__(**params)
        self._dynamic_range_params: set[str] = set()
        self.param.watch(self.merge, ["datas", "index_columns", "transformers"])
        
        self.param.watch(self.set_highlight_name_objects, ["merged_data"])
        self.param.watch(self.set_highlight_index_objects, ["index_columns"])
        self.param.watch(
            self.set_highlight_ranges,
            ["highlight_index", "index_columns", "merged_data"],
        )
        self.param.watch(
            self.set_highlight_series,
            ["highlight_index", "highlight_ranges", "highlight_name", "merged_data"],
        )
        self.param.watch(self.sync_highlight_range_params, ["highlight_ranges"])

        self._range_widgets = pn.Column()
        self.set_highlight_index_objects(None)
        self.view = self.make_view()

        self.merge()

    def set_highlight_series(self, event):
        if (
            self.merged_data is None
            or self.highlight_index is None
            or self.highlight_name is None
        ):
            return
        mask = pd.Series(True, index=self.merged_data.index)
        for col in self.highlight_ranges:
            rng = getattr(self, col, None)
            if rng is None:
                continue
            col_data = self.merged_data[col]
            is_dt = pd.api.types.is_datetime64_any_dtype(col_data)
            if is_dt:
                if hasattr(col_data.dtype, 'tz') and col_data.dtype.tz is not None:
                    col_data = col_data.dt.tz_localize(None)
                lo, hi = pd.Timestamp(rng[0]), pd.Timestamp(rng[1])
            else:
                lo, hi = rng[0], rng[1]
            mask = mask & (col_data >= lo) & (col_data <= hi)
        x = (
            self.merged_data.loc[mask]
            .groupby(self.highlight_index)[self.highlight_name]
            .mean()
        )

        self.param.update(highlight_series=x)

    def sync_highlight_range_params(self, event=None):
        current_keys = set(self.highlight_ranges.keys())

        for name in list(self._dynamic_range_params - current_keys):
            delattr(type(self), name)
            self._dynamic_range_params.discard(name)

        widgets = []
        for name, value in self.highlight_ranges.items():
            sanitized = tuple(None if pd.isna(v) else v for v in value)
            if not all(v is not None for v in sanitized):
                continue
            is_datetime = pd.api.types.is_datetime64_any_dtype(self.merged_data[name])
            if is_datetime:
                sanitized = tuple(v.tz_localize(None) if hasattr(v, 'tz_localize') and v.tzinfo is not None else v for v in sanitized)
            if name not in self._dynamic_range_params:
                if is_datetime:
                    p = param.DateRange(default=sanitized, bounds=sanitized)
                else:
                    p = param.Range(default=sanitized, bounds=sanitized)
                self.param.add_parameter(name, p)
                self._dynamic_range_params.add(name)
                self.param.watch(self.set_highlight_series, [name])
            else:
                self.param[name].bounds = sanitized
                self.param.update(**{name: sanitized})

            if isinstance(self.param[name], param.DateRange):
                widgets.append(pn.widgets.DateRangeSlider.from_param(
                    self.param[name], start=sanitized[0], end=sanitized[1], throttled=True
                ))
            else:
                widgets.append(pn.widgets.RangeSlider.from_param(self.param[name], throttled=True))

        self._range_widgets.objects = widgets

    def set_highlight_ranges(self, event):
        if self.merged_data is None or self.highlight_index is None:
            return
        dynamic_indexers = set(self.index_columns)
        dynamic_indexers.remove(self.highlight_index)
        _ranges = dict()
        for idxr in dynamic_indexers:
            _ranges[idxr] = (self.merged_data[idxr].min(), self.merged_data[idxr].max())
        self.param.update(highlight_ranges=_ranges)

    def set_highlight_name_objects(self, event):
        if self.merged_data is not None:
            columns = set(self.merged_data.columns)
            options = list(columns.difference(set(self.index_columns)))
            self.param["highlight_name"].objects = sorted(options)
            if self.highlight_name == None:
                self.param.update(
                    highlight_name=options[0] if len(options) > 0 else None
                )

    def set_highlight_index_objects(self, event):
        self.param["highlight_index"].objects = self.index_columns
        if self.highlight_index == None:
            self.param.update(
                highlight_index=(
                    self.index_columns[0] if len(self.index_columns) > 0 else None
                )
            )

    def merge(self, *events):
        left = self.datas[0].set_index(self.index_columns)
        for right in self.datas[1:]:
            right = right.set_index(self.index_columns)
            left = left.join(right, how="outer")
        merged_data = left.reset_index()

        for fn in self.transformers:
            unwrapped = fn.func if isinstance(fn, functools.partial) else fn
            params = list(inspect.signature(unwrapped).parameters)
            if params and params[0] == "self":
                merged_data = fn(self, merged_data)
            else:
                merged_data = fn(merged_data)

        self.param.update(merged_data=merged_data)

    def make_view(self):
        view = pn.Column(
            self.param["highlight_name"],
            # self.param["highlight_index"],
            self._range_widgets,
        )
        return view
