#!/usr/bin/env python3
"""
validate_waveform.py — Per-effect waveform validator for the streaming-chain E2E.

Compares the client-recorded stream (``received_<id>.bin``, RCV1 format written
by chain-client) against the analytic ground truth (rebuilt from gen_data) and,
when present, the fed-reference tap (``tap_<id>.bin``, MARTe binary). It also
rolls in the client behavioural checks (``checks_<id>.json``) and emits
``metrics_<id>.json`` with an overall pass/fail (exit 0/1).

Oracle (per signal)
-------------------
* **Fidelity** (always): every received value must lie within ``tol`` of *some*
  ground-truth value. ``tol`` is 0 (bit-exact) for un-quantised integers, a tiny
  float epsilon for un-quantised floats, and ``quant_step/2 + 1e-6·range`` for
  quantised floats. Catches type corruption and out-of-range quantisation.
* **Shape** (sine signals, ≥8 points): least-squares fit of
  ``a·sin(ωt)+b·cos(ωt)+c`` at the scenario frequency ω=2πf. A high correlation
  (≥0.99) and low normalised RMSE confirm the received waveform is the expected
  sinusoid at the expected frequency (the fit recovers the unknown wall-clock
  phase offset automatically). nRMSE tolerance is relaxed by the quant step.
* **Fed reference** (when ``--tap`` given): each received value must be within
  ``tol`` of some tap value too.
"""
import argparse
import json
import os
import struct
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import scenarios as S  # noqa: E402
import gen_data as G    # noqa: E402

CODE_TO_TYPE = {v: k for k, v in S.TYPE_CODES.items()}


# ── readers ──────────────────────────────────────────────────────────────────

def read_received(path):
    """Read RCV1 → {key: (t ndarray, v ndarray)}."""
    with open(path, "rb") as f:
        data = f.read()
    if data[:4] != b"RCV1":
        raise ValueError(f"{path}: bad magic")
    off = 4
    nsig = struct.unpack_from("<I", data, off)[0]
    off += 4
    out = {}
    for _ in range(nsig):
        klen = struct.unpack_from("<H", data, off)[0]
        off += 2
        key = data[off:off + klen].decode()
        off += klen
        n = struct.unpack_from("<I", data, off)[0]
        off += 4
        t = np.frombuffer(data, "<f8", n, off).copy()
        off += 8 * n
        v = np.frombuffer(data, "<f8", n, off).copy()
        off += 8 * n
        out[key] = (t, v)
    return out


def read_marte_binary(path):
    """Read a MARTe binary file → {name: ndarray[n_rows, elements]} (typed)."""
    with open(path, "rb") as f:
        data = f.read()
    ns = struct.unpack_from("<I", data, 0)[0]
    off = 4
    descs = []
    for _ in range(ns):
        tc = struct.unpack_from("<H", data, off)[0]
        name = data[off + 2:off + 34].rstrip(b"\0").decode(errors="replace")
        ne = struct.unpack_from("<I", data, off + 34)[0]
        descs.append((name, tc, ne))
        off += 38
    row_bytes = sum(struct.calcsize(_npfmt(tc)) * ne for _, tc, ne in descs)
    nrows = (len(data) - off) // row_bytes if row_bytes else 0
    cols = {name: [] for name, _, _ in descs}
    for r in range(nrows):
        for name, tc, ne in descs:
            dt = np.dtype(S.NP_DTYPE[CODE_TO_TYPE[tc]])
            vals = np.frombuffer(data, dt, ne, off)
            cols[name].append(vals.astype(np.float64))
            off += dt.itemsize * ne
    return {name: np.array(cols[name]) for name, _, _ in descs}


def _npfmt(tc):
    return {"<i1": "b", "<u1": "B", "<i2": "h", "<u2": "H", "<i4": "i",
            "<u4": "I", "<i8": "q", "<u8": "Q", "<f4": "f",
            "<f8": "d"}[S.NP_DTYPE[CODE_TO_TYPE[tc]]]


# ── metrics ──────────────────────────────────────────────────────────────────

def _tol(gt):
    if gt["quant"] and gt["quant"] != "none":
        levels = S.QUANT_LEVELS[gt["quant"]]
        rng = gt["range_max"] - gt["range_min"]
        step = rng / levels
        return step / 2.0 + 1e-6 * abs(rng), step
    if gt["type"] in S.FLOAT_TYPES:
        return 1e-3, 0.0           # float round-trip epsilon
    return 0.5, 0.0                # integer: rounding-exact (within 0.5)


def nearest_err(recv_v, truth_v):
    """Max distance from each received value to the nearest truth value."""
    sv = np.unique(np.sort(truth_v.astype(np.float64)))
    idx = np.searchsorted(sv, recv_v)
    idx = np.clip(idx, 1, len(sv) - 1) if len(sv) > 1 else np.zeros_like(idx)
    lo = sv[np.clip(idx - 1, 0, len(sv) - 1)]
    hi = sv[np.clip(idx, 0, len(sv) - 1)]
    d = np.minimum(np.abs(recv_v - lo), np.abs(recv_v - hi))
    return float(np.max(d)) if d.size else 0.0


def sine_shape(t, v, freq):
    """Return (corr, nrmse, amp_fit) for a sinusoid fit at ``freq``."""
    w = 2.0 * np.pi * freq
    A = np.column_stack([np.sin(w * t), np.cos(w * t), np.ones_like(t)])
    coef, *_ = np.linalg.lstsq(A, v, rcond=None)
    fit = A @ coef
    span = float(np.max(v) - np.min(v)) or 1.0
    nrmse = float(np.sqrt(np.mean((v - fit) ** 2)) / span)
    corr = float(np.corrcoef(v, fit)[0, 1]) if np.std(v) > 0 else 1.0
    amp = float(np.hypot(coef[0], coef[1]))
    return corr, nrmse, amp


def compare_signal(gt, t_recv, v_recv, tap_v=None):
    tol, step = _tol(gt)
    truth_v = gt["v"].astype(np.float64)
    m = {
        "type": gt["type"], "quant": gt["quant"], "formula": gt["formula"],
        "n_truth": int(truth_v.size), "n_recv": int(v_recv.size),
        "quant_step": step, "tol": tol,
    }
    if v_recv.size == 0:
        m["pass"] = False
        m["reason"] = "no received samples"
        return m

    max_err = nearest_err(v_recv, truth_v)
    m["max_abs_err"] = max_err
    fidelity_ok = max_err <= tol
    m["fidelity_ok"] = bool(fidelity_ok)

    shape_ok = True
    if gt["formula"] == "sine" and v_recv.size >= 8 and gt["freq"]:
        corr, nrmse, amp = sine_shape(t_recv, v_recv, gt["freq"])
        nrmse_tol = 0.05 + (step / (gt["range_max"] - gt["range_min"])
                            if gt["quant"] != "none" else 0.0)
        shape_ok = corr >= 0.99 and nrmse <= nrmse_tol
        m.update(corr=corr, nrmse=nrmse, amp_fit=amp,
                 nrmse_tol=nrmse_tol, shape_ok=bool(shape_ok))

    fed_ok = True
    if tap_v is not None and tap_v.size:
        fed_err = nearest_err(v_recv, tap_v.astype(np.float64))
        fed_ok = fed_err <= tol
        m.update(fed_err=fed_err, fed_ok=bool(fed_ok))

    m["pass"] = bool(fidelity_ok and shape_ok and fed_ok)
    return m


# ── driver ────────────────────────────────────────────────────────────────────

def validate(scenario, received, tap, checks):
    gt = G.build_ground_truth(scenario)
    recv = read_received(received)
    tap_cols = read_marte_binary(tap) if tap and os.path.exists(tap) else None

    sigs = {}
    overall = True
    for key, (t, v) in sorted(recv.items()):
        # key is "src:sig" or "src:sig[i]" (per-element array push)
        base = key.split("[")[0]
        if base not in gt:
            sigs[key] = {"pass": True, "note": "no ground truth (skipped)",
                         "n_recv": int(v.size)}
            continue
        tap_v = None
        if tap_cols is not None:
            sig_name = base.split(":", 1)[1]
            if sig_name in tap_cols:
                tap_v = tap_cols[sig_name].reshape(-1)
        m = compare_signal(gt[base], t, v, tap_v)
        sigs[key] = m
        overall = overall and m.get("pass", False)

    # roll in client checks
    client = {}
    if checks and os.path.exists(checks):
        with open(checks) as f:
            client = json.load(f)
        live_ok = client.get("live", {}).get("ok", False)
        zoom_ok = all(z.get("inrange", False) for z in client.get("zoom", [])) \
            if client.get("zoom") else True
        win = client.get("window", {})
        window_ok = win.get("ok", True) if win.get("returned", 0) else True
        trig = client.get("trigger", [])
        trig_ok = all(t.get("fired", False) and t.get("windowOk", False)
                      for t in trig) if trig else True
        overall = overall and live_ok and zoom_ok and window_ok and trig_ok
        client["_rollup"] = {"live_ok": live_ok, "zoom_ok": zoom_ok,
                             "window_ok": window_ok, "trigger_ok": trig_ok}

    return {"scenario": scenario["id"], "signals": sigs,
            "client": client, "pass": bool(overall)}


def _selftest():
    import math
    t = np.linspace(0, 1, 200)
    truth = np.sin(2 * np.pi * 3 * t)
    gt = {"v": truth, "type": "float32", "quant": "uint16", "formula": "sine",
          "freq": 3.0, "range_min": -1.0, "range_max": 1.0, "elements": 1,
          "rows": 200, "is_time": False, "t": t, "dt": t[1] - t[0]}
    _, step = _tol(gt)
    good = truth + (np.random.rand(200) - 0.5) * step      # ≤ step/2
    bad = truth + (np.random.rand(200) - 0.5) * step * 8    # ≫ step
    mg = compare_signal(gt, t, good)
    mb = compare_signal(gt, t, bad)
    assert mg["pass"], f"good should pass: {mg}"
    assert not mb["pass"], f"bad should fail: {mb}"
    # bit-exact integer
    gi = dict(gt); gi.update(type="uint32", quant="none",
                             v=np.arange(200, dtype=np.float64), formula="counter")
    mi = compare_signal(gi, t, np.arange(50, 150, dtype=np.float64))
    assert mi["pass"], f"int subset should pass: {mi}"
    mi2 = compare_signal(gi, t, np.array([1.5, 250.0]))
    assert not mi2["pass"], f"int off-grid should fail: {mi2}"
    print("selftest OK")


def main():
    p = argparse.ArgumentParser(description="Validate received waveform")
    p.add_argument("--selftest", action="store_true")
    p.add_argument("--scenario")
    p.add_argument("--received")
    p.add_argument("--tap", default=None)
    p.add_argument("--checks", default=None)
    p.add_argument("--out", default=None)
    args = p.parse_args()

    if args.selftest:
        _selftest()
        sys.exit(0)

    sc = next((s for s in S.SCENARIOS if s["id"] == args.scenario), None)
    if sc is None:
        print(f"unknown scenario {args.scenario}", file=sys.stderr)
        sys.exit(2)
    res = validate(sc, args.received, args.tap, args.checks)
    if args.out:
        with open(args.out, "w") as f:
            json.dump(res, f, indent=2)
    print(f"{sc['id']}: {'PASS' if res['pass'] else 'FAIL'}")
    for k, m in res["signals"].items():
        print(f"  {k}: pass={m.get('pass')} "
              f"err={m.get('max_abs_err','-')} corr={m.get('corr','-')}")
    sys.exit(0 if res["pass"] else 1)


if __name__ == "__main__":
    main()