#!/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 sys

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

import calibrate
import capture

# 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", "X axis", "#2a78d6"),
    ("y", "Y axis", "#eb6834"),
    ("z", "Z axis", "#1baf7a"),
    ("total", "Norm  |B| = sqrt(X^2 + Y^2 + Z^2)", TEXT_PRIMARY),
]


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()

    try:
        # Either format: a raw capture, or one already converted to nT by
        # calibrate.py. Which it is, is a property of the file.
        cap = calibrate.load_any(args.csv)
    except (OSError, capture.CaptureError,
            calibrate.CalibrationError) as exc:
        sys.exit(str(exc))

    # Drift-corrected time: the nominal grid would be ~6% off real seconds.
    t = cap.elapsed
    series = dict(cap.axes())
    duration = cap.duration
    rate = cap.true_rate_hz
    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("nT", 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():,.0f} nT   sd {v.std():.0f} nT   "
                    f"span {v.max() - v.min():,.0f} nT",
                    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)
    # A calibrated file has no nominal rate to be wrong against -- it has a
    # gain factor instead, which is the thing worth stating about it.
    provenance = (cap.describe() if hasattr(cap, "describe") else
                  f"nominal {cap.nominal_rate_hz:g} Hz, "
                  f"{cap.rate_error * 100:+.1f}%")
    fig.text(0.5, 0.955,
             f"{len(t):,} samples over {duration:.1f} s at {rate:.2f} Hz "
             f"({provenance}), cycle count {cap.cycle_count} "
             f"(1 LSB = {cap.lsb_nt:.1f} nT)"
             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():+12,.1f} nT   "
              f"sd {v.std():8.1f} nT   span {v.max()-v.min():10,.1f} nT")


if __name__ == "__main__":
    main()