"""Live sensor traces and the exact-timestamp gradiometer difference.

Three stacked plots are always shown: the first selected sensor, the second
selected sensor, and their difference. Any two of the connected sensors A–H
can occupy those slots; every trace is drawn in QuSpin golden rod.
"""

from __future__ import annotations

import math
from collections import deque

import pyqtgraph as pg
from PyQt5.QtCore import QTimer, pyqtSignal
from PyQt5.QtWidgets import (
    QComboBox,
    QFrame,
    QHBoxLayout,
    QLabel,
    QPushButton,
    QVBoxLayout,
)

from sensors import SENSOR_KEYS, difference_label
from serial_connection import QTFMMeasurement
from theme import (
    GOLDEN_ROD,
    PLOT_AXIS,
    PLOT_BACKGROUND,
    PLOT_GRID_ALPHA,
    PLOT_TEXT,
)


class SensorTimebase:
    """Map device milliseconds onto the app timeline with rollover handling."""

    TIMESTAMP_MODULUS = 2**32
    MAX_EXPECTED_STEP_MS = 10_000

    def __init__(self) -> None:
        self._last_timestamp_ms: int | None = None
        self._device_elapsed_ms = 0
        self._timeline_anchor = 0.0

    def clear(self) -> None:
        self._last_timestamp_ms = None
        self._device_elapsed_ms = 0
        self._timeline_anchor = 0.0

    def map_time(
        self,
        measurement: QTFMMeasurement,
        app_origin: float,
    ) -> float | None:
        host_elapsed = measurement.host_time - app_origin
        timestamp_ms = measurement.timestamp_ms
        if timestamp_ms is None:
            return None

        if self._last_timestamp_ms is None:
            self._last_timestamp_ms = timestamp_ms
            self._timeline_anchor = host_elapsed
            return self._timeline_anchor

        step_ms = (
            timestamp_ms - self._last_timestamp_ms
        ) % self.TIMESTAMP_MODULUS
        self._last_timestamp_ms = timestamp_ms

        # A sensor reboot/reset is a new clock epoch, not a 49-day jump.
        if step_ms > self.MAX_EXPECTED_STEP_MS:
            self._device_elapsed_ms = 0
            self._timeline_anchor = host_elapsed
        else:
            self._device_elapsed_ms += step_ms

        return self._timeline_anchor + self._device_elapsed_ms / 1000.0


class ExactTimestampMatcher:
    """Subtract sensors from the reference only on identical raw timestamps.

    Every gradiometer channel is ``reference − other``, so the default pair
    still produces ``A − B``.
    """

    MAX_PENDING = 4096

    def __init__(
        self,
        reference: str = "A",
        sensors: tuple[str, ...] = ("A", "B"),
    ) -> None:
        self._reference = reference
        self._sensors = tuple(sensors)
        self._pending: dict[str, dict[int, tuple[float, float]]] = {}
        self._rebuild()

    def _rebuild(self) -> None:
        self._pending = {key: {} for key in self._sensors}
        self._pending.setdefault(self._reference, {})

    @property
    def reference(self) -> str:
        return self._reference

    def configure(self, reference: str, sensors: tuple[str, ...]) -> None:
        """Set the reference sensor and the sensors it is subtracted from."""
        self._reference = reference
        self._sensors = tuple(sensors)
        self._rebuild()

    def clear(self) -> None:
        self._rebuild()

    def _others(self) -> tuple[str, ...]:
        return tuple(key for key in self._sensors if key != self._reference)

    def add(
        self,
        sensor: str,
        timestamp_ms: int,
        value_nt: float,
        elapsed: float,
    ) -> list[tuple[str, float, float]]:
        """Return the ``(other, elapsed, reference − other)`` pairs completed."""
        if sensor not in self._pending:
            return []

        if sensor == self._reference:
            self._store(sensor, timestamp_ms, value_nt, elapsed)
            completed: list[tuple[str, float, float]] = []
            for other in self._others():
                waiting = self._pending[other].pop(timestamp_ms, None)
                if waiting is not None:
                    completed.append((other, elapsed, value_nt - waiting[0]))
            return completed

        reference_sample = self._pending[self._reference].get(timestamp_ms)
        if reference_sample is not None:
            reference_value, reference_elapsed = reference_sample
            return [(sensor, reference_elapsed, reference_value - value_nt)]

        self._store(sensor, timestamp_ms, value_nt, elapsed)
        return []

    def _store(
        self,
        sensor: str,
        timestamp_ms: int,
        value_nt: float,
        elapsed: float,
    ) -> None:
        own = self._pending[sensor]
        own[timestamp_ms] = (value_nt, elapsed)
        while len(own) > self.MAX_PENDING:
            own.pop(next(iter(own)))


class TimeSeriesPanel(QFrame):
    """Two selected sensor plots plus their timestamp-paired difference."""

    gradiometer_sample = pyqtSignal(str, float, float)
    MAX_POINTS = 70_000

    def __init__(self, parent=None) -> None:
        super().__init__(parent)
        self.setObjectName("plotCard")

        self._origin: float | None = None
        self._latest_time = 0.0
        self._active: tuple[str, ...] = ("A", "B")
        self._reference = "A"
        self._compare = "B"
        self._samples: dict[str, deque[tuple[float, float]]] = {
            key: deque(maxlen=self.MAX_POINTS) for key in SENSOR_KEYS
        }
        self._difference: deque[tuple[float, float]] = deque(
            maxlen=self.MAX_POINTS
        )
        self._timestamp_matcher = ExactTimestampMatcher()
        self._gradiometer_enabled = False
        self._gradiometer_reason = "Waiting for synchronized sensors."
        self._pair_count = 0
        self._timebases = {key: SensorTimebase() for key in SENSOR_KEYS}
        self._latest_values: dict[str, float | None] = {
            key: None for key in SENSOR_KEYS
        }
        self._latest_difference: float | None = None
        self._timestamp_states = {key: "unknown" for key in SENSOR_KEYS}

        self._build_ui()
        self._apply_plot_selection()

        self._update_timer = QTimer(self)
        self._update_timer.setInterval(40)
        self._update_timer.timeout.connect(self._refresh_curves)
        self._update_timer.start()

    # ------------------------------------------------------------------ setup
    def _build_ui(self) -> None:
        layout = QVBoxLayout(self)
        layout.setContentsMargins(14, 11, 14, 12)
        layout.setSpacing(7)

        toolbar = QHBoxLayout()
        toolbar.setSpacing(10)
        title = QLabel("LIVE TIME SERIES")
        title.setObjectName("plotTitle")
        toolbar.addWidget(title)

        self.value_a_label = QLabel("A: — nT")
        self.value_a_label.setObjectName("sensorAValue")
        toolbar.addWidget(self.value_a_label)

        self.value_b_label = QLabel("B: — nT")
        self.value_b_label.setObjectName("sensorBValue")
        toolbar.addWidget(self.value_b_label)

        self.value_difference_label = QLabel("A − B: — nT")
        self.value_difference_label.setObjectName("differenceValue")
        toolbar.addWidget(self.value_difference_label)

        self.sync_controls_layout = QHBoxLayout()
        self.sync_controls_layout.setSpacing(8)
        toolbar.addLayout(self.sync_controls_layout, 1)

        window_label = QLabel("WINDOW")
        window_label.setObjectName("fieldLabel")
        toolbar.addWidget(window_label)
        self.history_combo = QComboBox()
        for label, seconds in (
            ("10 s", 10),
            ("30 s", 30),
            ("60 s", 60),
            ("300 s", 300),
        ):
            self.history_combo.addItem(label, seconds)
        self.history_combo.setMinimumWidth(90)
        toolbar.addWidget(self.history_combo)

        clear_button = QPushButton("CLEAR")
        clear_button.setObjectName("secondaryButton")
        clear_button.clicked.connect(self.clear)
        toolbar.addWidget(clear_button)
        layout.addLayout(toolbar)

        self.graphics = pg.GraphicsLayoutWidget()
        self.graphics.setBackground(PLOT_BACKGROUND)
        layout.addWidget(self.graphics, 1)

        self.plot_a, self.curve_a = self._add_plot(row=0, title="Sensor A")
        self.plot_b, self.curve_b = self._add_plot(row=1, title="Sensor B")
        self.plot_difference, self.curve_difference = self._add_plot(
            row=2,
            title="Gradiometer A − B — WAITING",
        )

        self.plot_b.setXLink(self.plot_a)
        self.plot_difference.setXLink(self.plot_a)
        self.plot_difference.setLabel("bottom", "Elapsed device time (s)")
        self.plot_a.getAxis("bottom").setStyle(showValues=False)
        self.plot_b.getAxis("bottom").setStyle(showValues=False)

    def _add_plot(
        self,
        row: int,
        title: str,
    ) -> tuple[pg.PlotItem, pg.PlotDataItem]:
        plot = self.graphics.addPlot(row=row, col=0)
        plot.setTitle(title, color=PLOT_TEXT, size="10pt")
        plot.setLabel("left", "Field (nT)")
        plot.getAxis("left").enableAutoSIPrefix(False)
        for axis in ("left", "bottom"):
            plot.getAxis(axis).setPen(pg.mkPen(color=PLOT_AXIS, width=2))
            plot.getAxis(axis).setTextPen(PLOT_TEXT)
        plot.showGrid(x=True, y=True, alpha=PLOT_GRID_ALPHA)
        plot.setMouseEnabled(x=False, y=False)
        plot.hideButtons()
        plot.setMenuEnabled(False)
        plot.setClipToView(True)
        plot.setDownsampling(auto=True, mode="peak")
        curve = plot.plot(
            pen=pg.mkPen(GOLDEN_ROD, width=1.5),
            connect="finite",
        )
        return plot, curve

    def add_sync_controls(
        self,
        timing_status: QLabel,
        sync_button: QPushButton,
    ) -> None:
        """Place timing health and the sync action in the graph header."""
        self.sync_controls_layout.addWidget(timing_status, 1)
        self.sync_controls_layout.addWidget(sync_button)

    # ------------------------------------------------------ sensor selection
    @property
    def active_sensors(self) -> tuple[str, ...]:
        return self._active

    @property
    def reference(self) -> str:
        return self._reference

    @property
    def compare(self) -> str:
        return self._compare

    @property
    def pair_count(self) -> int:
        """Matched exact-timestamp pairs since the difference last restarted."""
        return self._pair_count

    def set_sensor_layout(
        self,
        active: tuple[str, ...],
        reference: str,
        compare: str,
    ) -> None:
        """Choose the buffered sensors and the pair occupying the plots."""
        active = tuple(key for key in SENSOR_KEYS if key in active)
        if reference not in active:
            reference = active[0] if active else "A"
        if compare not in active or compare == reference:
            compare = next(
                (key for key in active if key != reference),
                reference,
            )
        if (
            active == self._active
            and reference == self._reference
            and compare == self._compare
        ):
            return

        pair_changed = (
            reference != self._reference or compare != self._compare
        )
        for key in SENSOR_KEYS:
            if key not in active:
                self._samples[key].clear()
                self._timebases[key].clear()
                self._latest_values[key] = None

        self._active = active
        self._reference = reference
        self._compare = compare
        if pair_changed:
            # The buffered difference belongs to the previous pair.
            self._difference.clear()
            self.curve_difference.clear()
            self._latest_difference = None
            self._pair_count = 0
        self._timestamp_matcher.configure(reference, (reference, compare))
        self._apply_plot_selection()

    def _apply_plot_selection(self) -> None:
        self._apply_plot_timestamp_state(self._reference, self.plot_a)
        self._apply_plot_timestamp_state(self._compare, self.plot_b)
        self._apply_gradiometer_title()
        self._update_value_labels()

    # ------------------------------------------------------------------- data
    def add_measurement(
        self,
        sensor: str,
        measurement: QTFMMeasurement,
    ) -> None:
        if sensor not in self._active:
            return
        if measurement.timestamp_ms is None:
            return
        if self._origin is None:
            self._origin = measurement.host_time

        elapsed = self._timebases[sensor].map_time(measurement, self._origin)
        if elapsed is None:
            return
        self._latest_time = max(self._latest_time, elapsed)
        value = measurement.value_nt if measurement.valid else math.nan
        self._samples[sensor].append((elapsed, value))

        if measurement.valid:
            self._latest_values[sensor] = measurement.value_nt

        if not self._gradiometer_enabled:
            return

        for other, pair_elapsed, difference in self._timestamp_matcher.add(
            sensor,
            measurement.timestamp_ms,
            value,
            elapsed,
        ):
            if other != self._compare:
                continue
            first_finite = self._latest_difference is None
            self._difference.append((pair_elapsed, difference))
            self._pair_count += 1
            if math.isfinite(difference):
                self._latest_difference = difference
            self.gradiometer_sample.emit(other, pair_elapsed, difference)
            # The title is a static banner; re-render it only on the
            # waiting → plotting transition, not once per matched packet.
            if first_finite and self._latest_difference is not None:
                self._apply_gradiometer_title()

    def set_timestamp_state(self, sensor: str, state: str) -> None:
        if sensor not in self._timestamp_states:
            return
        self._timestamp_states[sensor] = state
        if state != "available":
            self._samples[sensor].clear()
            self._timebases[sensor].clear()
            self._latest_values[sensor] = None
            if sensor == self._reference:
                self.curve_a.clear()
            elif sensor == self._compare:
                self.curve_b.clear()
        if sensor == self._reference:
            self._apply_plot_timestamp_state(sensor, self.plot_a)
        elif sensor == self._compare:
            self._apply_plot_timestamp_state(sensor, self.plot_b)

    def set_gradiometer_status(self, enabled: bool, reason: str) -> None:
        """Allow exact-packet subtraction only while timing requirements pass."""
        changed = enabled != self._gradiometer_enabled
        self._gradiometer_enabled = enabled
        self._gradiometer_reason = reason

        if changed:
            self._timestamp_matcher.clear()
            self._pair_count = 0
        if not enabled and changed:
            self._difference.clear()
            self.curve_difference.clear()
            self._latest_difference = None

        self._apply_gradiometer_title()

    # ------------------------------------------------------------- rendering
    def _apply_gradiometer_title(self) -> None:
        channel = difference_label(self._reference, self._compare)
        if self._reference == self._compare:
            title = "Gradiometer — SELECT TWO DIFFERENT SENSORS"
        elif not self._gradiometer_enabled:
            title = (
                f"Gradiometer {channel} — NOT PLOTTING: "
                f"{self._gradiometer_reason} — SYNC SENSORS"
            )
        elif self._latest_difference is None:
            title = (
                f"Gradiometer {channel} — WAITING FOR MATCHED PACKETS; "
                "SYNC SENSORS IF THIS CONTINUES"
            )
        else:
            title = f"Gradiometer {channel} — EXACT MS TIMESTAMP PAIRS"
        color = (
            PLOT_TEXT
            if self._gradiometer_enabled and self._latest_difference is not None
            else GOLDEN_ROD
        )
        self.plot_difference.setTitle(title, color=color, size="10pt")
        self.plot_difference.setToolTip(self._gradiometer_reason)

    def _apply_plot_timestamp_state(
        self,
        sensor: str,
        plot: pg.PlotItem,
    ) -> None:
        state = self._timestamp_states[sensor]
        name = f"Sensor {sensor}"

        if state == "available":
            title = name
            color = PLOT_TEXT
        elif state == "missing":
            title = f"{name} — CANNOT PLOT: MS TIMESTAMP UNAVAILABLE"
            color = GOLDEN_ROD
        elif state == "enabling":
            title = f"{name} — ENABLING MS TIMESTAMP…"
            color = GOLDEN_ROD
        else:
            title = f"{name} — WAITING FOR MS TIMESTAMP"
            color = PLOT_TEXT
        plot.setTitle(title, color=color, size="10pt")

    def clear(self) -> None:
        self._origin = None
        self._latest_time = 0.0
        for key in SENSOR_KEYS:
            self._samples[key].clear()
            self._timebases[key].clear()
            self._latest_values[key] = None
        self._difference.clear()
        self._latest_difference = None
        self._pair_count = 0
        self._timestamp_matcher.clear()
        self.curve_a.clear()
        self.curve_b.clear()
        self.curve_difference.clear()
        self._update_value_labels()

    def _refresh_curves(self) -> None:
        history = float(self.history_combo.currentData())
        cutoff = self._latest_time - history

        self._set_curve_data(
            self.curve_a,
            self._samples[self._reference],
            cutoff,
        )
        self._set_curve_data(
            self.curve_b,
            self._samples[self._compare],
            cutoff,
        )
        self._set_curve_data(self.curve_difference, self._difference, cutoff)
        self._update_value_labels()

        right = max(history, self._latest_time)
        left = max(0.0, right - history)
        self.plot_a.setXRange(left, right, padding=0.0)

    def _update_value_labels(self) -> None:
        self.value_a_label.setText(
            self._value_text(self._reference, self._latest_values[self._reference])
        )
        self.value_b_label.setText(
            self._value_text(self._compare, self._latest_values[self._compare])
        )

        channel = difference_label(self._reference, self._compare)
        if self._reference == self._compare:
            self.value_difference_label.setText("— : select two sensors")
        elif not self._gradiometer_enabled:
            self.value_difference_label.setText(
                f"{channel}: not plotting — SYNC SENSORS"
            )
        elif self._latest_difference is None:
            # The pair counter shows that gradiometer data is arriving even
            # before a plottable (finite) difference exists.
            if self._pair_count:
                self.value_difference_label.setText(
                    f"{channel}: receiving — {self._pair_count:,} "
                    "packets matched"
                )
            else:
                self.value_difference_label.setText(
                    f"{channel}: no matched packets — SYNC SENSORS"
                )
        else:
            self.value_difference_label.setText(
                f"{channel}: {self._latest_difference:,.3f} nT  "
                f"({self._pair_count:,} pairs)"
            )

    def _value_text(self, sensor: str, value: float | None) -> str:
        state = self._timestamp_states[sensor]
        if state == "missing":
            return f"{sensor}: timestamp required"
        if state == "enabling":
            return f"{sensor}: enabling timestamp…"
        return f"{sensor}: — nT" if value is None else f"{sensor}: {value:,.3f} nT"

    @staticmethod
    def _set_curve_data(
        curve: pg.PlotDataItem,
        samples: deque[tuple[float, float]],
        cutoff: float,
    ) -> None:
        visible = [(x, y) for x, y in samples if x >= cutoff]
        if not visible:
            curve.clear()
            return
        x_values, y_values = zip(*visible)
        curve.setData(x_values, y_values)
