"""Generate the documented averaged ASD figure from a gradiometer CSV."""

from __future__ import annotations

import argparse
import csv
import math
from pathlib import Path
import sys

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from asd_panel import calculate_asd


COLORS = {
    "A": "#F3A712",
    "B": "#62B6CB",
    "D": "#8FE3B0",
}


def _find_column(fieldnames: list[str], prefix: str) -> str:
    matches = [name for name in fieldnames if name.startswith(prefix)]
    if len(matches) != 1:
        raise ValueError(f"Expected one {prefix!r} column; found {matches!r}.")
    return matches[0]


def _serial_from_header(header: str) -> str:
    if "[" in header and header.endswith("]"):
        return header.rsplit("[", 1)[1][:-1]
    return "unknown"


def load_recording(
    path: Path,
) -> tuple[np.ndarray, dict[str, np.ndarray], dict[str, str]]:
    with path.open(newline="", encoding="utf-8") as handle:
        reader = csv.DictReader(handle)
        fieldnames = reader.fieldnames or []
        timestamp_column = "timestamp_ms"
        if timestamp_column not in fieldnames:
            raise ValueError("The recording has no timestamp_ms column.")
        columns = {
            "A": _find_column(fieldnames, "sensor_a_nt"),
            "B": _find_column(fieldnames, "sensor_b_nt"),
            "D": "a_minus_b_nt",
        }
        if columns["D"] not in fieldnames:
            raise ValueError("The recording has no a_minus_b_nt column.")

        timestamps: list[int] = []
        values: dict[str, list[float]] = {"A": [], "B": [], "D": []}
        for row in reader:
            timestamps.append(int(row[timestamp_column]))
            for key, column in columns.items():
                values[key].append(float(row[column]))

    if len(timestamps) < 4:
        raise ValueError("At least four exact timestamp pairs are required.")

    timestamp_array = np.asarray(timestamps, dtype=np.int64)
    deltas_ms = np.diff(timestamp_array)
    interval_ms = float(np.median(deltas_ms))
    if (
        interval_ms <= 0
        or np.any(deltas_ms <= 0)
        or np.max(np.abs(deltas_ms - interval_ms)) > interval_ms * 0.05
    ):
        raise ValueError("Sample timestamps are not uniformly spaced.")

    times = (timestamp_array - timestamp_array[0]).astype(float) / 1000.0
    series = {
        key: np.asarray(stream, dtype=float)
        for key, stream in values.items()
    }
    serials = {
        "A": _serial_from_header(columns["A"]),
        "B": _serial_from_header(columns["B"]),
    }
    return times, series, serials


def averaged_asd(
    times: np.ndarray,
    values: np.ndarray,
    point_count: int,
    overlap: float,
) -> tuple[np.ndarray, np.ndarray, int]:
    if not 0.0 <= overlap < 1.0:
        raise ValueError("Overlap must be at least 0 and less than 1.")
    if values.size < point_count:
        raise ValueError(
            f"Recording has {values.size:,} rows; {point_count:,} are required."
        )

    step = max(1, int(round(point_count * (1.0 - overlap))))
    power_sum: np.ndarray | None = None
    frequencies: np.ndarray | None = None
    count = 0
    for start in range(0, values.size - point_count + 1, step):
        stop = start + point_count
        samples = list(zip(times[start:stop], values[start:stop]))
        result = calculate_asd(samples, point_count)
        if result is None:
            continue
        window_frequencies, window_asd = result
        if frequencies is None:
            frequencies = window_frequencies
            power_sum = np.zeros_like(window_asd)
        elif not np.allclose(
            frequencies,
            window_frequencies,
            rtol=1e-9,
            atol=1e-12,
        ):
            raise ValueError("ASD windows do not share the same frequency bins.")
        power_sum += window_asd**2
        count += 1

    if frequencies is None or power_sum is None or count == 0:
        raise ValueError("No valid uniformly spaced ASD windows were found.")
    return frequencies, np.sqrt(power_sum / count), count


def create_plot(
    input_path: Path,
    output_path: Path,
    point_count: int = 4096,
    overlap: float = 0.5,
) -> None:
    times, series, serials = load_recording(input_path)
    spectra: dict[str, tuple[np.ndarray, np.ndarray]] = {}
    window_count = 0
    for key, values in series.items():
        frequencies, asd, count = averaged_asd(
            times,
            values,
            point_count,
            overlap,
        )
        spectra[key] = frequencies, asd
        window_count = count

    sample_interval = float(np.median(np.diff(times)))
    sample_rate = 1.0 / sample_interval
    duration = float(times[-1] - times[0])

    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 10,
            "axes.titleweight": "bold",
        }
    )
    figure, axis = plt.subplots(figsize=(12, 6.75), dpi=160)
    figure.patch.set_facecolor("#202328")
    axis.set_facecolor("#17191C")

    labels = {
        "A": f"Sensor A  ({serials['A']})",
        "B": f"Sensor B  ({serials['B']})",
        "D": "Gradiometer  A − B",
    }
    widths = {"A": 1.35, "B": 1.35, "D": 1.8}
    for key in ("A", "B", "D"):
        frequency, asd = spectra[key]
        axis.plot(
            frequency,
            asd,
            color=COLORS[key],
            linewidth=widths[key],
            label=labels[key],
            alpha=0.96,
        )

    axis.set_xlim(0.0, sample_rate / 2.0)
    axis.set_yscale("log")
    axis.set_xlabel("Frequency (Hz)", color="#E7E9EC", labelpad=9)
    axis.set_ylabel("Amplitude spectral density (nT/√Hz)", color="#E7E9EC")
    axis.tick_params(colors="#CDD1D6")
    for spine in axis.spines.values():
        spine.set_color("#697078")
    axis.grid(True, which="major", color="#66707A", alpha=0.30, linewidth=0.7)
    axis.grid(True, which="minor", color="#66707A", alpha=0.13, linewidth=0.45)

    title = "QTFM Sample Data — Averaged Amplitude Spectral Density"
    axis.set_title(title, color="#FFFFFF", fontsize=17, loc="left", pad=34)
    subtitle = (
        f"{sample_rate:,.0f} Hz sample rate  •  {duration:,.1f} s recording  •  "
        f"{point_count:,}-point Hann windows  •  {overlap:.0%} overlap  •  "
        f"{window_count} power averages"
    )
    axis.text(
        0.0,
        1.025,
        subtitle,
        transform=axis.transAxes,
        color="#B9C0C8",
        fontsize=10,
        va="bottom",
    )

    legend = axis.legend(
        loc="upper right",
        frameon=True,
        facecolor="#2A2E34",
        edgecolor="#697078",
        framealpha=0.94,
    )
    for text in legend.get_texts():
        text.set_color("#F4F5F6")

    diff_frequency, diff_asd = spectra["D"]
    peak_mask = diff_frequency >= 1.0
    if np.any(peak_mask):
        masked_indices = np.flatnonzero(peak_mask)
        peak_index = int(masked_indices[np.argmax(diff_asd[peak_mask])])
        peak_frequency = float(diff_frequency[peak_index])
        peak_value = float(diff_asd[peak_index])
        if math.isfinite(peak_value):
            axis.annotate(
                f"A−B peak: {peak_frequency:.1f} Hz",
                xy=(peak_frequency, peak_value),
                xytext=(min(peak_frequency + 10.0, sample_rate * 0.38), peak_value * 2.2),
                color=COLORS["D"],
                fontsize=9,
                arrowprops={
                    "arrowstyle": "->",
                    "color": COLORS["D"],
                    "lw": 1.0,
                },
            )

    figure.text(
        0.012,
        0.012,
        "A−B uses sensor packets paired by the exact QTFM millisecond timestamp.",
        color="#9EA5AD",
        fontsize=8.5,
    )
    figure.tight_layout(rect=(0.0, 0.035, 1.0, 0.98))
    output_path.parent.mkdir(parents=True, exist_ok=True)
    figure.savefig(
        output_path,
        facecolor=figure.get_facecolor(),
        bbox_inches="tight",
    )
    plt.close(figure)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--points", type=int, default=4096)
    parser.add_argument("--overlap", type=float, default=0.5)
    arguments = parser.parse_args()
    create_plot(
        arguments.input,
        arguments.output,
        arguments.points,
        arguments.overlap,
    )
    print(f"Wrote {arguments.output}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
