#!/usr/bin/env python3
"""Characterise an RM3100 capture: noise floor, spectrum, stability.

    ./.venv/bin/python characterize.py capture_60s.csv

Produces a four-panel figure and a text summary:

  Amplitude spectral density  nT/sqrt(Hz) against the 1.2 nT/sqrt(Hz) the
                              manual quotes (Table 3-1), and against the
                              white-noise level implied by the sample sd.
  Allan deviation             where averaging stops helping and drift takes
                              over -- the honest measure of a noise floor.
  Residual distribution       after removing a slow trend, so a non-Gaussian
                              tail or quantisation shows up.
  Read latency                host-side diagnostics only -- the measurement
                              grid itself is uniform regardless.

The sample grid is exact, not assumed: the sensor samples on its own internal
schedule (manual sections 5.7.2, 5.8.2), logger.py aborts rather than skip an
interval, and capture.py refuses a non-contiguous sample index. Frequencies use
the rate measured against the host clock, not the nominal table value, which is
~6% out on this unit.
"""

import argparse
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

import capture

SURFACE = "#fcfcfb"
TEXT_PRIMARY = "#0b0b0b"
TEXT_SECONDARY = "#52514e"
GRID = "#e3e2df"
REFERENCE = "#8a8880"

# Categorical slots 1-3; three peer axes, validated all-pairs in light mode.
AXES = [("x", "X", "#2a78d6"), ("y", "Y", "#eb6834"), ("z", "Z", "#1baf7a")]

# Total field is derived from the three axes rather than a fourth peer, so it
# takes neutral ink instead of a competing hue -- which also keeps the
# categorical set at the three slots that validate.
TOTAL = ("total", "Total", TEXT_PRIMARY)
SERIES = AXES + [TOTAL]

# Table 3-1: "Noise Density @ Max. Single-Axis Sample Rate".
SPEC_ASD_NT = 1.2


def welch_asd(v, fs, nperseg=4096):
    """Amplitude spectral density in units/sqrt(Hz) via Welch's method."""
    nperseg = min(nperseg, len(v) // 4 * 2 or len(v))
    step = nperseg // 2
    window = np.hanning(nperseg)
    # Normalisation for a one-sided PSD with this window.
    scale = 1.0 / (fs * (window ** 2).sum())

    segments = []
    for start in range(0, len(v) - nperseg + 1, step):
        seg = v[start:start + nperseg]
        # Linear detrend: removes DC and any slow ramp that would smear
        # energy across the low-frequency bins.
        seg = seg - np.polyval(np.polyfit(np.arange(nperseg), seg, 1),
                               np.arange(nperseg))
        spectrum = np.abs(np.fft.rfft(seg * window)) ** 2 * scale
        spectrum[1:-1] *= 2.0          # fold negative frequencies
        segments.append(spectrum)

    psd = np.mean(segments, axis=0)
    freqs = np.fft.rfftfreq(nperseg, 1.0 / fs)
    return freqs[1:], np.sqrt(psd[1:])   # drop DC bin


def allan_deviation(v, fs, points=40):
    """Overlapping Allan deviation of the signal against averaging time tau."""
    n = len(v)
    max_m = n // 4
    ms = np.unique(np.geomspace(1, max(max_m, 2), points).astype(int))
    taus, devs = [], []
    cumulative = np.concatenate([[0.0], np.cumsum(v)])

    for m in ms:
        # Bin means of length m, taken at every offset (overlapping).
        means = (cumulative[m:] - cumulative[:-m]) / m
        diffs = means[m:] - means[:-m]
        if diffs.size < 2:
            continue
        taus.append(m / fs)
        devs.append(np.sqrt(0.5 * np.mean(diffs ** 2)))
    return np.array(taus), np.array(devs)


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("csv")
    ap.add_argument("-o", "--output", default=None)
    ap.add_argument("--start", type=float, default=0.0,
                    help="ignore samples before this elapsed time (s)")
    ap.add_argument("--end", type=float, default=None,
                    help="ignore samples after this elapsed time (s)")
    args = ap.parse_args()

    try:
        cap = capture.load(args.csv)
        if args.start or args.end is not None:
            cap = cap.restrict(args.start or None, args.end)
    except (OSError, capture.CaptureError) as exc:
        sys.exit(str(exc))

    # Total field is derived from the three axes; its noise is dominated by
    # whichever axis carries most of the field, weighted by direction cosines,
    # so it tracks that axis rather than being an independent measurement.
    data = dict(cap.axes())
    t = cap.elapsed
    # Rate measured against the host clock, not the nominal table value: a 6%
    # error would put every spectral feature 6% off.
    fs = cap.true_rate_hz
    lsb_nt = cap.lsb_nt
    # Host read latency, kept apart from the sample grid it does not affect.
    latency = np.diff(cap.system_time)  # local read spacing, not the fit residual

    # Normalising by |B| makes noise comparable across captures whose scale
    # differs -- a gain change carries the noise with it, so absolute sd alone
    # will read as a noise difference when only the scale moved.
    field = data["total"].mean()

    print(cap.summary() + "\n")
    print("series     sd    sd/|B|      p2p    sd/LSB   white-noise ASD   median ASD")

    fig, axs = plt.subplots(2, 2, figsize=(13.5, 9), dpi=150)
    fig.patch.set_facecolor(SURFACE)
    for ax in axs.flat:
        ax.set_facecolor(SURFACE)

    for key, label, color in SERIES:
        v = data[key]
        sd = v.std()
        # A flat (white) spectrum of this sd would sit at this level.
        implied = sd / np.sqrt(fs / 2)

        freqs, asd = welch_asd(v, fs)
        axs[0, 0].loglog(freqs, asd, color=color, linewidth=1.2,
                         label=label, alpha=0.85)

        taus, devs = allan_deviation(v, fs)
        axs[0, 1].loglog(taus, devs, color=color, linewidth=1.6, label=label)

        # Detrend before the histogram so slow drift does not masquerade as
        # a fat tail.
        resid = v - np.polyval(np.polyfit(t, v, 3), t)
        axs[1, 0].hist(resid, bins=120, histtype="step", linewidth=1.4,
                       color=color, label=label, density=True)

        print(f"{label:6s} {sd:8.1f} {sd/field*1e6:7.1f}ppm {v.max()-v.min():9.1f} nT "
              f"{sd/lsb_nt:7.2f}   {implied:9.2f} nT/rtHz "
              f"{np.median(asd):9.2f} nT/rtHz")

    a = axs[0, 0]
    a.axhline(SPEC_ASD_NT, color=REFERENCE, linestyle="--", linewidth=1.2)
    a.annotate(f"Table 3-1 spec {SPEC_ASD_NT} nT/√Hz", xy=(freqs[1], SPEC_ASD_NT),
               xytext=(0, 5), textcoords="offset points",
               color=REFERENCE, fontsize=9)
    a.set_title("Amplitude spectral density", loc="left",
                color=TEXT_PRIMARY, fontsize=12, fontweight="bold", pad=8)
    a.set_xlabel("frequency (Hz)"); a.set_ylabel("nT/√Hz")

    a = axs[0, 1]
    a.set_title("Allan deviation", loc="left",
                color=TEXT_PRIMARY, fontsize=12, fontweight="bold", pad=8)
    a.set_xlabel("averaging time τ (s)"); a.set_ylabel("σ (nT)")
    a.annotate("slope −½ = white noise; upturn = drift",
               xy=(0.02, 0.04), xycoords="axes fraction",
               color=TEXT_SECONDARY, fontsize=9)

    a = axs[1, 0]
    a.set_title("Residual distribution (cubic trend removed)", loc="left",
                color=TEXT_PRIMARY, fontsize=12, fontweight="bold", pad=8)
    a.set_xlabel("nT"); a.set_ylabel("density")

    a = axs[1, 1]
    a.hist(latency * 1000, bins=120, color=REFERENCE)
    a.set_yscale("log")
    a.axvline(cap.dt_true * 1000, color=TEXT_PRIMARY, linestyle="--", linewidth=1.2)
    a.annotate(f"grid {cap.dt_true * 1000:.2f} ms", xy=(cap.dt_true * 1000, 1),
               xytext=(6, 0), textcoords="offset points",
               color=TEXT_PRIMARY, fontsize=9)
    a.set_title("Host read latency", loc="left",
                color=TEXT_PRIMARY, fontsize=12, fontweight="bold", pad=8)
    a.annotate("spread is host-side only;\nthe sample grid is exact",
               xy=(0.98, 0.94), xycoords="axes fraction", ha="right", va="top",
               color=TEXT_SECONDARY, fontsize=9)
    a.set_xlabel("interval between reads (ms)"); a.set_ylabel("count")

    for ax in axs.flat:
        ax.grid(True, which="both", color=GRID, linewidth=0.7)
        ax.set_axisbelow(True)
        for side in ("top", "right"):
            ax.spines[side].set_visible(False)
        for side in ("left", "bottom"):
            ax.spines[side].set_color(GRID)
        ax.tick_params(colors=TEXT_SECONDARY, labelsize=9, length=0)
        ax.xaxis.label.set_color(TEXT_SECONDARY)
        ax.yaxis.label.set_color(TEXT_SECONDARY)
    # Explicit placement: "best" puts the ASD legend on top of the spec-line
    # annotation in the lower left.
    for ax, loc in ((axs[0, 0], "upper right"), (axs[0, 1], "upper right"),
                    (axs[1, 0], "upper right")):
        ax.legend(frameon=False, fontsize=9, labelcolor=TEXT_SECONDARY, loc=loc)

    fig.suptitle("RM3100 noise characterisation", color=TEXT_PRIMARY,
                 fontsize=15, fontweight="bold", y=0.985)
    fig.text(0.5, 0.945,
             f"{len(t):,} samples, {cap.duration:.1f} s at {fs:.2f} Hz measured "
             f"(nominal {cap.nominal_rate_hz:g} Hz, {cap.rate_error * 100:+.1f}%), "
             f"cycle count {cap.cycle_count}, 1 LSB = {lsb_nt:.2f} nT. "
             f"Sample grid is exact.",
             color=TEXT_SECONDARY, fontsize=10, ha="center")

    fig.tight_layout(rect=[0, 0, 1, 0.935])
    out = args.output or args.csv.rsplit(".", 1)[0] + "_noise.png"
    fig.savefig(out, facecolor=SURFACE)
    print(f"\n-> {out}")


if __name__ == "__main__":
    main()