#!/usr/bin/env python3
"""Compare Pi-GPIO ground-truth HS edges against the two USB observation paths.

Inputs:
  hs_gpio_machine.jsonl    kernel-stamped edges from hs_gpio_logger.py (truth)
  hs_timing_machine.jsonl  uartd trace (optional): us_hs_obs = GET_REG poll,
                           k_gpio_int = XR21 Table-8 interrupt endpoint

All three sources use CLOCK_MONOTONIC on the same host, so the deltas below are
real and not cross-domain guesses.

Answers three questions:
  1. Do HS windows shorter than the ~2.75 ms poll resolution exist? (census)
  2. How late is the GET_REG poll, and how many windows does it miss entirely?
  3. How wrong is the 8 ms-quantized interrupt endpoint against truth?

Usage:
  python3 hs_gpio_analyze.py sessions/<stamp>/5/data/
  python3 hs_gpio_analyze.py --truth hs_gpio_machine.jsonl --usb hs_timing_machine.jsonl
  python3 hs_gpio_analyze.py --self-test
"""

from __future__ import annotations

import argparse
import json
import math
import sys
from pathlib import Path
from typing import Any, Iterable, Optional, Sequence

WIDTH_BUCKETS_US = (50, 100, 250, 500, 1000, 2000, 2750, 3000, 5000, 8000)
MATCH_WINDOW_NS = 200_000_000


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as fh:
        for lineno, line in enumerate(fh, 1):
            line = line.strip()
            if not line:
                continue
            try:
                out.append(json.loads(line))
            except json.JSONDecodeError as exc:
                raise SystemExit(f"{path}:{lineno}: invalid JSON: {exc}") from exc
    return out


def percentile(values: Sequence[float], pct: float) -> Optional[float]:
    if not values:
        return None
    if len(values) == 1:
        return float(values[0])
    k = (len(values) - 1) * (pct / 100.0)
    lo, hi = math.floor(k), math.ceil(k)
    if lo == hi:
        return float(values[int(k)])
    return float(values[lo] * (hi - k) + values[hi] * (k - lo))


def summarize_us(samples: Iterable[float]) -> dict[str, Any]:
    s = sorted(float(x) for x in samples)
    if not s:
        return {"n": 0}
    return {
        "n": len(s),
        "min_us": round(s[0] / 1000.0, 2),
        "p50_us": round(percentile(s, 50) / 1000.0, 2),  # type: ignore[operator]
        "p95_us": round(percentile(s, 95) / 1000.0, 2),  # type: ignore[operator]
        "p99_us": round(percentile(s, 99) / 1000.0, 2),  # type: ignore[operator]
        "max_us": round(s[-1] / 1000.0, 2),
    }


class Truth:
    """Ground-truth HS edges and the windows derived from them."""

    def __init__(self, records: list[dict[str, Any]]) -> None:
        self.header: dict[str, Any] = {}
        self.edges: list[tuple[int, int, int]] = []  # (ts_ns, hs, seqno)
        self.drain: list[int] = []
        self.seq_gaps: list[dict[str, Any]] = []
        self.stats: list[dict[str, Any]] = []

        for rec in records:
            kind = rec.get("type")
            if kind == "hdr":
                self.header = rec
            elif kind == "hs_edge":
                if rec.get("ts_ns") is None:
                    continue
                ts = int(rec["ts_ns"])
                self.edges.append((ts, int(rec.get("hs", 0)), int(rec.get("seqno", -1))))
                if rec.get("recv_ns") is not None:
                    lag = int(rec["recv_ns"]) - ts
                    if 0 <= lag < 5_000_000_000:
                        self.drain.append(lag)
            elif kind == "seq_gap":
                self.seq_gaps.append(rec)
            elif kind == "stat":
                self.stats.append(rec)
        self.edges.sort()

        self.gap_after: set[int] = set()
        for gap in self.seq_gaps:
            self.gap_after.add(int(gap.get("expected_seqno", -1)) - 1)

    @property
    def lost(self) -> int:
        return sum(int(g.get("lost", 0) or 0) for g in self.seq_gaps)

    def windows(self, level: int) -> list[tuple[int, int]]:
        """Closed [start, end] intervals where logical HS == level.

        Pairs spanning a dropped-event gap or a non-alternating edge are skipped:
        their width would be a fiction.
        """
        out: list[tuple[int, int]] = []
        for (t0, h0, s0), (t1, h1, _s1) in zip(self.edges, self.edges[1:]):
            if h0 != level or h1 == h0:
                continue
            if s0 in self.gap_after:
                continue
            if t1 > t0:
                out.append((t0, t1))
        return out

    def non_alternating(self) -> int:
        return sum(1 for a, b in zip(self.edges, self.edges[1:]) if a[1] == b[1])


class UsbTrace:
    """uartd GET_REG poll observations and XR Table-8 interrupt records."""

    def __init__(self, records: list[dict[str, Any]]) -> None:
        self.header: dict[str, Any] = {}
        self.obs: list[tuple[int, int]] = []
        self.kernel: list[tuple[int, int]] = []
        for rec in records:
            kind = rec.get("type")
            if kind == "hdr":
                self.header = rec
            elif kind == "us_hs_obs":
                if rec.get("mono_ns") is None:
                    continue
                self.obs.append((int(rec["mono_ns"]), int(rec.get("a", 0) or 0)))
            elif kind == "k_gpio_int":
                if rec.get("c") is None:
                    continue
                flags = int(rec.get("a", 0) or 0)
                if flags & 0x01:
                    self.kernel.append((int(rec["c"]), 1 if (flags & 0x08) else 0))
        self.obs.sort()
        self.kernel.sort()

    def level_times(self, source: str, level: int) -> list[int]:
        src = self.obs if source == "poll" else self.kernel
        return [t for t, lv in src if lv == level]


def first_at_or_after(times: Sequence[int], target: int) -> Optional[int]:
    import bisect

    idx = bisect.bisect_left(times, target)
    if idx < len(times):
        return times[idx]
    return None


def any_in_range(times: Sequence[int], lo: int, hi: int) -> bool:
    import bisect

    idx = bisect.bisect_left(times, lo)
    return idx < len(times) and times[idx] <= hi


def width_census(widths_ns: Sequence[int], label: str) -> list[str]:
    lines: list[str] = []
    total = len(widths_ns)
    s = sorted(widths_ns)
    lines.append(f"  {label}: {summarize_us(s)}")
    if not total:
        return lines
    for bucket in WIDTH_BUCKETS_US:
        n = sum(1 for w in s if w < bucket * 1000)
        if n:
            lines.append(f"      < {bucket / 1000.0:>6.3f} ms : {n:6d}  ({100.0 * n / total:.3f}%)")
    lines.append(f"      narrowest 10 (ms): {[round(w / 1e6, 3) for w in s[:10]]}")
    return lines


def analyze(truth: Truth, usb: Optional[UsbTrace], *, miss_margin_us: float) -> list[str]:
    out: list[str] = []
    hdr = truth.header
    out.append("HS GPIO ground-truth analysis")
    out.append(
        f"  tap: {hdr.get('signal', '?')} on {hdr.get('chip_path', '?')} "
        f"({hdr.get('chip_label', '?')}) line {hdr.get('line', '?')} "
        f"active_low={hdr.get('active_low')}"
    )
    out.append(f"  eg={hdr.get('eg', '?')} pos={hdr.get('fixture_pos', '?')} clock={hdr.get('event_clock', '?')}")

    if not truth.edges:
        out.append("  NO EDGES CAPTURED — check wiring, polarity, and that the bus is active.")
        return out

    span_ns = truth.edges[-1][0] - truth.edges[0][0]
    out.append(f"  edges: {len(truth.edges)}  span: {span_ns / 1e9:.1f} s")
    out.append(f"  seq gaps: {len(truth.seq_gaps)} (lost {truth.lost})   non-alternating: {truth.non_alternating()}")
    if truth.seq_gaps:
        out.append("  WARNING: event buffer overflowed; widths across gaps excluded.")
    out.append(f"  delivery drain lag: {summarize_us(truth.drain)}")
    out.append("    (drain lag does not affect edge timestamps — kernel stamps at IRQ)")

    highs = truth.windows(1)
    lows = truth.windows(0)
    hw = [b - a for a, b in highs]
    lw = [b - a for a, b in lows]

    out.append("")
    out.append("== TRUE HS WIDTHS (the question USB cannot answer) ==")
    out.extend(width_census(hw, "HS active (high) width"))
    out.extend(width_census(lw, "HS inactive (low) width"))

    sub = sum(1 for w in hw if w < 3_000_000)
    verdict = (
        f"  VERDICT: {sub} of {len(hw)} HS-active windows are < 3 ms"
        if hw
        else "  VERDICT: no complete HS windows captured"
    )
    out.append("")
    out.append(verdict)
    if hw:
        if sub == 0:
            out.append("    => no sub-3 ms structure seen; the USB poll floor is not hiding windows.")
        else:
            out.append("    => sub-3 ms structure EXISTS and the USB poll cannot resolve it.")

    if usb is None:
        out.append("")
        out.append("  (no USB trace supplied — skipping poll / interrupt comparison)")
        return out

    margin_ns = int(miss_margin_us * 1000)
    for source, name in (("poll", "GET_REG poll (us_hs_obs)"), ("kernel", "XR Table-8 INT (k_gpio_int)")):
        times1 = usb.level_times(source, 1)
        if not times1:
            out.append("")
            out.append(f"== {name}: no samples in trace ==")
            continue
        lags: list[int] = []
        missed = 0
        missed_widths: list[int] = []
        for start, end in highs:
            if any_in_range(times1, start, end + margin_ns):
                nxt = first_at_or_after(times1, start)
                if nxt is not None and nxt - start < MATCH_WINDOW_NS:
                    lags.append(nxt - start)
            else:
                missed += 1
                missed_widths.append(end - start)
        out.append("")
        out.append(f"== {name} vs truth ==")
        out.append(f"  observation lag after true HS rise: {summarize_us(lags)}")
        total_h = len(highs)
        pct = (100.0 * missed / total_h) if total_h else 0.0
        out.append(f"  HS-active windows never observed: {missed} of {total_h} ({pct:.3f}%)")
        if missed_widths:
            out.append(f"    missed-window widths: {summarize_us(missed_widths)}")
    return out


def resolve_inputs(args: argparse.Namespace) -> tuple[Path, Optional[Path]]:
    if args.data_dir:
        d = Path(args.data_dir)
        if not d.is_dir():
            raise SystemExit(f"FATAL: not a directory: {d}")
        truth = Path(args.truth) if args.truth else d / "hs_gpio_machine.jsonl"
        usb = Path(args.usb) if args.usb else d / "hs_timing_machine.jsonl"
        return truth, (usb if usb.exists() else None)
    if not args.truth:
        raise SystemExit("FATAL: supply a data dir, or --truth (and optionally --usb)")
    usb_path = Path(args.usb) if args.usb else None
    if usb_path is not None and not usb_path.exists():
        usb_path = None
    return Path(args.truth), usb_path


def self_test() -> int:
    """Synthetic edges exercise width, gap exclusion, and miss detection."""
    errors: list[str] = []
    ms = 1_000_000

    recs: list[dict[str, Any]] = [
        {"type": "hdr", "signal": "SLV_HS", "line": 9, "active_low": True, "event_clock": "monotonic"}
    ]
    # 16 ms active windows at t=0 and t=100 ms, then a deliberate 800 us window.
    seq = 0
    for start, width in ((0, 16 * ms), (100 * ms, 16 * ms), (200 * ms, 800_000)):
        recs.append({"type": "hs_edge", "ts_ns": start, "recv_ns": start + 50_000, "hs": 1, "seqno": seq})
        seq += 1
        recs.append(
            {"type": "hs_edge", "ts_ns": start + width, "recv_ns": start + width + 50_000, "hs": 0, "seqno": seq}
        )
        seq += 1

    truth = Truth(recs)
    if len(truth.edges) != 6:
        errors.append(f"edges={len(truth.edges)} != 6")
    highs = truth.windows(1)
    widths = sorted(b - a for a, b in highs)
    if widths != [800_000, 16 * ms, 16 * ms]:
        errors.append(f"high widths={widths}")
    if truth.non_alternating() != 0:
        errors.append("unexpected non-alternating edges")

    gapped = list(recs) + [
        {"type": "seq_gap", "expected_seqno": 6, "got_seqno": 9, "lost": 3},
        {"type": "hs_edge", "ts_ns": 300 * ms, "recv_ns": 300 * ms, "hs": 1, "seqno": 9},
        {"type": "hs_edge", "ts_ns": 300 * ms + 4 * ms, "recv_ns": 300 * ms, "hs": 0, "seqno": 10},
    ]
    tg = Truth(gapped)
    if tg.lost != 3:
        errors.append(f"lost={tg.lost} != 3")
    if len(tg.windows(1)) != 4:
        errors.append(f"gapped high windows={len(tg.windows(1))} != 4")

    # Poll sees the two 16 ms windows but not the 800 us one.
    usb_recs = [
        {"type": "us_hs_obs", "mono_ns": 2 * ms, "a": 1},
        {"type": "us_hs_obs", "mono_ns": 17 * ms, "a": 0},
        {"type": "us_hs_obs", "mono_ns": 103 * ms, "a": 1},
        {"type": "us_hs_obs", "mono_ns": 118 * ms, "a": 0},
        {"type": "k_gpio_int", "c": 8 * ms, "a": 0x09},
        {"type": "k_gpio_int", "c": 24 * ms, "a": 0x01},
    ]
    usb = UsbTrace(usb_recs)
    if len(usb.obs) != 4:
        errors.append(f"usb obs={len(usb.obs)} != 4")
    if usb.level_times("kernel", 1) != [8 * ms]:
        errors.append(f"kernel level1={usb.level_times('kernel', 1)}")

    report = "\n".join(analyze(truth, usb, miss_margin_us=0.0))
    if "1 of 3" not in report:
        errors.append("expected the 800 us window to be reported as missed by the poll")
    if "sub-3 ms structure EXISTS" not in report:
        errors.append("expected sub-3ms verdict")

    if errors:
        for e in errors:
            print(f"FAIL: {e}")
        return 1
    print("hs_gpio_analyze self-test OK")
    return 0


def main(argv: Optional[list[str]] = None) -> int:
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("data_dir", nargs="?", default=None, help="session data dir containing both jsonl files")
    ap.add_argument("--truth", default=None, help="hs_gpio_machine.jsonl")
    ap.add_argument("--usb", default=None, help="hs_timing_machine.jsonl")
    ap.add_argument(
        "--miss-margin-us",
        type=float,
        default=2750.0,
        help="grace added to a true window before calling it missed (default = one poll period)",
    )
    ap.add_argument("--self-test", action="store_true")
    args = ap.parse_args(argv)

    if args.self_test:
        return self_test()

    truth_path, usb_path = resolve_inputs(args)
    if not truth_path.exists():
        raise SystemExit(f"FATAL: truth log not found: {truth_path}")
    truth = Truth(load_jsonl(truth_path))
    usb = UsbTrace(load_jsonl(usb_path)) if usb_path else None

    print(f"truth: {truth_path}")
    print(f"usb:   {usb_path if usb_path else '(none)'}")
    print()
    for line in analyze(truth, usb, miss_margin_us=args.miss_margin_us):
        print(line)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())