import holoviews as hv
import holoviews.streams as streams
import param
import numpy as np
import panel as pn
from bokeh.models import BoxAnnotation, ColumnDataSource, CustomJS
from bokeh.models.annotations.geometry import AreaVisuals

hv.extension("bokeh")


class DimensionRangePickerGraph(param.Parameterized):

    graph = param.ClassSelector(class_=hv.DynamicMap, constant=True)
    axis = param.Selector(default="x", objects=["x", "y", "both"], constant=True)
    box_style = param.Dict(
        default=dict(
            fill_color="blue",
            fill_alpha=0.2,
            line_color="blue",
            line_alpha=0.5,
            use_handles=False,
        )
    )  # See BoxAnnotation properties: https://docs.bokeh.org/en/latest/docs/reference/models/annotations.html#bokeh.models.BoxAnnotation

    _throttle_ms = param.Integer(default=200, bounds=(0, None))
    throttled = param.Boolean(default=False)
    tap_to_clear = param.Boolean(default=False)
    default_span = param.Number(default=0.2, bounds=(0, 1))
    default_offset = param.Number(default=0.15, bounds=(0, 1))
    clamp_to_axis = param.Boolean(default=True, constant=True)

    selection = param.Tuple(default=None, length=2, allow_None=True)

    _box_left = param.Number(default=None, allow_None=True)
    _box_right = param.Number(default=None, allow_None=True)
    _box_bottom = param.Number(default=None, allow_None=True)
    _box_top = param.Number(default=None, allow_None=True)
    _last_bounds = param.Tuple(default=None, length=4, allow_None=True)

    def __init__(self, **params):
        # Convert any user-supplied _box_* values to numeric before super().__init__
        user_set_box = False
        for p in ("_box_left", "_box_right", "_box_bottom", "_box_top"):
            if p in params and params[p] is not None:
                params[p] = self._to_numeric(params[p])
                user_set_box = True

        super().__init__(**params)

        self._user_set_box = user_set_box
        self._box_ref = None
        self._source_ref = None
        self._syncing_from_python = False
        self._inside_box_hook = False
        self._x_range_limits = (None, None)
        self._y_range_limits = (None, None)
        self._from_numeric = lambda v: v  # identity; set in _box_hook for datetime axes

        if self.axis == "both":
            self.param.selection.length = 4

        # Determine tool and stream based on axis mode.
        tool_map = {"x": "xbox_select", "y": "ybox_select", "both": "box_select"}
        self._select_tool = tool_map[self.axis]

        # Add a redraw stream for forcing re-renders from Python.
        self._redraw_stream = streams.Stream.define("Redraw")()
        self._box_hook_ref = self._box_hook

        def _append_box_opts(
            element, _box_hook=self._box_hook, _tool=self._select_tool
        ):
            # Extract existing opts so we append rather than replace.
            try:
                plot_opts = hv.Store.lookup_options(
                    hv.Store.current_backend, element, "plot"
                ).kwargs
            except Exception:
                plot_opts = {}
            existing_hooks = plot_opts.get("hooks", [])
            existing_tools = plot_opts.get("tools", [])
            hooks = (
                existing_hooks
                if _box_hook in existing_hooks
                else existing_hooks + [_box_hook]
            )
            tools = (
                existing_tools if _tool in existing_tools else existing_tools + [_tool]
            )
            return element.opts(tools=tools, hooks=hooks)

        with param.edit_constant(self):
            self.graph = self.graph.clone(
                streams=self.graph.streams + [self._redraw_stream]
            ).apply(_append_box_opts)

        # Watch bounds and SingleTap on the graph (not as DynamicMap streams,
        # so the user's callback doesn't need to accept extra args).
        self._tap = streams.SingleTap(source=self.graph)

        if self.axis == "x":
            self._bounds_stream = streams.BoundsX(source=self.graph)
            self._bounds_stream.param.watch(self._on_bounds, ["boundsx"])
        elif self.axis == "y":
            self._bounds_stream = streams.BoundsY(source=self.graph)
            self._bounds_stream.param.watch(self._on_bounds, ["boundsy"])
        else:
            self._bounds_stream = streams.BoundsXY(source=self.graph)
            self._bounds_stream.param.watch(self._on_bounds, ["bounds"])

        self._tap.param.watch(self._on_tap, ["x", "y"])

        watched_params = self._active_box_params()
        self.param.watch(self._sync_box, watched_params)
        self.param.watch(self._update_selection, watched_params)
        self._update_selection()

    def _active_box_params(self):
        if self.axis == "x":
            return ["_box_left", "_box_right"]
        elif self.axis == "y":
            return ["_box_bottom", "_box_top"]
        else:
            return ["_box_left", "_box_right", "_box_bottom", "_box_top"]

    def _has_box(self):
        if self.axis == "x":
            return self._box_left is not None
        elif self.axis == "y":
            return self._box_bottom is not None
        else:
            return self._box_left is not None

    def _clear_box_params(self):
        updates = {}
        if self.axis in ("x", "both"):
            updates["_box_left"] = None
            updates["_box_right"] = None
        if self.axis in ("y", "both"):
            updates["_box_bottom"] = None
            updates["_box_top"] = None
        return updates

    def _on_bounds(self, event):
        if event.new is None:
            return
        if self.axis == "x":
            l, r = event.new
            new_bounds = (self._to_numeric(l), None, self._to_numeric(r), None)
        elif self.axis == "y":
            b, t = event.new
            new_bounds = (None, self._to_numeric(b), None, self._to_numeric(t))
        else:
            new_bounds = tuple(self._to_numeric(v) for v in event.new)

        if self.clamp_to_axis:
            l, b, r, t = new_bounds
            if self.axis in ("x", "both") and self._x_range_limits[0] is not None:
                x_lo, x_hi = sorted(self._x_range_limits)
                if l is not None:
                    l = max(l, x_lo)
                if r is not None:
                    r = min(r, x_hi)
            if self.axis in ("y", "both") and self._y_range_limits[0] is not None:
                y_lo, y_hi = sorted(self._y_range_limits)
                if b is not None:
                    b = max(b, y_lo)
                if t is not None:
                    t = min(t, y_hi)
            new_bounds = (l, b, r, t)

        if new_bounds == self._last_bounds:
            return

        updates = {"_last_bounds": new_bounds}
        if self.axis in ("x", "both"):
            updates["_box_left"] = self._clip_to_bounds("_box_left", new_bounds[0])
            updates["_box_right"] = self._clip_to_bounds("_box_right", new_bounds[2])
        if self.axis in ("y", "both"):
            updates["_box_bottom"] = self._clip_to_bounds("_box_bottom", new_bounds[1])
            updates["_box_top"] = self._clip_to_bounds("_box_top", new_bounds[3])
        self.param.update(**updates)

    def _on_tap(self, *events):
        if not self.tap_to_clear:
            return
        ev = {e.name: e.new for e in events}
        x, y = self._to_numeric(ev.get("x")), self._to_numeric(ev.get("y"))
        if not self._has_box():
            return
        inside = True
        if self.axis in ("x", "both") and x is not None:
            if not (self._box_left <= x <= self._box_right):
                inside = False
        if self.axis in ("y", "both") and y is not None:
            if not (self._box_bottom <= y <= self._box_top):
                inside = False
        if not inside:
            self.param.update(**self._clear_box_params())

    def _source_keys(self):
        keys = []
        if self.axis in ("x", "both"):
            keys += ["left", "right"]
        if self.axis in ("y", "both"):
            keys += ["bottom", "top"]
        return keys

    def _box_values(self):
        vals = {}
        if self.axis in ("x", "both"):
            vals["left"] = self._box_left
            vals["right"] = self._box_right
        if self.axis in ("y", "both"):
            vals["bottom"] = self._box_bottom
            vals["top"] = self._box_top
        return vals

    @staticmethod
    def _to_numeric(value):
        """Convert datetime64/Timestamp to float (ms since epoch), pass through others."""
        if value is None:
            return value
        if isinstance(value, np.datetime64):
            return value.astype("datetime64[ms]").astype(np.float64)
        if hasattr(value, "timestamp"):  # pd.Timestamp
            return value.timestamp() * 1000
        return value

    def _clip_to_bounds(self, param_name, value):
        if value is None:
            return value
        bounds = getattr(self.param, param_name).bounds
        if bounds is not None:
            lo, hi = bounds
            if lo is not None:
                value = max(value, lo)
            if hi is not None:
                value = min(value, hi)
        return value

    def _clamp_values(self, vals):
        clamped = dict(vals)
        if self.axis in ("x", "both") and self._x_range_limits[0] is not None:
            x_lo, x_hi = sorted(self._x_range_limits)
            if clamped.get("left") is not None:
                clamped["left"] = max(clamped["left"], x_lo)
            if clamped.get("right") is not None:
                clamped["right"] = min(clamped["right"], x_hi)
        if self.axis in ("y", "both") and self._y_range_limits[0] is not None:
            y_lo, y_hi = sorted(self._y_range_limits)
            if clamped.get("bottom") is not None:
                clamped["bottom"] = max(clamped["bottom"], y_lo)
            if clamped.get("top") is not None:
                clamped["top"] = min(clamped["top"], y_hi)
        return clamped

    def _update_selection(self, *events):
        if not self._has_box():
            self.selection = None
            return
        fn = self._from_numeric
        if self.axis == "x":
            self.selection = (fn(self._box_left), fn(self._box_right))
        elif self.axis == "y":
            self.selection = (fn(self._box_bottom), fn(self._box_top))
        else:
            self.selection = (fn(self._box_left), fn(self._box_bottom),
                              fn(self._box_right), fn(self._box_top))

    def _sync_box(self, *events):
        """Push box param changes directly to the Bokeh models."""
        if self._box_ref is None:
            if self._has_box() and not self._inside_box_hook:
                self._redraw_stream.event()
            return
        try:
            if not self._has_box():
                self._box_ref.visible = False
            else:
                self._box_ref.visible = True
                if self._source_ref is not None:
                    self._syncing_from_python = True
                    try:
                        vals = self._box_values()
                        if self.clamp_to_axis:
                            vals = self._clamp_values(vals)
                        self._source_ref.data = {k: [v] for k, v in vals.items()}
                    finally:
                        self._syncing_from_python = False
        except RuntimeError:
            # In server context, direct Bokeh model modification fails
            # inside a Bokeh event callback chain. Trigger a redraw instead.
            self._redraw_stream.event()

    def _box_hook(self, plot, element):
        # Detect axis type and set converter on first render.
        if self.axis in ("x", "both"):
            sample = plot.state.x_range.start
        else:
            sample = plot.state.y_range.start
        if isinstance(sample, np.datetime64) or hasattr(sample, "timestamp"):
            self._from_numeric = lambda v: np.datetime64(int(v), "ms") if v is not None else None
        # Always remove any existing BoxAnnotation
        plot.state.center = [
            r for r in plot.state.center if not isinstance(r, BoxAnnotation)
        ]

        # Compute axis range limits early (needed for defaults and clamping).
        if self.axis in ("x", "both"):
            x_lo, x_hi = sorted((
                self._to_numeric(plot.state.x_range.start),
                self._to_numeric(plot.state.x_range.end),
            ))
            self._x_range_limits = (x_lo, x_hi)
            if self.clamp_to_axis:
                self.param._box_left.bounds = (x_lo, x_hi)
                self.param._box_right.bounds = (x_lo, x_hi)
        if self.axis in ("y", "both"):
            y_lo, y_hi = sorted((
                self._to_numeric(plot.state.y_range.start),
                self._to_numeric(plot.state.y_range.end),
            ))
            self._y_range_limits = (y_lo, y_hi)
            if self.clamp_to_axis:
                self.param._box_bottom.bounds = (y_lo, y_hi)
                self.param._box_top.bounds = (y_lo, y_hi)

        # Auto-generate a default selection when tap_to_clear=False and no
        # box values were provided by the user.
        if not self._has_box() and not self.tap_to_clear and not self._user_set_box:
            defaults = {}
            if self.axis in ("x", "both"):
                span = self.default_span * (x_hi - x_lo)
                left = x_lo + self.default_offset * (x_hi - x_lo)
                defaults["_box_left"] = left
                defaults["_box_right"] = min(left + span, x_hi)
            if self.axis in ("y", "both"):
                span = self.default_span * (y_hi - y_lo)
                bottom = y_lo + self.default_offset * (y_hi - y_lo)
                defaults["_box_bottom"] = bottom
                defaults["_box_top"] = min(bottom + span, y_hi)
            self._inside_box_hook = True
            try:
                self.param.update(**defaults)
            finally:
                self._inside_box_hook = False

        if not self._has_box():
            self._box_ref = None
            self._source_ref = None
            return

        vals = self._box_values()
        box_kwargs = dict(editable=True, **self.box_style)
        if self.axis == "x":
            box_kwargs.update(
                left=vals["left"], right=vals["right"], movable="x", resizable="x"
            )
        elif self.axis == "y":
            box_kwargs.update(
                bottom=vals["bottom"], top=vals["top"], movable="y", resizable="y"
            )
        else:
            box_kwargs.update(
                left=vals["left"],
                right=vals["right"],
                bottom=vals["bottom"],
                top=vals["top"],
                movable="both",
                resizable="all",
            )
        box = BoxAnnotation(**box_kwargs)
        if self.clamp_to_axis:
            if self.axis in ("x", "both"):
                box.left_limit = x_lo
                box.right_limit = x_hi
            if self.axis in ("y", "both"):
                box.bottom_limit = y_lo
                box.top_limit = y_hi
        box.handles.corners = AreaVisuals(
            fill_alpha=0, line_alpha=0, hover_fill_alpha=0, hover_line_alpha=0
        )
        box.handles.sides = AreaVisuals(
            fill_color="gray",
            fill_alpha=0.3,
            line_alpha=0,
            hover_fill_color="gray",
            hover_fill_alpha=0.5,
            hover_line_alpha=0,
        )

        # Use a ColumnDataSource as a JS→Python bridge since
        # box.on_change doesn't fire for editable interactions.
        keys = self._source_keys()
        source_data = {k: [getattr(box, k)] for k in keys}
        source = ColumnDataSource(data=source_data)

        # JS: build assignment lines for sync_from_python and data emission.
        py_to_box_lines = "\n".join(f"box.{k} = source.data.{k}[0];" for k in keys)
        box_to_source = "{" + ", ".join(f"{k}: [box.{k}]" for k in keys) + "}"

        sync_from_python_js = CustomJS(
            args=dict(box=box, source=source),
            code=f"""
                source._from_python = true;
                {py_to_box_lines}
                source._from_python = false;
            """,
        )
        source.js_on_change("data", sync_from_python_js)

        if self.throttled:
            js_code = f"""
                if (source._from_python) {{ return; }}
                clearTimeout(source._timeout_id);
                source._timeout_id = setTimeout(() => {{
                    source.data = {box_to_source};
                    source.change.emit();
                }}, throttle_ms);
            """
        else:
            js_code = f"""
                if (source._from_python) {{ return; }}
                const now = Date.now();
                clearTimeout(source._timeout_id);
                if (now - (source._last_fired || 0) >= throttle_ms) {{
                    source._last_fired = now;
                    source.data = {box_to_source};
                    source.change.emit();
                }} else {{
                    source._timeout_id = setTimeout(() => {{
                        source._last_fired = Date.now();
                        source.data = {box_to_source};
                        source.change.emit();
                    }}, throttle_ms);
                }}
            """

        js_cb = CustomJS(
            args=dict(source=source, box=box, throttle_ms=self._throttle_ms),
            code=js_code,
        )
        for k in keys:
            box.js_on_change(k, js_cb)

        def on_source_change(attr, old, new):
            if self._syncing_from_python:
                return
            updates = {}
            if "left" in new:
                updates["_box_left"] = self._clip_to_bounds("_box_left", self._to_numeric(new["left"][0]))
            if "right" in new:
                updates["_box_right"] = self._clip_to_bounds("_box_right", self._to_numeric(new["right"][0]))
            if "bottom" in new:
                updates["_box_bottom"] = self._clip_to_bounds("_box_bottom", self._to_numeric(new["bottom"][0]))
            if "top" in new:
                updates["_box_top"] = self._clip_to_bounds("_box_top", self._to_numeric(new["top"][0]))
            self.param.update(**updates)

        source.on_change("data", on_source_change)

        self._box_ref = box
        self._source_ref = source

        plot.state.add_layout(box)


if __name__ == "__main__":

    def make_graph(amplitude, frequency):
        x = np.linspace(0, 4 * np.pi, 1000)
        y = amplitude * np.sin(frequency * x)
        return hv.Curve((x, y), kdims="x", vdims="y").opts(
            responsive=True,
            height=400,
            line_width=2,
            framewise=True,
        )

    amplitude = pn.widgets.FloatSlider(name="Amplitude", value=2, start=1, end=10)
    frequency = pn.widgets.FloatSlider(name="Frequency", value=1, start=1, end=6)
    dmap = hv.DynamicMap(
        make_graph,
        streams=[
            streams.Params(amplitude, ["value"], rename={"value": "amplitude"}),
            streams.Params(frequency, ["value"], rename={"value": "frequency"}),
        ],
    )

    picker = DimensionRangePickerGraph(graph=dmap, axis="x", clamp_to_axis=True, tap_to_clear=True, _box_left=5, _box_right=10)
    view = pn.Column(amplitude, frequency, picker.graph, picker.param["selection"])

    # --- Datetime x-axis example ---
    import pandas as pd

    def make_datetime_graph(amplitude, frequency):
        dates = pd.date_range("2024-01-01", periods=365, freq="D")
        days = np.arange(len(dates))
        y = amplitude * np.sin(frequency * 2 * np.pi * days / 365)
        return hv.Curve((dates, y), kdims="date", vdims="value").opts(
            responsive=True,
            height=400,
            line_width=2,
            framewise=True,
        )

    amplitude2 = pn.widgets.FloatSlider(name="Amplitude", value=2, start=1, end=10)
    frequency2 = pn.widgets.FloatSlider(name="Frequency", value=1, start=1, end=6)
    dmap2 = hv.DynamicMap(
        make_datetime_graph,
        streams=[
            streams.Params(amplitude2, ["value"], rename={"value": "amplitude"}),
            streams.Params(frequency2, ["value"], rename={"value": "frequency"}),
        ],
    )

    picker2 = DimensionRangePickerGraph(graph=dmap2, axis="both", clamp_to_axis=True, tap_to_clear=False, default_offset=0.1)
    view2 = pn.Column(amplitude2, frequency2, picker2.graph, picker2.param["selection"])
