#!/usr/bin/env python3
"""Compare captures taken under different conditions, e.g. two supplies.

    ./.venv/bin/python compare.py a.csv b.csv c.csv ...
    ./.venv/bin/python compare.py --group note *.csv        # label from header
    ./.venv/bin/python compare.py LDO/cc100=a.csv 3V3/cc100=b.csv ...

A label may be given as `LABEL=path`, which is the only way to identify a run
whose capture was recorded without `--note`. Everything before the first `/` in
a label is the **condition**, the thing under test, and captures sharing one
are treated as the same configuration. `--supply LDO=3.0` then attaches a rail
voltage to every capture in that condition.

Three rules are built into how it reports.

**Compare |B|, not the axes.** Swapping a supply means touching the rig, and the
sensor comes back somewhere close to where it was. |B| survives a rotation, so a
per-axis difference measures the handling. The axes appear here as movement
diagnostics.

**Compare fractionally, or over a stated band.** A gain change carries the noise
with it, so a run reading 6% larger also reads ~6% noisier while being physically
identical. And a broadband figure is comparable over a band both captures cover,
so one band is chosen for the whole invocation and printed.

**A magnitude difference can be movement.** |B| is preserved under rotation, and
translation through a field gradient changes it, so a sensor that moved can shift
magnitude on its own. Two diagnostics bound that:

  rotation angle      the angle between mean field directions.
  gain residual       what is left after the best single scale factor, as a
                      fraction of |B|. Zero means one number relates the two
                      captures, which is what a gain change looks like.

Both size the doubt; the comparison runs either way. The reproducibility floor
from repeats of one condition is the honest error bar, and this prints it when
there are repeats to use.
"""

import argparse
import itertools
import math
import os
import sys
from collections import defaultdict

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

import calibrate
import capture
import characterize as ch
import rm3100

# A pure gain change scales every axis by the same factor. Allow a little for
# noise on the means before calling it movement.
RATIO_SPREAD_OK = 0.01      # 1%
ROTATION_OK_DEG = 0.5

# An axis whose mean field is near zero has a ratio dominated by noise, so the
# spread test drops it. The rotation check works on the whole vector's direction,
# where a small component contributes little, and needs no such guard.
RATIO_AXIS_MIN_FRACTION = 0.05      # of |B|

# Two captures whose rates differ by less than this are treated as the same
# rate, so a line holding a fixed fraction of fs across them proves nothing.
RATE_DISTINCT = 0.01        # 1%

# Condition and variant are separated by this in a label: `LDO/cc100`.
CONDITION_SEP = "/"

# Hue carries the condition, the thing under test, and the cycle count goes to
# line style below, so a chart keeps its colours when a run is added. Validated
# all-pairs in light mode against the #fcfcfb surface.
CONDITION_COLORS = ["#2a78d6", "#eb6834", "#1baf7a"]
# Cycle count is the secondary encoding, carried by line style and marker, so
# identity never rests on colour alone.
VARIANT_STYLES = ["-", "--", ":", "-."]
VARIANT_MARKERS = ["o", "s", "^", "D"]


# --------------------------------------------------------------------------
# Arguments
# --------------------------------------------------------------------------

def split_label(argument):
    """Split `LABEL=path` into (label, path); (None, path) when unlabelled.

    A path may itself contain '=', so a bare existing path wins over the
    labelled reading. When neither reading names a file that exists, the
    labelled one is returned, which puts the likelier path in the error message.
    """
    label, separator, path = argument.partition("=")
    if not separator:
        return None, argument
    if not label:
        raise ValueError(f"{argument!r} has an empty label before '='")
    if not os.path.exists(path) and os.path.exists(argument):
        return None, argument
    return label, path


def parse_supply(argument):
    """Parse `CONDITION=VOLTS` for --supply."""
    condition, separator, volts = argument.partition("=")
    if not separator or not condition:
        raise argparse.ArgumentTypeError(
            f"--supply {argument!r} must be CONDITION=VOLTS, e.g. LDO=3.0")
    try:
        value = float(volts)
    except ValueError:
        raise argparse.ArgumentTypeError(
            f"--supply {argument!r}: {volts!r} is not a number")
    if not 0 < value < 100:
        raise argparse.ArgumentTypeError(
            f"--supply {argument!r}: {value} V is not a plausible rail")
    return condition, value


def parse_band(text):
    """Parse `LO,HI` for --band."""
    parts = text.split(",")
    if len(parts) != 2:
        raise argparse.ArgumentTypeError(f"--band {text!r} must be LO,HI in Hz")
    try:
        lo, hi = (float(p) for p in parts)
    except ValueError:
        raise argparse.ArgumentTypeError(f"--band {text!r} must be two numbers")
    if not 0 < lo < hi:
        raise argparse.ArgumentTypeError(
            f"--band {text!r} needs 0 < LO < HI")
    return lo, hi


def label_for(cap, path, group_by, explicit):
    if explicit:
        return explicit
    if group_by == "note":
        return cap.meta.get("note", "(no note)")
    return os.path.basename(path)


def condition_of(label):
    """The part of a label before the first '/': the thing under test."""
    return label.split(CONDITION_SEP, 1)[0]


def variant_of(label):
    """The part after the first '/', or '' when the label carries none."""
    _, separator, variant = label.partition(CONDITION_SEP)
    return variant if separator else ""


def conditions_in_order(records):
    """Distinct conditions, in the order they were given on the command line.

    Command-line order is the one thing the caller controls, so contrasts read
    in the direction they asked for and a colour stays put across runs.
    """
    seen = []
    for label, _, _ in records:
        condition = condition_of(label)
        if condition not in seen:
            seen.append(condition)
    return seen


# --------------------------------------------------------------------------
# Per-capture summary
# --------------------------------------------------------------------------

def summarise(cap, band):
    d = dict(cap.axes())
    mean = np.array([d["x"].mean(), d["y"].mean(), d["z"].mean()])
    field = float(np.linalg.norm(mean))
    fs = cap.true_rate_hz
    out = {"cap": cap, "mean": mean, "field": field, "axes": d, "fs": fs,
           "cycle_count": cap.cycle_count}
    for key, _, _ in ch.SERIES:
        v = d[key]
        median_asd, band_rms, peak, peak_hz = ch.band_stats(v, fs, band)
        # The absolute figures are nanotesla, and nanotesla are what a gain
        # change moves, so every cross-condition comparison uses the columns
        # divided by that condition's own |B|.
        out[key] = {
            "sd": v.std(),
            "white": ch.white_sd(v),
            "ppm": v.std() / field * 1e6,
            "asd": median_asd,
            "asd_ppm": median_asd / field * 1e6,
            "white_ppm": ch.white_sd(v) / field * 1e6,
            "band_rms": band_rms,
            "band_rms_ppm": band_rms / field * 1e6,
            "peak": peak,
            "peak_hz": peak_hz,
        }
    taus, devs = ch.allan_deviation(d["total"], fs)
    i = int(np.argmin(devs))
    out["allan"] = (float(devs[i]), float(taus[i]))
    out["allan_curve"] = (taus, devs)
    out["lines"] = ch.sample_locked_lines(d["total"])
    return out


# --------------------------------------------------------------------------
# Cross-capture analysis
# --------------------------------------------------------------------------

def axis_ratio_spread(a, b, ratio):
    """Spread of the per-axis scale ratios, or None if it cannot be measured.

    A pure gain change scales X, Y and Z identically, so the spread is the
    movement test. Dividing two near-zero means amplifies their noise without
    bound, so axes below RATIO_AXIS_MIN_FRACTION of |B| in *either* capture are
    dropped, and with fewer than two left the test abstains.
    """
    usable = [i for i in range(3)
              if abs(a["mean"][i]) >= RATIO_AXIS_MIN_FRACTION * a["field"]
              and abs(b["mean"][i]) >= RATIO_AXIS_MIN_FRACTION * b["field"]]
    if len(usable) < 2:
        return None
    kept = ratio[usable]
    return float(np.ptp(kept) / np.abs(np.mean(kept)))


def gain_and_movement(a, b):
    """(scale, residual, rotation deg) relating two mean field vectors.

    `scale` is the single factor that best maps a onto b, and `residual` is what
    that factor leaves over, as a fraction of |b|. A pure gain change leaves a
    residual of zero at any angle; a rotation leaves a large one. It stays finite
    when an axis passes through zero, which is when a re-oriented sensor most
    needs the diagnostic.
    """
    ma, mb = a["mean"], b["mean"]
    scale = float(ma @ mb / (ma @ ma))
    residual = float(np.linalg.norm(mb - scale * ma) / np.linalg.norm(mb))
    ua, ub = ma / a["field"], mb / b["field"]
    angle = float(np.degrees(np.arccos(np.clip(ua @ ub, -1, 1))))
    return scale, residual, angle


def fit_rate_model(caps):
    """Solve rm3100's timing model for a set of captures: (counts/s, overhead).

    The model is `period = AXES * (cycle_count / C + h)`, so period/AXES is
    linear in cycle count with slope 1/C and intercept h. Two cycle counts
    determine both exactly; more are a least-squares fit.

    Fitting it per condition separates an oscillator that changed from a model
    that was wrong to begin with: a supply that only scales the clock moves C and
    h together by one factor, leaving the *shape* of the model alone.
    """
    try:
        # The uncertainty on each period is how far the oscillator drifted
        # during that capture; the worst of them bounds the fit.
        drift = max(ch.rate_stability(c) for c in caps)
        return calibrate.fit_timing([c.cycle_count for c in caps],
                                    [c.dt_true for c in caps],
                                    period_sd=drift or 1e-9)
    except calibrate.CalibrationError as exc:
        # Translated because callers here already handle ValueError. The fit
        # itself stays in calibrate.py, shared with sweep.py.
        raise ValueError(str(exc)) from exc


def decimation_paths(fast, slow, k, band):
    """Noise of a fast capture decimated by k, against a natively slow one.

    Answers whether sampling fast and filtering down is equivalent to having
    sampled slowly, which is the basis for recommending a low cycle count. Both
    decimations are shown: the boxcar is what the sensor's own integration does,
    poor stopband and all, and the FIR rejects what the boxcar folds in.

    Compare the ASD column, which is quoted over a common band. Each path has
    its own effective noise bandwidth, which the per-sample sd carries with it.
    """
    v_fast = fast["axes"]["total"]
    v_slow = slow["axes"]["total"]
    fs_fast, fs_slow = fast["fs"], slow["fs"]
    rows = [("raw fast", v_fast, fs_fast, fast["field"])]
    for method in ("boxcar", "fir"):
        rows.append((f"{method} /{k}", ch.decimate(v_fast, k, method),
                     fs_fast / k, fast["field"]))
    rows.append(("native slow", v_slow, fs_slow, slow["field"]))
    out = []
    for name, v, fs, field in rows:
        median_asd, band_rms, _, _ = ch.band_stats(v, fs, band)
        # The two cycle counts have their own gains, so the fractional column
        # is the comparable one here too.
        out.append({"name": name, "fs": fs, "asd": median_asd,
                    "asd_ppm": median_asd / field * 1e6,
                    "band_rms": band_rms, "white": ch.white_sd(v),
                    "lines": ch.sample_locked_lines(v)})
    return out


def decimation_pairs(entries):
    """[(condition, fast, slow, k)] for every pair a decimation can bridge.

    `entries` are (label, summary-or-capture-like) with .cycle_count and a rate.
    A pair qualifies when one cycle count is an integer multiple of the other,
    which is when averaging k samples of the fast run buys the same integration
    time as one sample of the slow one.
    """
    by_condition = defaultdict(list)
    for label, item in entries:
        by_condition[condition_of(label)].append((label, item))
    pairs = []
    for condition, group in by_condition.items():
        for (la, a), (lb, b) in itertools.combinations(group, 2):
            fast, slow, lf, ls = ((a, b, la, lb) if a["fs"] > b["fs"]
                                  else (b, a, lb, la))
            if not fast["cycle_count"] or slow["cycle_count"] % fast["cycle_count"]:
                continue
            k = slow["cycle_count"] // fast["cycle_count"]
            if k >= 2:
                pairs.append((condition, (lf, fast), (ls, slow), k))
    return pairs


def comparable_rates(loaded):
    """Every rate that will be scored, decimated paths included.

    The band has to fit inside the narrowest of these. Decimating by k lands at
    fs/k, and the per-axis overhead puts that a little *below* the natively-slow
    rate; two percent of bandwidth is enough to score the decimated path on its
    own filter rolloff.

    Runs before the summaries exist, so it takes captures.
    """
    entries = [(label, {"fs": cap.true_rate_hz,
                        "cycle_count": cap.cycle_count})
               for label, cap in loaded]
    rates = [item["fs"] for _, item in entries]
    rates += [fast["fs"] / k for _, (_, fast), _, k in decimation_pairs(entries)]
    return rates


def locked_line_verdict(records):
    """Which rational-fraction lines hold their fraction across rates.

    A field line sits at a fixed frequency, so it lands on a different fraction
    of fs when the rate changes. One that keeps the same fraction is locked to
    the sampling, and so is an artifact of measuring. Deciding this needs two
    captures at genuinely different rates, which is why it lives here.
    """
    seen = defaultdict(list)
    for label, _, s in records:
        for line in s["lines"]:
            seen[(line.numerator, line.period)].append((label, s["fs"], line))
    verdicts = []
    for fraction, hits in sorted(seen.items(),
                                 key=lambda kv: -max(h[2].sigma for h in kv[1])):
        rates = [fs for _, fs, _ in hits]
        distinct = (len(rates) > 1
                    and (max(rates) / min(rates) - 1) > RATE_DISTINCT)
        verdicts.append((fraction, hits, distinct))
    return verdicts


# --------------------------------------------------------------------------
# Output sections
# --------------------------------------------------------------------------

def print_overview(records, band, supplies):
    print(f"{'capture':<20} {'cc':>5} {'fs Hz':>9} {'drift':>7} {'|B| nT':>10} "
          f"{'sd':>7} {'white':>7} {'ASD':>6} {'Allan':>6} | {'white':>7} "
          f"{'ASD':>7}")
    for label, _, s in records:
        t = s["total"]
        print(f"{label[:20]:<20} {s['cycle_count']:5d} {s['fs']:9.3f} "
              f"{ch.rate_stability(s['cap']) * 1e6:6.0f}p "
              f"{s['field']:10,.0f} {t['sd']:7.1f} {t['white']:7.2f} "
              f"{t['asd']:6.2f} {s['allan'][0]:6.2f} | "
              f"{t['white_ppm']:7.1f} {t['asd_ppm']:7.1f}")
    print(f"  Left of the bar, absolute: sd/white/Allan in nT, ASD in nT/rtHz "
          f"over {band[0]:g}-{band[1]:.1f} Hz.")
    print("  Right of the bar, the same two divided by that capture's own |B|: "
          "ppm and ppm/rtHz.")
    print("  Compare conditions on the right, where each capture is divided by "
          "its own |B|, so a\n  gain difference stays out of the noise figure. "
          "Use the left against Table 3-1,\n  which is quoted in nT at one "
          "supply.")
    print("  white = sd(diff)/sqrt(2), which rejects drift. "
          "Allan = best sigma by averaging.")
    print("  drift is how far the oscillator moved across the run, in ppm: the "
          "error bar on that\n  capture's rate, and on anything derived from it. "
          "A tenth of a percent is ordinary\n  for an LR relaxation oscillator. "
          "It is a timing figure.")
    if supplies:
        rails = "  ".join(f"{c} {v:g} V" for c, v in sorted(supplies.items()))
        print(f"  rails: {rails}")


def print_rate_model(records):
    by_condition = defaultdict(list)
    for label, _, s in records:
        by_condition[condition_of(label)].append(s["cap"])
    fits = {}
    for condition in conditions_in_order(records):
        try:
            fits[condition] = fit_rate_model(by_condition[condition])
        except ValueError:
            continue
    if not fits:
        return
    print("\n=== timing model: period = 3 x (cycle_count / C + overhead) ===")
    for condition, fit in fits.items():
        print(f"  {condition:<12} C {fit.counts_per_second:10,.0f} "
              f"+/- {fit.counts_per_second_sd:5,.0f} counts/s   overhead "
              f"{fit.axis_overhead_s * 1e6:6.2f} +/- "
              f"{fit.axis_overhead_s_sd * 1e6:4.2f} us   "
              f"({fit.points} points, "
              + ("exactly determined, so no residual)"
                 if fit.points == 2 else
                 f"residual {fit.residual_ppm:.0f} ppm)"))
    print(f"  {'rm3100.py':<12} C {rm3100.COUNTS_PER_SECOND:10,.0f} counts/s   "
          f"overhead {rm3100.AXIS_OVERHEAD_S * 1e6:6.2f} us   (the model's own)")
    if len(fits) < 2:
        return
    print("\n  Everything the chip times comes off one oscillator, so a "
          "condition that only")
    print("  changes its frequency scales C up and the overhead down by the "
          "same factor:")
    for (ca, fa), (cb, fb) in itertools.combinations(fits.items(), 2):
        count_ratio = fb.counts_per_second / fa.counts_per_second
        overhead_ratio = fb.axis_overhead_s / fa.axis_overhead_s
        print(f"    {ca} -> {cb}   C x{count_ratio:.5f} "
              f"({(count_ratio - 1) * 100:+.2f}%)   "
              f"overhead x{overhead_ratio:.5f} "
              f"({(overhead_ratio - 1) * 100:+.2f}%)")
        print(f"      pure clock scaling would put the overhead at "
              f"x{1 / count_ratio:.5f}; it is "
              f"{(overhead_ratio * count_ratio - 1) * 100:+.2f}% off that.")


def print_locked_lines(records):
    verdicts = locked_line_verdict(records)
    if not verdicts:
        return
    print("\n=== lines at simple fractions of the sample rate (|B|) ===")
    for (numerator, period), hits, distinct in verdicts:
        print(f"  {numerator}/{period} of fs:")
        for label, fs, line in hits:
            print(f"    {label[:24]:<24} {line.numerator/line.period*fs:8.3f} Hz"
                  f"   {line.amplitude:6.2f} nT   {line.sigma:5.1f} sigma")
        if distinct:
            print("    -> SAMPLE-LOCKED. It holds this fraction of fs across "
                  "captures whose rates\n       differ, so it moves with the "
                  "sampling. Decimating a faster capture\n       removes it.")
        elif len(hits) > 1:
            print("    -> seen in several captures, at rates too close together "
                  "to tell a\n       sample-locked line from a fixed one. "
                  "Re-run with rates further apart.")
        else:
            print("    -> seen in one capture only. Deciding whether it is "
                  "locked to the sampling\n       or to a frequency needs a "
                  "second capture at another rate.")


def print_decimation(records, band):
    pairs = decimation_pairs([(label, s) for label, _, s in records])
    if not pairs:
        return
    print("\n=== does filtering and decimating the fast capture match the "
          "slow one? ===")
    for condition, (lf, fast), (ls, slow), k in pairs:
        print(f"\n  {condition}: {lf} decimated by {k} vs {ls}")
        try:
            rows = decimation_paths(fast, slow, k, band)
        except ValueError as exc:
            print(f"    cannot decimate: {exc}")
            continue
        print(f"    {'path':<14} {'fs Hz':>8} {'ASD':>7} {'ASD ppm':>8} "
              f"{'white sd':>9}  sample-locked lines")
        for row in rows:
            lines = (", ".join(f"{l.numerator}/{l.period} {l.amplitude:.2f} nT"
                               for l in row["lines"]) or "none")
            print(f"    {row['name']:<14} {row['fs']:8.2f} {row['asd']:7.2f} "
                  f"{row['asd_ppm']:8.1f} {row['white']:9.2f}  {lines}")
        raw = rows[0]
        fir = next(r for r in rows if r["name"].startswith("fir"))
        native = rows[-1]
        # Two ratios. Decimating leaves the gain and the sensor position
        # untouched, so the first measures the filter alone; the second brings
        # in a separate run, carrying every difference between the two.
        print(f"    -> decimating this capture changes its own floor by "
              f"{(fir['asd_ppm'] / raw['asd_ppm'] - 1) * 100:+.1f}%, "
              f"this capture against itself.")
        print(f"       Against the separately-recorded slow capture it is "
              f"{fir['asd_ppm'] / native['asd_ppm']:.3f}x, which also carries "
              f"whatever\n       differs between two runs.")
    print("\n  Compare the ASD ppm column, which is quoted over a common band "
          "and divided by\n  each capture's own |B|. Each path carries its own "
          "effective noise bandwidth in\n  the per-sample sd, and its own gain "
          "in the absolute nT.")


# The candidate explanations for a gain that moves with the supply, in the
# order reported. `both` stacks the two, and its residual shows what that costs:
# the rate change *is* the supply change seen through the oscillator, so the two
# together count the same physics twice.
CORRECTIONS = [("raw", "raw"),
               ("rate", "x measured fs"),
               ("supply", "x rail volts"),
               ("both", "x fs x volts")]


def correction_names(cal):
    """Row labels, naming what the rate row is actually referenced to.

    Uncalibrated it is `1 + rate_error`, measured against rm3100's *nominal*
    model, whose misfit is cycle-count dependent, so the row carries meaning
    inside a matched cycle-count contrast, where that misfit cancels. With a
    calibration it is the measured oscillator against a stated reference, which
    means the same thing everywhere.
    """
    renamed = {"rate": "x oscillator", "both": "x osc x volts"}
    return [(key, renamed.get(key, name) if cal else name)
            for key, name in CORRECTIONS]


def corrected_fields(s, supplies, reference_volts, cal=None):
    """|B| under each candidate correction, keyed by name; None if unavailable.

    rate    the chip's own clock is measured every run, so a gain that tracks
            the clock is correctable from the capture alone. With a calibration
            this is the measured oscillator against a stated reference; without
            one it falls back to `1 + rate_error`, which carries meaning inside
            a matched cycle-count contrast.
    supply  a gain simply ratiometric with the rail. Needs the rail measured,
            and nothing from the capture.
    both    the two together, which holds if they are independent mechanisms.
    """
    raw = s["field"]
    if cal is not None:
        by_rate = raw * calibrate.gain_factor(s["cap"], cal).value
    else:
        by_rate = raw * (1.0 + s["cap"].rate_error)
    volts = supplies.get(condition_of(s["label"]))
    ratio = (volts / reference_volts) if volts else None
    return {"raw": raw,
            "rate": by_rate,
            "supply": raw * ratio if ratio else None,
            "both": by_rate * ratio if ratio else None}


def print_contrasts(records, supplies, cal=None):
    """Condition contrasts at matched cycle count: the designed comparison."""
    conditions = conditions_in_order(records)
    if len(conditions) < 2:
        return
    reference_volts = min(supplies.values()) if supplies else None
    by_cc = defaultdict(dict)
    for label, _, s in records:
        s["label"] = label
        by_cc[s["cycle_count"]][condition_of(label)] = s

    print("\n=== |B| by condition, at matched cycle count ===")
    print("The comparison the interleaving was for: same cycle count, "
          "so only the condition differs.")
    contrasts = defaultdict(list)
    exponents = defaultdict(list)
    for cc in sorted(by_cc):
        present = by_cc[cc]
        if len(present) < 2:
            continue
        for ca, cb in itertools.combinations(conditions, 2):
            if ca not in present or cb not in present:
                continue
            a, b = present[ca], present[cb]
            left = corrected_fields(a, supplies, reference_volts, cal)
            right = corrected_fields(b, supplies, reference_volts, cal)
            print(f"\n  cycle count {cc}: {ca} -> {cb}")
            for key, name in correction_names(cal):
                if left[key] is None or right[key] is None:
                    continue
                print(f"    {name:<14} {left[key]:10,.1f} -> "
                      f"{right[key]:10,.1f} nT   "
                      f"{right[key] - left[key]:+9,.1f} nT   "
                      f"{(right[key]/left[key] - 1) * 100:+6.2f}%")
            contrasts[(ca, cb)].append((cc, math.log(right["raw"] / left["raw"])))
            # Fit the power gain actually goes as. -1 is exactly ratiometric,
            # 0 is no supply dependence. The rate gets the same treatment: same
            # oscillator, seen from the other end.
            va, vb = supplies.get(ca), supplies.get(cb)
            if va and vb and va != vb:
                dv = math.log(vb / va)
                exponents[(ca, cb)].append(
                    (cc, math.log(right["raw"] / left["raw"]) / dv,
                     math.log(b["fs"] / a["fs"]) / dv))

    if cal is not None:
        print("\n  The x oscillator rows above hold for any "
              "gain_offset_counts, including the\n  "
              f"{cal.gain_offset_counts:+.2f} this calibration carries. The "
              "shape term depends only on the cycle\n  count, so at a matched "
              "cycle count both conditions carry the same one and it\n  "
              "cancels. A contrast spanning cycle counts is what measures the "
              "offset.")

    for (ca, cb), points in exponents.items():
        va, vb = supplies[ca], supplies[cb]
        print(f"\n  how {ca} -> {cb} ({va:g} V -> {vb:g} V) scales with the rail:")
        for cc, field_power, rate_power in points:
            print(f"    cycle count {cc}: |B| ~ V^{field_power:+.3f}   "
                  f"rate ~ V^{rate_power:+.3f}")
        print("    V^-1 on |B| is exactly ratiometric, gain proportional to "
              "the rail; V^0 is\n    no supply dependence. Read the residual "
              "of the `x rail volts` row above\n    against the "
              "reproducibility figure below to see what these support.")

    for (ca, cb), points in contrasts.items():
        if len(points) < 2:
            continue
        values = [v for _, v in points]
        print(f"\n  {ca} -> {cb} across {len(points)} cycle counts: "
              f"mean {np.mean(values) * 100:+.2f}%, "
              f"spread {(max(values) - min(values)) * 100:.2f}%")
        print("    A small spread means the condition acts the same at every "
              "cycle count, which\n    is the signature of a scale factor. An "
              "environmental step spreads wider.")

    # A condition measured twice bounds everything this experiment cannot
    # separate, handling and gradient and drift, and is the honest error bar.
    print("\n  reproducibility within one condition (the error bar on all of "
          "the above):")
    any_repeat = False
    for condition in conditions:
        runs = [(s["cycle_count"], s["field"]) for l, _, s in records
                if condition_of(l) == condition]
        if len(runs) < 2:
            continue
        any_repeat = True
        fields = [f for _, f in runs]
        print(f"    {condition:<12} n={len(runs)}  |B| "
              + ", ".join(f"{f:,.0f}" for f in fields)
              + f"   spread {(max(fields)/min(fields) - 1) * 100:.2f}%")
    if any_repeat:
        print("    These repeats differ in cycle count as well, so the spread "
              "is an upper bound\n    on run-to-run reproducibility.")
    else:
        print("    none: no condition was captured twice, so there is no error "
              "bar at all.")


def print_pairwise(records):
    if len(records) < 2:
        return
    print("\n=== pairwise: |B|, and how far the sensor moved ===")
    print("|B| survives a rotation, so a difference is movement where the "
          "sensor also\ntranslated through a gradient. These size that doubt.")
    for (la, _, a), (lb, _, b) in itertools.combinations(records, 2):
        ratio = b["mean"] / a["mean"]
        spread = axis_ratio_spread(a, b, ratio)
        scale, residual, angle = gain_and_movement(a, b)
        print(f"\n{la[:24]} -> {lb[:24]}")
        print(f"   |B|              {a['field']:,.0f} -> {b['field']:,.0f} nT   "
              f"{b['field'] - a['field']:+,.1f} nT   "
              f"({(b['field']/a['field'] - 1) * 100:+.2f}%)")
        print(f"   fractional noise {a['total']['ppm']:.1f} -> "
              f"{b['total']['ppm']:.1f} ppm  "
              f"({(b['total']['ppm']/a['total']['ppm'] - 1) * 100:+.1f}%)")
        print(f"   best-fit scale   {scale:.5f}   gain residual "
              f"{residual * 100:5.2f}%   rotation {angle:.3f} deg")
        if spread is not None:
            print(f"   per-axis ratios  X {ratio[0]:.5f}  Y {ratio[1]:.5f}  "
                  f"Z {ratio[2]:.5f}   spread {spread * 100:.2f}%")
        else:
            print(f"   per-axis ratios  fewer than two axes carry "
                  f"{RATIO_AXIS_MIN_FRACTION:.0%} of |B|, so this "
                  "test abstains")
        moved = angle > ROTATION_OK_DEG or residual > RATIO_SPREAD_OK
        if moved:
            print("   -> the sensor moved between these two, so some of the "
                  "|B| difference may be\n      a different place in the "
                  "field.")
        else:
            print("   -> the sensor held still, so the |B| difference is a "
                  "scale change.")


# --------------------------------------------------------------------------
# Figure
# --------------------------------------------------------------------------

def binned_asd(freqs, asd, bins=140):
    """Median ASD in log-spaced frequency bins, for a trace that can be read.

    A raw periodogram at these lengths draws tens of thousands of points into a
    few hundred pixels and renders as a solid band whose thickness is chi-square
    scatter. Binning by the median keeps the floor where the quoted numbers put
    it, those being medians too, and keeps a line a line.
    """
    edges = np.geomspace(freqs[0], freqs[-1], bins + 1)
    index = np.digitize(freqs, edges) - 1
    centres, values = [], []
    for b in range(bins):
        inside = index == b
        if inside.any():
            centres.append(np.sqrt(edges[b] * edges[b + 1]))
            values.append(np.median(asd[inside]))
    return np.array(centres), np.array(values)


def make_figure(records, band, supplies, path, cal=None):
    conditions = conditions_in_order(records)
    variants = sorted({variant_of(l) for l, _, _ in records})
    color = {c: CONDITION_COLORS[i % len(CONDITION_COLORS)]
             for i, c in enumerate(conditions)}
    dash = {v: VARIANT_STYLES[i % len(VARIANT_STYLES)]
            for i, v in enumerate(variants)}
    marker = {v: VARIANT_MARKERS[i % len(VARIANT_MARKERS)]
              for i, v in enumerate(variants)}

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

    a = axs[0, 0]
    a.axvspan(band[0], band[1], color=ch.GRID, alpha=0.55, zorder=0,
              linewidth=0)
    for label, _, s in records:
        freqs, asd = binned_asd(*ch.welch_asd(
            s["axes"]["total"], s["fs"], nperseg=ch.segment_length(s["fs"])))
        # Fractional, because each condition has its own gain: in absolute
        # nT/rtHz a condition reading 7% smaller plots 7% lower while being
        # equally quiet.
        asd = asd / s["field"] * 1e6
        a.loglog(freqs, asd, color=color[condition_of(label)],
                 linestyle=dash[variant_of(label)], linewidth=1.4, alpha=0.9,
                 label=label)
    a.set_xlim(*ch.FRACTIONAL_ASD_HZ); a.set_ylim(*ch.FRACTIONAL_ASD_PPM)
    ch.decade_ticks(a)
    ch.title(a, "Spectral density of |B|, as a fraction of |B|")
    a.set_xlabel("frequency (Hz)"); a.set_ylabel("ppm of |B| per √Hz")
    a.annotate(f"shaded: {band[0]:g}–{band[1]:.0f} Hz, the compared band."
               "\nnormalised by each capture's own |B|, which keeps"
               "\na gain difference out of the noise figure",
               xy=(0.02, 0.05), xycoords="axes fraction",
               color=ch.TEXT_SECONDARY, fontsize=9)
    a.legend(frameon=False, fontsize=9, labelcolor=ch.TEXT_SECONDARY,
             loc="upper right")

    a = axs[0, 1]
    plotted = _plot_decimation(a, records, band, color, dash)
    ch.title(a, "Fast capture, filtered and decimated")
    a.set_xlabel("frequency (Hz)"); a.set_ylabel("ppm of |B| per √Hz")
    if plotted:
        a.set_xlim(*ch.FRACTIONAL_ASD_HZ); a.set_ylim(*ch.FRACTIONAL_ASD_PPM)
        ch.decade_ticks(a)
        a.legend(frameon=False, fontsize=9, labelcolor=ch.TEXT_SECONDARY,
                 loc="upper right")
    else:
        a.annotate("no pair of captures differs by an integer cycle-count "
                   "factor,\nso there is nothing to decimate onto",
                   xy=(0.5, 0.5), xycoords="axes fraction", ha="center",
                   color=ch.TEXT_SECONDARY, fontsize=10)

    a = axs[1, 0]
    for label, _, s in records:
        taus, devs = s["allan_curve"]
        a.loglog(taus, devs / s["field"] * 1e6,
                 color=color[condition_of(label)],
                 linestyle=dash[variant_of(label)], linewidth=1.6, label=label)
    a.set_xlim(*ch.ALLAN_TAU_S); a.set_ylim(*ch.FRACTIONAL_ALLAN_PPM)
    ch.decade_ticks(a)
    ch.title(a, "Allan deviation of |B|, as a fraction of |B|")
    a.set_xlabel("averaging time τ (s)"); a.set_ylabel("σ (ppm of |B|)")
    a.annotate("slope −½ = white noise; upturn = drift",
               xy=(0.02, 0.04), xycoords="axes fraction",
               color=ch.TEXT_SECONDARY, fontsize=9)
    a.legend(frameon=False, fontsize=9, labelcolor=ch.TEXT_SECONDARY,
             loc="upper right")

    _plot_corrections(axs[1, 1], records, supplies, color, marker, cal)

    fig.suptitle("RM3100 capture comparison", color=ch.TEXT_PRIMARY,
                 fontsize=15, fontweight="bold", y=0.985)
    fig.text(0.5, 0.945,
             f"{len(records)} captures; colour is the condition, line style "
             f"the cycle count. Broadband figures over "
             f"{band[0]:g}–{band[1]:.1f} Hz, inside every capture's Nyquist.",
             color=ch.TEXT_SECONDARY, fontsize=10, ha="center")
    fig.tight_layout(rect=[0, 0, 1, 0.935])
    fig.savefig(path, facecolor=ch.SURFACE)
    return path


def _plot_decimation(ax, records, band, color, dash):
    """Decimated-fast against natively-slow, per condition. True if anything."""
    ax.axvspan(band[0], band[1], color=ch.GRID, alpha=0.55, zorder=0,
               linewidth=0)
    plotted = False
    for condition, (_, fast), (_, slow), k in decimation_pairs(
            [(label, s) for label, _, s in records]):
        try:
            v = ch.decimate(fast["axes"]["total"], k, "fir")
        except ValueError:
            continue
        hue = color[condition]
        freqs, asd = binned_asd(*ch.welch_asd(v, fast["fs"] / k,
                                             nperseg=ch.segment_length(fast["fs"] / k)))
        # Fractional again: the two cycle counts have their own gains too.
        asd = asd / fast["field"] * 1e6
        ax.loglog(freqs, asd, color=hue, linestyle="--", linewidth=1.5,
                  label=f"{condition} cc{fast['cycle_count']} ÷{k}")
        freqs, asd = binned_asd(*ch.welch_asd(
            slow["axes"]["total"], slow["fs"],
            nperseg=ch.segment_length(slow["fs"])))
        asd = asd / slow["field"] * 1e6
        ax.loglog(freqs, asd, color=hue, linestyle="-", linewidth=1.5,
                  alpha=0.7,
                  label=f"{condition} cc{slow['cycle_count']} native")
        plotted = True
    if plotted:
        ax.annotate("dashed: fast capture decimated. solid: natively slow.",
                    xy=(0.02, 0.05), xycoords="axes fraction",
                    color=ch.TEXT_SECONDARY, fontsize=9)
    return plotted


def _plot_corrections(ax, records, supplies, color, marker, cal=None):
    """|B| under each correction, so the spread each one leaves is visible."""
    reference_volts = min(supplies.values()) if supplies else None
    for label, _, s in records:
        s["label"] = label
    corrected = [(label, corrected_fields(s, supplies, reference_volts, cal))
                 for label, _, s in records]
    rows = [(key, name) for key, name in correction_names(cal)
            if all(c[key] is not None for _, c in corrected)]

    for row, (key, name) in enumerate(rows):
        values = []
        for label, correction in corrected:
            values.append(correction[key])
            ax.plot(correction[key], -row, marker=marker[variant_of(label)],
                    color=color[condition_of(label)], markersize=9,
                    markeredgecolor=ch.SURFACE, markeredgewidth=2,
                    linestyle="none", zorder=3)
        ax.plot([min(values), max(values)], [-row, -row],
                color=ch.GRID, linewidth=6, solid_capstyle="round", zorder=1)
        ax.annotate(f"spread {(max(values)/min(values) - 1) * 100:.1f}%",
                    xy=(max(values), -row), xytext=(10, 0),
                    textcoords="offset points", va="center",
                    color=ch.TEXT_SECONDARY, fontsize=9)
    ax.set_yticks([-r for r in range(len(rows))], [n for _, n in rows])
    ax.set_ylim(-len(rows) + 0.5, 0.85)
    ax.margins(x=0.24)
    ch.title(ax, "Total field under each candidate correction")
    ax.set_xlabel("|B| (nT)")
    ax.annotate("a correction that is right collapses the spread",
                xy=(0.02, 0.06), xycoords="axes fraction",
                color=ch.TEXT_SECONDARY, fontsize=9)
    # Identity never rests on colour alone: the marker repeats the cycle count.
    handles = [plt.Line2D([], [], color=color[c], marker="o", markersize=9,
                          linestyle="none", label=c)
               for c in conditions_in_order(records)]
    handles += [plt.Line2D([], [], color=ch.TEXT_SECONDARY,
                           marker=marker[v], markersize=8, linestyle="none",
                           label=v or "(one variant)")
                for v in sorted({variant_of(l) for l, _, _ in records})]
    ax.legend(handles=handles, frameon=False, fontsize=9, ncol=2,
              labelcolor=ch.TEXT_SECONDARY, loc="upper left")


# --------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("csv", nargs="+", metavar="[LABEL=]CSV")
    ap.add_argument("--group", choices=["note", "file"], default="file",
                    help="where unlabelled captures take their name from "
                         "(default: %(default)s)")
    ap.add_argument("--trim", type=float, default=0.0, metavar="SECONDS",
                    help="drop this many seconds from BOTH ends of every "
                         "capture (default: %(default)s)")
    ap.add_argument("--band", type=parse_band, default=None, metavar="LO,HI",
                    help=f"broadband comparison band in Hz; the default is "
                         f"{ch.BAND_LO_HZ:g} Hz to "
                         f"{ch.BAND_NYQUIST_FRACTION * 100:g}%% of the Nyquist "
                         f"of the slowest path compared, decimation included")
    ap.add_argument("--supply", type=parse_supply, action="append",
                    default=[], metavar="CONDITION=VOLTS",
                    help="rail voltage for a condition, enabling the "
                         "ratiometric correction; repeatable")
    ap.add_argument("--calibration", default=None, metavar="JSON",
                    help="reference the rate correction to a measured "
                         "calibration in place of rm3100's nominal model")
    ap.add_argument("-o", "--output", default=None, metavar="PNG",
                    help="write a comparison figure here")
    args = ap.parse_args()
    supplies = dict(args.supply)
    try:
        cal = (calibrate.load_calibration(args.calibration)
               if args.calibration else None)
    except (OSError, calibrate.CalibrationError) as exc:
        sys.exit(str(exc))

    loaded = []
    for argument in args.csv:
        try:
            explicit, path = split_label(argument)
            cap = capture.load(path)
            cap, note = ch.trimmed(cap, args.trim)
        except (OSError, ValueError, capture.CaptureError) as exc:
            print(f"skipping {argument}: {exc}", file=sys.stderr)
            continue
        loaded.append((label_for(cap, path, args.group, explicit), path, cap,
                       note))
    if not loaded:
        sys.exit("nothing to compare")

    unknown = supplies.keys() - {condition_of(l) for l, _, _, _ in loaded}
    if unknown:
        print(f"warning: --supply names no capture: {', '.join(sorted(unknown))}",
              file=sys.stderr)

    # One band for every capture, so no run is scored over more bandwidth than
    # another; otherwise the difference in bandwidth reads as a noise difference.
    slowest = min(comparable_rates([(l, cap) for l, _, cap, _ in loaded]))
    band = args.band or ch.band_for(slowest)
    if band[1] > ch.BAND_NYQUIST_FRACTION * slowest / 2:
        print(f"warning: --band reaches {band[1]:g} Hz, past "
              f"{ch.BAND_NYQUIST_FRACTION * slowest / 2:.2f} Hz where the "
              f"slowest path compared ({slowest:.2f} Hz, decimation "
              f"included) rolls off", file=sys.stderr)

    records = [(label, path, summarise(cap, band)) for label, path, cap, _ in loaded]

    if args.trim:
        for label, _, _, note in loaded:
            if note:
                print(f"{label[:24]:<24} {note}")
        print()
    print_overview(records, band, supplies)
    print_rate_model(records)
    print_locked_lines(records)
    print_decimation(records, band)
    print_contrasts(records, supplies, cal)
    print_pairwise(records)

    if args.output:
        print(f"\n-> {make_figure(records, band, supplies, args.output, cal)}")
    return 0


if __name__ == "__main__":
    sys.exit(main())