#!/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     # group by header note

Built around one hard lesson: a naive comparison of absolute noise is wrong when
the measured scale differs between runs. A gain change carries the noise with it,
so a run that reads 6% larger also reads ~6% noisier while being physically
identical. Everything here is therefore reported **fractionally**, in ppm of the
field magnitude.

The second trap is attributing a scale change to the variable under test when
the sensor simply moved. Two diagnostics separate them:

  per-axis ratio spread   a pure gain change scales X, Y and Z identically,
                          so the spread is ~0. Anything larger means the sensor
                          moved.
  rotation angle          the angle between mean field directions. ~0 deg means
                          the sensor held still.

Both must be small before a magnitude difference can be blamed on gain. Note
that |B| is preserved under rotation but *not* under translation through a field
gradient, so a moved sensor can change magnitude on its own.
"""

import argparse
import itertools
import sys
from collections import defaultdict

import numpy as np

import capture
import characterize as ch

# 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


def summarise(cap):
    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}
    for key, _, _ in ch.SERIES:
        v = d[key]
        freqs, asd = ch.welch_asd(v, fs, nperseg=min(4096, len(v) // 4 * 2))
        band = freqs > min(3.0, fs / 8)
        out[key] = {
            "sd": v.std(),
            "ppm": v.std() / field * 1e6,
            "asd": float(np.median(asd[band])) if band.any() else float("nan"),
            "peak": float(asd[band].max()) if band.any() else float("nan"),
            "peak_hz": float(freqs[band][np.argmax(asd[band])]) if band.any() else float("nan"),
        }
    taus, devs = ch.allan_deviation(d["total"], fs)
    i = int(np.argmin(devs))
    out["allan"] = (float(devs[i]), float(taus[i]))
    return out


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


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("csv", nargs="+")
    ap.add_argument("--group", choices=["note", "file"], default="file",
                    help="group captures by header note, averaging repeats "
                         "(default: %(default)s)")
    args = ap.parse_args()

    recs = []
    for path in args.csv:
        try:
            cap = capture.load(path)
        except (OSError, capture.CaptureError) as exc:
            print(f"skipping {path}: {exc}", file=sys.stderr)
            continue
        recs.append((label_for(cap, path, args.group), path, summarise(cap)))
    if len(recs) < 1:
        sys.exit("nothing to compare")

    print(f"{'capture':<26} {'|B| nT':>10} {'sd nT':>8} {'sd ppm':>9} "
          f"{'ASD':>7} {'peak':>8} {'@Hz':>7} {'Allan':>7}")
    for label, path, s in recs:
        t = s["total"]
        print(f"{label[:26]:<26} {s['field']:10,.0f} {t['sd']:8.1f} {t['ppm']:9.1f} "
              f"{t['asd']:7.2f} {t['peak']:8.1f} {t['peak_hz']:7.3f} "
              f"{s['allan'][0]:7.2f}")
    print("            ASD/peak in nT/rtHz; Allan = best sigma by averaging")

    # Averaged per group, which is the number to compare when runs are repeated.
    if args.group == "note":
        groups = defaultdict(list)
        for label, _, s in recs:
            groups[label].append(s)
        if len(groups) > 1:
            print("\n=== group means (fractional -- the comparable figure) ===")
            for label, ss in groups.items():
                ppm = [s["total"]["ppm"] for s in ss]
                print(f"{label[:26]:<26} n={len(ss)}  "
                      f"|B| {np.mean([s['field'] for s in ss]):9,.0f} nT   "
                      f"sd {np.mean(ppm):7.1f} ppm"
                      + (f" +/- {np.std(ppm):.1f}" if len(ss) > 1 else ""))

    if len(recs) < 2:
        return 0

    print("\n=== pairwise: did the sensor hold still? ===")
    print("A magnitude difference only means gain if BOTH checks pass.")
    for (la, _, a), (lb, _, b) in itertools.combinations(recs, 2):
        ratio = b["mean"] / a["mean"]
        spread = float(np.ptp(ratio) / np.abs(np.mean(ratio)))
        ua, ub = a["mean"] / a["field"], b["mean"] / b["field"]
        angle = float(np.degrees(np.arccos(np.clip(np.dot(ua, ub), -1, 1))))
        moved = spread > RATIO_SPREAD_OK or angle > ROTATION_OK_DEG
        print(f"\n{la[:24]} -> {lb[:24]}")
        print(f"   |B| ratio        {b['field']/a['field']:.5f}  "
              f"({(b['field']/a['field'] - 1) * 100:+.2f}%)")
        print(f"   fractional noise {a['total']['ppm']:.1f} -> {b['total']['ppm']:.1f} ppm  "
              f"({(b['total']['ppm']/a['total']['ppm'] - 1) * 100:+.1f}%)")
        print(f"   per-axis ratios  X {ratio[0]:.5f}  Y {ratio[1]:.5f}  Z {ratio[2]:.5f}"
              f"   spread {spread * 100:.2f}%")
        print(f"   rotation         {angle:.3f} deg")
        if moved:
            print("   -> SENSOR MOVED. The magnitude difference cannot be "
                  "attributed to gain;\n      re-run with the sensor clamped.")
        else:
            print(f"   -> held still. The {(b['field']/a['field'] - 1) * 100:+.2f}% "
                  "magnitude difference is a real gain change.")
    return 0


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