#!/usr/bin/env python3
"""Plot an RM3100 capture: X, Y, Z and the norm of the three.

    ./.venv/bin/python plot.py capture_60s.csv
    ./.venv/bin/python plot.py capture_60s.csv -o out.png

Small multiples rather than one shared axis: the three axes sit at very
different DC offsets, so a single scale would flatten the variation that
matters. Each panel therefore has its own y-scale -- read the panels
independently, and note the per-panel mean/sd annotation for context.
"""

import argparse
import csv
import sys

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

# Light-mode design tokens.
SURFACE = "#fcfcfb"
TEXT_PRIMARY = "#0b0b0b"
TEXT_SECONDARY = "#52514e"
GRID = "#e3e2df"

# Categorical slots 1-3 for the three peer axes. Validated all-pairs in light
# mode (worst CVD dE 9.2, normal-vision 24.0). The norm is a derived quantity,
# not 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 for small
# multiples.
SERIES = [
    ("x_uT", "X axis", "#2a78d6"),
    ("y_uT", "Y axis", "#eb6834"),
    ("z_uT", "Z axis", "#1baf7a"),
    (None,   "Norm  |B| = sqrt(X^2 + Y^2 + Z^2)", TEXT_PRIMARY),
]


def load(path):
    t, x, y, z, gains = [], [], [], [], []
    with open(path, newline="") as fh:
        for row in csv.DictReader(fh):
            t.append(float(row["elapsed_s"]))
            x.append(float(row["x_uT"]))
            y.append(float(row["y_uT"]))
            z.append(float(row["z_uT"]))
            # Recover the gain from the raw/uT ratio so the caption reports the
            # settings actually used rather than an assumption.
            for axis in "xyz":
                ut = float(row[f"{axis}_uT"])
                if abs(ut) > 1.0:
                    gains.append(float(row[f"{axis}_raw"]) / ut)
    if not t:
        sys.exit(f"{path} contains no samples")
    gain = float(np.median(gains)) if gains else float("nan")
    return np.array(t), np.array(x), np.array(y), np.array(z), gain


def rolling_mean(v, window):
    """Centred moving average that stays smooth all the way to both ends.

    A plain convolution tapers toward zero at the edges. Dividing by the number
    of samples that actually contributed gives a true partial-window mean
    instead, so the ends carry no artefact.
    """
    if window < 2:
        return v
    kernel = np.ones(window)
    total = np.convolve(v, kernel, mode="same")
    count = np.convolve(np.ones_like(v), kernel, mode="same")
    return total / count


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("csv", help="capture written by logger.py")
    ap.add_argument("-o", "--output", default=None, help="PNG path (default: <csv>.png)")
    ap.add_argument("--smooth", type=float, default=1.0,
                    help="moving-average window in seconds, 0 to disable "
                         "(default: %(default)s)")
    args = ap.parse_args()

    t, x, y, z, gain = load(args.csv)
    # Inverse of rm3100.gain_lsb_per_ut().
    cycle_count = (gain - 1.5) / 0.3671
    norm = np.sqrt(x**2 + y**2 + z**2)
    series = {"x_uT": x, "y_uT": y, "z_uT": z, None: norm}

    duration = t[-1] - t[0]
    rate = len(t) / duration if duration > 0 else float("nan")
    window = max(1, int(round(args.smooth * rate))) if args.smooth > 0 else 0

    fig, axes = plt.subplots(4, 1, figsize=(12, 9.5), sharex=True, dpi=150)
    fig.patch.set_facecolor(SURFACE)

    for ax, (key, label, color) in zip(axes, SERIES):
        v = series[key]
        ax.set_facecolor(SURFACE)

        # Raw trace kept thin and translucent: at ~376 Hz there are far more
        # samples than pixels, so a full-weight line would read as a solid band.
        ax.plot(t, v, color=color, linewidth=0.4, alpha=0.30,
                solid_capstyle="round", rasterized=True)
        if window > 1:
            ax.plot(t, rolling_mean(v, window), color=color, linewidth=1.6,
                    solid_capstyle="round")

        ax.set_ylabel("µT", color=TEXT_SECONDARY, fontsize=10)
        # Direct label instead of a legend: one series per panel, so the title
        # names it. This is also the relief the palette's contrast WARN requires.
        ax.set_title(label, color=TEXT_PRIMARY, fontsize=12, loc="left",
                     pad=8, fontweight="bold")
        ax.annotate(f"mean {v.mean():.3f}   sd {v.std() * 1000:.0f} nT   "
                    f"span {v.max() - v.min():.3f} µT",
                    xy=(1.0, 1.0), xycoords="axes fraction",
                    xytext=(0, 8), textcoords="offset points",
                    ha="right", va="bottom",
                    color=TEXT_SECONDARY, fontsize=9)

        ax.grid(True, axis="y", color=GRID, linewidth=0.8)
        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.spines[side].set_linewidth(0.8)
        ax.tick_params(colors=TEXT_SECONDARY, labelsize=9, length=0)

    axes[-1].set_xlabel("elapsed (s)", color=TEXT_SECONDARY, fontsize=10)
    axes[-1].set_xlim(t[0], t[-1])

    smooth_note = (f"; {args.smooth:g} s moving average over translucent raw trace"
                   if window > 1 else "")
    fig.suptitle("RM3100 magnetometer capture", color=TEXT_PRIMARY,
                 fontsize=15, fontweight="bold", x=0.5, y=0.985)
    fig.text(0.5, 0.955,
             f"{len(t):,} samples over {duration:.1f} s ({rate:.0f} Hz), "
             f"cycle count {cycle_count:.0f} ({gain:.1f} LSB/µT)"
             f"{smooth_note}. Panels have independent y-scales.",
             color=TEXT_SECONDARY, fontsize=10, ha="center")

    fig.tight_layout(rect=[0, 0, 1, 0.945])
    out = args.output or args.csv.rsplit(".", 1)[0] + ".png"
    fig.savefig(out, facecolor=SURFACE)
    print(f"{len(t):,} samples, {duration:.2f} s, {rate:.1f} Hz -> {out}")
    for key, label, _ in SERIES:
        v = series[key]
        print(f"  {label.split()[0]:5s} mean {v.mean():+9.3f} uT   "
              f"sd {v.std()*1000:6.1f} nT   span {v.max()-v.min():6.3f} uT")


if __name__ == "__main__":
    main()