#!/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, since the three axes sit at very different DC offsets and a
shared scale would flatten the variation that matters. Each panel has its own
y-scale, with a 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
import characterize
from characterize import GRID, SURFACE, TEXT_PRIMARY, TEXT_SECONDARY

# The three peer axes take characterize's categorical slots, so a capture keeps
# the same colours in every figure. The norm is derived from them, so it takes
# neutral ink.
SERIES = [(key, f"{name} axis", color)
          for key, name, color in characterize.AXES]
SERIES.append(("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 artifact.
    """
    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.
        cap = calibrate.load_any(args.csv)
    except (OSError, capture.CaptureError,
            calibrate.CalibrationError) as exc:
        sys.exit(str(exc))

    t = cap.elapsed        # drift-corrected; the nominal grid is ~6% off
    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)

        # Thin and translucent, there being far more samples than pixels.
        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: one series per panel.
        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 carries a gain factor, which is what is worth stating
    # about it; a raw one carries a nominal rate to compare against.
    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()