from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Dict, Optional, Tuple

import numpy as np
from PyQt5 import QtCore, QtWidgets

try:
    import pyqtgraph as pg
except Exception as exc:
    pg = None
    PYQTGRAPH_IMPORT_ERROR = exc
else:
    PYQTGRAPH_IMPORT_ERROR = None


if pg is None:
    InteractivePlotViewBox = None
else:
    class InteractivePlotViewBox(pg.ViewBox):
        def wheelEvent(self, ev, axis=None):
            modifiers = QtWidgets.QApplication.keyboardModifiers()
            target_axis = 0 if modifiers & QtCore.Qt.ControlModifier else 1
            super().wheelEvent(ev, axis=target_axis)


MAX_RENDER_POINTS = 600
REFRESH_INTERVAL_MS = 80
LEGEND_LIMIT = 8
POSITION_Y_RANGE = (-3.141592653589793, 3.141592653589793)
VELOCITY_Y_RANGE = (-5.0, 5.0)
TIME_WINDOW_OPTIONS = (
    ("1 min", 60.0),
    ("5 min", 300.0),
    ("All", None),
)


@dataclass
class MultiResolutionSeries:
    levels: list[list[tuple[float, float]]] = field(default_factory=lambda: [[]])
    carry: list[list[tuple[float, float]]] = field(default_factory=lambda: [[]])

    def clear(self) -> None:
        self.levels = [[]]
        self.carry = [[]]

    def append(self, t_value: float, y_value: float) -> None:
        self._append_level(0, (t_value, y_value))

    def sample(self, max_points: int, min_time: Optional[float] = None) -> tuple[np.ndarray, np.ndarray]:
        raw_points = self.levels[0]
        if not raw_points:
            return np.empty(0, dtype=float), np.empty(0, dtype=float)

        points = self._slice_points(raw_points, min_time)
        if not points:
            return np.empty(0, dtype=float), np.empty(0, dtype=float)

        for level_points in self.levels:
            sliced_points = self._slice_points(level_points, min_time)
            if not sliced_points:
                continue
            points = sliced_points
            if len(sliced_points) <= max_points:
                break

        if len(points) > max_points:
            step = max(1, len(points) // max_points)
            points = points[::step]

        sampled = list(points)
        first_raw = self._slice_points(raw_points, min_time)[0]
        last_raw = raw_points[-1]
        if sampled[0][0] > first_raw[0]:
            sampled.insert(0, first_raw)
        if sampled[-1][0] < last_raw[0]:
            sampled.append(last_raw)

        x_values = np.fromiter((point[0] for point in sampled), dtype=float, count=len(sampled))
        y_values = np.fromiter((point[1] for point in sampled), dtype=float, count=len(sampled))
        return x_values, y_values

    def _append_level(self, level: int, point: tuple[float, float]) -> None:
        while len(self.levels) <= level:
            self.levels.append([])
            self.carry.append([])

        self.levels[level].append(point)
        bucket = self.carry[level]
        bucket.append(point)
        if len(bucket) < 2:
            return

        point_a, point_b = bucket
        bucket.clear()
        merged = (
            (point_a[0] + point_b[0]) * 0.5,
            (point_a[1] + point_b[1]) * 0.5,
        )
        self._append_level(level + 1, merged)

    def _slice_points(
        self,
        points: list[tuple[float, float]],
        min_time: Optional[float],
    ) -> list[tuple[float, float]]:
        if not points:
            return []
        if min_time is None or min_time <= points[0][0]:
            return points

        left = 0
        right = len(points)
        while left < right:
            mid = (left + right) // 2
            if points[mid][0] < min_time:
                left = mid + 1
            else:
                right = mid

        start = max(0, left - 1)
        return points[start:]


class PlotPage(QtWidgets.QWidget):
    def __init__(self, joint_names: list[str], parent=None) -> None:
        super().__init__(parent)
        self.joint_names = joint_names
        self.pos_lines: Dict[str, object] = {}
        self.vel_lines: Dict[str, object] = {}
        self.pos_hist: Dict[str, MultiResolutionSeries] = {
            name: MultiResolutionSeries() for name in joint_names
        }
        self.vel_hist: Dict[str, MultiResolutionSeries] = {
            name: MultiResolutionSeries() for name in joint_names
        }
        self._plot_start = time.time()
        self._elapsed_time = 0.0
        self._last_plot_time = None
        self._pending_selected: set[str] = set()
        self._dirty = False
        self._last_legend_key: tuple[str, ...] = ()
        self._color_map = {}
        self._auto_follow_x = True
        self._applying_x_range = False
        self._build_ui()
        self._refresh_timer = QtCore.QTimer(self)
        self._refresh_timer.setSingleShot(True)
        self._refresh_timer.setInterval(REFRESH_INTERVAL_MS)
        self._refresh_timer.timeout.connect(self._flush_refresh)

    def _build_ui(self) -> None:
        layout = QtWidgets.QVBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(4)

        controls = QtWidgets.QHBoxLayout()
        controls.addWidget(QtWidgets.QLabel("Time Window:"))
        self.window_combo = QtWidgets.QComboBox()
        for label, seconds in TIME_WINDOW_OPTIONS:
            self.window_combo.addItem(label, seconds)
        self.window_combo.setCurrentIndex(0)
        self.window_combo.currentIndexChanged.connect(self._on_window_changed)
        controls.addWidget(self.window_combo)
        controls.addStretch(1)
        self.clear_btn = QtWidgets.QPushButton("Clear Plot")
        self.clear_btn.clicked.connect(self.reset_data)
        controls.addWidget(self.clear_btn)
        self.reset_view_btn = QtWidgets.QPushButton("Reset View")
        self.reset_view_btn.clicked.connect(self._reset_view)
        controls.addWidget(self.reset_view_btn)
        layout.addLayout(controls)

        if pg is None:
            message = QtWidgets.QLabel(f"pyqtgraph unavailable: {PYQTGRAPH_IMPORT_ERROR}")
            message.setStyleSheet("color: red; padding: 12px;")
            layout.addWidget(message)
            return

        pg.setConfigOptions(antialias=False)
        self.plot_splitter = QtWidgets.QSplitter(QtCore.Qt.Vertical)
        self.plot_splitter.setChildrenCollapsible(False)
        self.plot_splitter.setHandleWidth(8)
        layout.addWidget(self.plot_splitter, 1)

        self.pos_plot_widget = pg.PlotWidget(viewBox=InteractivePlotViewBox())
        self.pos_plot = self.pos_plot_widget.getPlotItem()
        self.pos_plot.setTitle("Position (rad)")
        self.plot_splitter.addWidget(self.pos_plot_widget)

        self.vel_plot_widget = pg.PlotWidget(viewBox=InteractivePlotViewBox())
        self.vel_plot = self.vel_plot_widget.getPlotItem()
        self.vel_plot.setTitle("Velocity (rad/s)")
        self.plot_splitter.addWidget(self.vel_plot_widget)
        self.plot_splitter.setStretchFactor(0, 1)
        self.plot_splitter.setStretchFactor(1, 1)

        self.vel_plot.setXLink(self.pos_plot)
        self.pos_plot.getViewBox().sigXRangeChanged.connect(self._on_plot_x_range_changed)
        self.vel_plot.getViewBox().sigXRangeChanged.connect(self._on_plot_x_range_changed)

        self._configure_plot(self.pos_plot, "Position", "rad", POSITION_Y_RANGE, show_bottom=False)
        self._configure_plot(self.vel_plot, "Velocity", "rad/s", VELOCITY_Y_RANGE, show_bottom=True)

        self.pos_legend = self.pos_plot.addLegend(offset=(-10, 10))
        self.vel_legend = self.vel_plot.addLegend(offset=(-10, 10))
        self.pos_legend.setVisible(False)
        self.vel_legend.setVisible(False)
        self._apply_time_window()

    def _configure_plot(self, plot, label: str, units: str, y_range: tuple[float, float], show_bottom: bool) -> None:
        plot.showGrid(x=True, y=True, alpha=0.25)
        plot.setLabel("left", label, units=units)
        plot.setMouseEnabled(x=True, y=True)
        plot.setMenuEnabled(False)
        plot.setYRange(y_range[0], y_range[1], padding=0.0)
        plot.getAxis("bottom").setStyle(showValues=show_bottom)
        if show_bottom:
            plot.setLabel("bottom", "Time", units="s")

    def update_data(
        self,
        data: Dict[str, Tuple[float, float]],
        selected: set[str],
        refresh: bool = True,
    ) -> None:
        sample_time = time.time() - self._plot_start
        if self._last_plot_time is not None and sample_time <= self._last_plot_time:
            sample_time = self._last_plot_time + 1e-6
        self._last_plot_time = sample_time
        self._elapsed_time = max(self._elapsed_time, sample_time)

        for name in selected:
            pos, vel = data.get(name, (0.0, 0.0))
            self.pos_hist[name].append(sample_time, pos)
            self.vel_hist[name].append(sample_time, vel)

        if refresh:
            self.refresh(selected)

    def refresh(self, selected: set[str]) -> None:
        self._pending_selected = set(selected)
        self._dirty = True
        if not self._refresh_timer.isActive():
            self._refresh_timer.start()

    def reset_data(self) -> None:
        for name in self.joint_names:
            self.pos_hist[name].clear()
            self.vel_hist[name].clear()
        self._plot_start = time.time()
        self._elapsed_time = 0.0
        self._last_plot_time = None
        self._pending_selected.clear()
        self._dirty = False
        self._last_legend_key = ()
        self._auto_follow_x = True

        if pg is None:
            return

        for line in self.pos_lines.values():
            line.setData([], [])
        for line in self.vel_lines.values():
            line.setData([], [])

        self._apply_time_window()
        self.pos_plot.setYRange(POSITION_Y_RANGE[0], POSITION_Y_RANGE[1], padding=0.0)
        self.vel_plot.setYRange(VELOCITY_Y_RANGE[0], VELOCITY_Y_RANGE[1], padding=0.0)
        self._sync_legends(set())

    def _flush_refresh(self) -> None:
        if not self._dirty:
            return
        self._dirty = False
        self._refresh_plot(self._pending_selected)

    def _refresh_plot(self, selected: set[str]) -> None:
        if pg is None:
            return

        min_time = self._window_min_time()

        for name in list(self.pos_lines.keys()):
            if name not in selected:
                self.pos_plot.removeItem(self.pos_lines[name])
                self.vel_plot.removeItem(self.vel_lines[name])
                del self.pos_lines[name]
                del self.vel_lines[name]

        for name in selected:
            if name not in self.pos_lines:
                color = self._color_for_joint(name)
                self.pos_lines[name] = self._create_curve(self.pos_plot, color)
                self.vel_lines[name] = self._create_curve(self.vel_plot, color)

            pos_x, pos_y = self.pos_hist[name].sample(MAX_RENDER_POINTS, min_time=min_time)
            vel_x, vel_y = self.vel_hist[name].sample(MAX_RENDER_POINTS, min_time=min_time)
            self.pos_lines[name].setData(pos_x, pos_y)
            self.vel_lines[name].setData(vel_x, vel_y)

        self._apply_time_window()

        self._sync_legends(selected)

    def _create_curve(self, plot, color):
        curve = plot.plot([], [], pen=pg.mkPen(color=color, width=1.6), antialias=False)
        curve.setClipToView(True)
        curve.setDownsampling(auto=True, method="peak")
        return curve

    def _color_for_joint(self, name: str):
        if name not in self._color_map:
            index = len(self._color_map)
            self._color_map[name] = pg.intColor(index, hues=max(8, len(self.joint_names)))
        return self._color_map[name]

    def _sync_legends(self, selected: set[str]) -> None:
        if pg is None:
            return

        selection_key = tuple(sorted(selected))
        if selection_key == self._last_legend_key:
            return

        self.pos_legend.clear()
        self.vel_legend.clear()

        show_legend = bool(selected) and len(selected) <= LEGEND_LIMIT
        self.pos_legend.setVisible(show_legend)
        self.vel_legend.setVisible(show_legend)
        if show_legend:
            for name in selection_key:
                pos_line = self.pos_lines.get(name)
                vel_line = self.vel_lines.get(name)
                if pos_line is not None:
                    self.pos_legend.addItem(pos_line, name)
                if vel_line is not None:
                    self.vel_legend.addItem(vel_line, name)
        self._last_legend_key = selection_key

    def _window_min_time(self) -> Optional[float]:
        window_seconds = self.window_combo.currentData() if hasattr(self, "window_combo") else None
        if window_seconds is None:
            return None
        return max(0.0, self._elapsed_time - float(window_seconds))

    def _apply_time_window(self) -> None:
        if pg is None or not hasattr(self, "window_combo"):
            return
        if not self._auto_follow_x:
            return

        window_seconds = self.window_combo.currentData()
        if window_seconds is None:
            min_time = 0.0
            max_time = max(1.0, self._elapsed_time)
        else:
            duration = float(window_seconds)
            min_time = max(0.0, self._elapsed_time - duration)
            max_time = duration if self._elapsed_time < duration else self._elapsed_time
            max_time = max(min_time + 1.0, max_time)

        self._set_x_range(min_time, max_time)

    def _set_x_range(self, min_time: float, max_time: float) -> None:
        self._applying_x_range = True
        try:
            self.pos_plot.setXRange(min_time, max_time, padding=0.0)
            self.vel_plot.setXRange(min_time, max_time, padding=0.0)
        finally:
            self._applying_x_range = False

    def _on_plot_x_range_changed(self, *_args) -> None:
        if self._applying_x_range:
            return
        self._auto_follow_x = False

    def _reset_view(self) -> None:
        if pg is None:
            return
        self._auto_follow_x = True
        self._apply_time_window()
        self.pos_plot.setYRange(POSITION_Y_RANGE[0], POSITION_Y_RANGE[1], padding=0.0)
        self.vel_plot.setYRange(VELOCITY_Y_RANGE[0], VELOCITY_Y_RANGE[1], padding=0.0)

    def _on_window_changed(self, _index: int) -> None:
        self._auto_follow_x = True
        active = self._pending_selected or set(self.pos_lines.keys())
        if active:
            self.refresh(active)
        else:
            self._apply_time_window()
