"""Amplitude spectral density for the two plotted sensors and their difference.

All three traces share one plot, so unlike the time series they need separate
colours: golden rod for the first sensor, QuSpin blue for the second, and green
for the difference. The colour identifies the plot slot, not the sensor letter.
"""

from __future__ import annotations

import math
from collections import deque

import numpy as np
import pyqtgraph as pg
from PyQt5.QtCore import Qt, QTimer
from PyQt5.QtWidgets import (
    QCheckBox,
    QFrame,
    QHBoxLayout,
    QLabel,
    QPushButton,
    QSpinBox,
    QVBoxLayout,
)

from sensors import SENSOR_KEYS, difference_label
from serial_connection import QTFMMeasurement
from theme import (
    COMPARE_TRACE,
    DIFFERENCE_COLOR,
    GOLDEN_ROD,
    PLOT_AXIS,
    PLOT_BACKGROUND,
    PLOT_GRID_ALPHA,
    PLOT_TEXT,
    SECONDARY_BLUE,
)
from time_series_panel import SensorTimebase

#: Stream key for the gradiometer difference of the currently plotted pair.
#: Deliberately not a single letter — sensor keys are "A".."H", so "D" would
#: collide with sensor D and mix its raw field into the difference spectrum.
DIFFERENCE_KEY = "difference"


def nearest_spectrum_value(
    frequencies: np.ndarray,
    values: np.ndarray,
    target_frequency: float,
) -> tuple[float, float] | None:
    """Return the spectrum bin nearest a frequency inside the data range."""
    if (
        frequencies.size == 0
        or frequencies.size != values.size
        or not math.isfinite(target_frequency)
        or target_frequency < frequencies[0]
        or target_frequency > frequencies[-1]
    ):
        return None

    insertion = int(np.searchsorted(frequencies, target_frequency))
    candidates = [
        index
        for index in (insertion - 1, insertion)
        if 0 <= index < frequencies.size
    ]
    if not candidates:
        return None
    nearest = min(
        candidates,
        key=lambda index: abs(frequencies[index] - target_frequency),
    )
    return float(frequencies[nearest]), float(values[nearest])


def calculate_asd(
    samples: deque[tuple[float, float]] | list[tuple[float, float]],
    point_count: int,
) -> tuple[np.ndarray, np.ndarray] | None:
    """Return a one-sided Hann-window ASD in nT/sqrt(Hz)."""
    if point_count < 4 or len(samples) < point_count:
        return None

    selected = list(samples)[-point_count:]
    times = np.asarray([sample[0] for sample in selected], dtype=float)
    values = np.asarray([sample[1] for sample in selected], dtype=float)
    if not np.all(np.isfinite(times)) or not np.all(np.isfinite(values)):
        return None

    deltas = np.diff(times)
    if deltas.size == 0 or np.any(deltas <= 0):
        return None
    sample_interval = float(np.median(deltas))
    if not math.isfinite(sample_interval) or sample_interval <= 0:
        return None

    # Reject a nonuniform packet sequence rather than showing a misleading FFT.
    tolerance = max(sample_interval * 0.05, 1e-9)
    if np.max(np.abs(deltas - sample_interval)) > tolerance:
        return None

    sample_rate_hz = 1.0 / sample_interval
    window = np.hanning(point_count)
    window_power = float(np.sum(window**2))
    if window_power <= 0:
        return None

    detrended = values - float(np.mean(values))
    spectrum = np.fft.rfft(detrended * window)
    asd = np.abs(spectrum) / math.sqrt(sample_rate_hz * window_power)

    # Convert the two-sided FFT normalization to a one-sided ASD.
    if point_count % 2 == 0:
        asd[1:-1] *= math.sqrt(2.0)
    else:
        asd[1:] *= math.sqrt(2.0)

    frequencies = np.fft.rfftfreq(point_count, sample_interval)
    usable = (frequencies > 0) & np.isfinite(asd) & (asd > 0)
    return frequencies[usable], asd[usable]


class ASDSpectrumPanel(QFrame):
    """One log-log ASD plot holding both plotted sensors and their difference."""

    MAX_POINTS = 70_000

    def __init__(self, parent=None) -> None:
        super().__init__(parent)
        self.setObjectName("plotCard")
        self._origin: float | None = None
        self._active: tuple[str, ...] = ("A", "B")
        self._reference = "A"
        self._compare = "B"
        self._series_keys = tuple(SENSOR_KEYS) + (DIFFERENCE_KEY,)
        self._streams: dict[str, deque[tuple[float, float]]] = {
            key: deque(maxlen=self.MAX_POINTS) for key in self._series_keys
        }
        self._timebases = {key: SensorTimebase() for key in SENSOR_KEYS}
        self._gradiometer_enabled = False
        self._gradiometer_reason = "Waiting for synchronized sensors."
        self._average_sums: dict[str, np.ndarray | None] = {
            key: None for key in self._series_keys
        }
        self._average_frequencies: dict[str, np.ndarray | None] = {
            key: None for key in self._series_keys
        }
        self._average_counts = {key: 0 for key in self._series_keys}
        self._last_average_marker: dict[str, tuple[int, float] | None] = {
            key: None for key in self._series_keys
        }
        self._display_spectra: dict[
            str, tuple[np.ndarray, np.ndarray] | None
        ] = {key: None for key in self._series_keys}
        self._build_ui()
        self._apply_plot_selection()

        self._update_timer = QTimer(self)
        self._update_timer.setInterval(250)
        self._update_timer.timeout.connect(self.refresh_spectra)
        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(9)
        title = QLabel("AMPLITUDE SPECTRAL DENSITY")
        title.setObjectName("plotTitle")
        toolbar.addWidget(title)
        toolbar.addStretch()

        points_label = QLabel("DATA POINTS")
        points_label.setObjectName("fieldLabel")
        toolbar.addWidget(points_label)
        self.point_count_spin = QSpinBox()
        self.point_count_spin.setRange(8, self.MAX_POINTS)
        self.point_count_spin.setSingleStep(100)
        self.point_count_spin.setValue(1000)
        self.point_count_spin.setMinimumWidth(95)
        self.point_count_spin.setToolTip(
            "Number of latest uniformly spaced samples used for each ASD."
        )
        self.point_count_spin.valueChanged.connect(
            self._on_point_count_changed
        )
        toolbar.addWidget(self.point_count_spin)

        self.log_x_checkbox = QCheckBox("LOG X")
        self.log_x_checkbox.setChecked(False)
        self.log_x_checkbox.setToolTip(
            "Use logarithmic frequency scaling. Linear X is the default."
        )
        self.log_x_checkbox.toggled.connect(self._set_log_x)
        toolbar.addWidget(self.log_x_checkbox)

        self.restart_average_button = QPushButton("RESTART AVERAGING")
        self.restart_average_button.setObjectName("deviceButton")
        self.restart_average_button.clicked.connect(self.restart_averaging)
        toolbar.addWidget(self.restart_average_button)

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

        self.cursor_readout_label = QLabel(self._cursor_placeholder())
        self.cursor_readout_label.setObjectName("cursorReadout")
        layout.addWidget(self.cursor_readout_label)

        self.status_label = QLabel("Waiting for samples.")
        self.status_label.setObjectName("asdStatus")
        self.status_label.setWordWrap(True)
        self.status_label.setMinimumHeight(30)
        layout.addWidget(self.status_label)

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

        self.plot = self.graphics.addPlot(row=0, col=0)
        self.plot.setTitle(
            "Averaged ASD — 1,000 data points per spectrum",
            color=PLOT_TEXT,
        )
        self.plot.setLabel("left", "ASD (nT/√Hz)")
        self.plot.setLabel("bottom", "Frequency (Hz)")
        for axis in ("left", "bottom"):
            self.plot.getAxis(axis).setPen(pg.mkPen(color=PLOT_AXIS, width=2))
            self.plot.getAxis(axis).setTextPen(PLOT_TEXT)
        self.plot.showGrid(x=True, y=True, alpha=PLOT_GRID_ALPHA)
        self.plot.setLogMode(x=False, y=True)
        self.plot.hideButtons()
        self.plot.setMenuEnabled(False)
        self.legend = self.plot.addLegend(offset=(12, 12))

        self.curve_a = self.plot.plot(pen=pg.mkPen(GOLDEN_ROD, width=1.7))
        self.curve_b = self.plot.plot(pen=pg.mkPen(COMPARE_TRACE, width=1.7))
        self.curve_difference = self.plot.plot(
            pen=pg.mkPen(DIFFERENCE_COLOR, width=1.9),
        )

        self.cursor_line = pg.InfiniteLine(
            angle=90,
            movable=False,
            pen=pg.mkPen(GOLDEN_ROD, width=1.2, style=Qt.DashLine),
        )
        self.cursor_line.setZValue(20)
        self.cursor_line.hide()
        self.plot.addItem(self.cursor_line, ignoreBounds=True)
        self.horizontal_cursor_line = pg.InfiniteLine(
            angle=0,
            movable=False,
            pen=pg.mkPen(SECONDARY_BLUE, width=1.2, style=Qt.DashLine),
        )
        self.horizontal_cursor_line.setZValue(20)
        self.horizontal_cursor_line.hide()
        self.plot.addItem(self.horizontal_cursor_line, ignoreBounds=True)
        self._mouse_proxy = pg.SignalProxy(
            self.graphics.scene().sigMouseMoved,
            rateLimit=40,
            slot=self._on_mouse_moved,
        )

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

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

    def set_sensor_layout(
        self,
        active: tuple[str, ...],
        reference: str,
        compare: str,
    ) -> None:
        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._reset_series(key)
                self._timebases[key].clear()

        self._active = active
        self._reference = reference
        self._compare = compare
        if pair_changed:
            self._reset_series(DIFFERENCE_KEY)
        self._apply_plot_selection()

    def _apply_plot_selection(self) -> None:
        self.legend.clear()
        self.legend.addItem(self.curve_a, f"Sensor {self._reference}")
        if self._compare != self._reference:
            self.legend.addItem(self.curve_b, f"Sensor {self._compare}")
            self.legend.addItem(
                self.curve_difference,
                difference_label(self._reference, self._compare),
            )
            self.curve_b.setVisible(True)
            self.curve_difference.setVisible(True)
        else:
            self.curve_b.setVisible(False)
            self.curve_difference.setVisible(False)
        self._reset_cursor()

    def _series_label(self, stream_key: str) -> str:
        if stream_key == DIFFERENCE_KEY:
            return difference_label(self._reference, self._compare)
        return stream_key

    def _reset_series(self, stream_key: str) -> None:
        self._streams[stream_key].clear()
        self._display_spectra[stream_key] = None
        self._reset_stream_average(stream_key)
        curve = self._curve_for(stream_key)
        if curve is not None:
            curve.clear()

    def _curve_for(self, stream_key: str) -> pg.PlotDataItem | None:
        if stream_key == DIFFERENCE_KEY:
            return self.curve_difference
        if stream_key == self._reference:
            return self.curve_a
        if stream_key == self._compare:
            return self.curve_b
        return None

    def _visible_series(self) -> list[str]:
        if self._compare == self._reference:
            return [self._reference]
        return [self._reference, self._compare, DIFFERENCE_KEY]

    # ------------------------------------------------------------------- data
    def add_measurement(
        self,
        sensor: str,
        measurement: QTFMMeasurement,
    ) -> None:
        if sensor not in self._active or 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
        value = measurement.value_nt if measurement.valid else math.nan
        self._streams[sensor].append((elapsed, value))

    def add_gradiometer_sample(
        self,
        sensor: str,
        elapsed: float,
        value_nt: float,
    ) -> None:
        if self._gradiometer_enabled and sensor == self._compare:
            self._streams[DIFFERENCE_KEY].append((elapsed, value_nt))

    def set_gradiometer_status(self, enabled: bool, reason: str) -> None:
        if enabled != self._gradiometer_enabled:
            self._reset_series(DIFFERENCE_KEY)
        self._gradiometer_enabled = enabled
        self._gradiometer_reason = reason

    def clear(self) -> None:
        self._origin = None
        for key in self._series_keys:
            self._reset_series(key)
        for timebase in self._timebases.values():
            timebase.clear()
        self._reset_cursor()
        self.status_label.setText("Waiting for samples.")

    def _on_point_count_changed(self, _value: int) -> None:
        self.restart_averaging()

    def _set_log_x(self, enabled: bool) -> None:
        self.plot.setLogMode(x=enabled, y=True)
        self.plot.enableAutoRange(axis="x", enable=True)
        self.plot.enableAutoRange(axis="y", enable=True)
        self._reset_cursor()

    def restart_averaging(self) -> None:
        """Discard accumulated spectra and begin a fresh running average."""
        for key in self._series_keys:
            self._reset_stream_average(key)
            self._display_spectra[key] = None
        for curve in (self.curve_a, self.curve_b, self.curve_difference):
            curve.clear()
        self._reset_cursor()
        self.status_label.setText("ASD averaging restarted.")
        self.refresh_spectra()

    def _reset_stream_average(self, stream_key: str) -> None:
        self._average_sums[stream_key] = None
        self._average_frequencies[stream_key] = None
        self._average_counts[stream_key] = 0
        self._last_average_marker[stream_key] = None

    def _add_to_average(
        self,
        stream_key: str,
        frequencies: np.ndarray,
        values: np.ndarray,
        marker: tuple[int, float],
    ) -> tuple[np.ndarray, np.ndarray]:
        stored_frequencies = self._average_frequencies[stream_key]
        if (
            stored_frequencies is None
            or stored_frequencies.shape != frequencies.shape
            or not np.allclose(
                stored_frequencies,
                frequencies,
                rtol=1e-8,
                atol=1e-12,
            )
        ):
            self._reset_stream_average(stream_key)
            self._average_frequencies[stream_key] = frequencies.copy()

        if marker != self._last_average_marker[stream_key]:
            power = values**2
            if self._average_sums[stream_key] is None:
                self._average_sums[stream_key] = power
            else:
                self._average_sums[stream_key] += power
            self._average_counts[stream_key] += 1
            self._last_average_marker[stream_key] = marker

        count = self._average_counts[stream_key]
        summed = self._average_sums[stream_key]
        averaged = (
            values
            if count == 0 or summed is None
            else np.sqrt(summed / count)
        )
        stored = self._average_frequencies[stream_key]
        return frequencies if stored is None else stored, averaged

    # ----------------------------------------------------------------- cursor
    def _cursor_placeholder(self) -> str:
        return "CURSOR  •  Move over the ASD plot to read every trace"

    def _on_mouse_moved(self, event: tuple[object, ...]) -> None:
        if not event:
            return
        scene_position = event[0]
        if not self.plot.sceneBoundingRect().contains(scene_position):
            self._reset_cursor()
            return

        view_position = self.plot.vb.mapSceneToView(scene_position)
        view_x = float(view_position.x())
        view_y = float(view_position.y())
        frequency = 10.0**view_x if self.log_x_checkbox.isChecked() else view_x
        cursor_asd = 10.0**view_y
        if (
            not math.isfinite(frequency)
            or frequency <= 0
            or not math.isfinite(cursor_asd)
            or cursor_asd <= 0
        ):
            self._reset_cursor()
            return

        self.cursor_line.setPos(view_x)
        self.cursor_line.show()
        self.horizontal_cursor_line.setPos(view_y)
        self.horizontal_cursor_line.show()
        self.cursor_readout_label.setText(
            self._cursor_text(frequency, cursor_asd)
        )

    def _cursor_text(
        self,
        frequency: float,
        cursor_asd: float | None = None,
    ) -> str:
        readings: dict[str, tuple[float, float] | None] = {}
        for stream_key in self._visible_series():
            spectrum = self._display_spectra[stream_key]
            readings[stream_key] = (
                None
                if spectrum is None
                else nearest_spectrum_value(
                    spectrum[0],
                    spectrum[1],
                    frequency,
                )
            )

        nearest_frequency = next(
            (
                reading[0]
                for reading in readings.values()
                if reading is not None
            ),
            frequency,
        )

        values = "  |  ".join(
            f"{self._series_label(stream_key)} "
            + ("—" if reading is None else f"{reading[1]:.4g}")
            for stream_key, reading in readings.items()
        )
        return (
            f"CURSOR  •  {nearest_frequency:.6g} Hz"
            + (
                ""
                if cursor_asd is None
                else f"  •  Y {cursor_asd:.4g} nT/√Hz"
            )
            + f"  |  ASD (nT/√Hz): {values}"
        )

    def _reset_cursor(self) -> None:
        if hasattr(self, "cursor_line"):
            self.cursor_line.hide()
        if hasattr(self, "horizontal_cursor_line"):
            self.horizontal_cursor_line.hide()
        if hasattr(self, "cursor_readout_label"):
            self.cursor_readout_label.setText(self._cursor_placeholder())

    # -------------------------------------------------------------- rendering
    def refresh_spectra(self) -> None:
        point_count = self.point_count_spin.value()
        self.plot.setTitle(
            f"Averaged ASD — {point_count:,} data points per spectrum",
            color=PLOT_TEXT,
        )

        ready_names: list[str] = []
        for stream_key in self._visible_series():
            curve = self._curve_for(stream_key)
            if curve is None:
                continue
            if stream_key == DIFFERENCE_KEY and not self._gradiometer_enabled:
                curve.clear()
                self._display_spectra[stream_key] = None
                continue
            result = calculate_asd(self._streams[stream_key], point_count)
            if result is None:
                curve.clear()
                self._display_spectra[stream_key] = None
                continue
            frequencies, values = result
            marker = (point_count, self._streams[stream_key][-1][0])
            averaged_frequencies, averaged_values = self._add_to_average(
                stream_key,
                frequencies,
                values,
                marker,
            )
            curve.setData(averaged_frequencies, averaged_values)
            self._display_spectra[stream_key] = (
                averaged_frequencies,
                averaged_values,
            )
            ready_names.append(self._series_label(stream_key))

        visible = self._visible_series()
        averages = "  •  ".join(
            f"{self._series_label(key)} {self._average_counts[key]:,}"
            for key in visible
        )
        counts = "  •  ".join(
            f"{self._series_label(key)} {len(self._streams[key]):,}"
            for key in visible
        )
        if ready_names:
            message = (
                f"Averaging {', '.join(ready_names)}"
                f"  •  Spectra averaged: {averages}"
                f"  •  Buffered: {counts}"
            )
        else:
            message = (
                f"Need {point_count:,} uniformly spaced valid samples"
                f"  •  Buffered: {counts}"
            )

        channel = difference_label(self._reference, self._compare)
        if not self._gradiometer_enabled:
            message += (
                f"  •  {channel} NOT PLOTTING — SYNC SENSORS: "
                f"{self._gradiometer_reason}"
            )
        elif self._average_counts[DIFFERENCE_KEY] == 0:
            message += (
                f"  •  {channel} has no matched packets — use SYNC SENSORS "
                "if this continues"
            )
        self.status_label.setText(message)
