#!/usr/bin/env python3
"""
jank_eval.py — deterministic 120 Hz scroll-jank evaluator.

Reads app-side perf logs (PRES/VSYNC lines from the vendored stg/vkprobe
instrumentation) and produces a strict, repeatable PASS/FAIL report against
the zero-jank definition (docs/BUGS/hypotheses-120hz.md):

  PASS iff, in every ACTIVE window (sustained rendering):
    1. no DROPPED slot  (present in slot k, next present in slot k+2+)
    2. no DOUBLE        (two presents in the same display slot)
    3. no LATE present  (edge-grid phase > +5.0 ms, past the SF latch)
    4. no STALL         (present-to-present > 20 ms)
    5. phase spread p10-p90 <= 0.5 ms (edge grid: phase within slot)

SLOT GRID (the part that takes care to get right):
  * PRIMARY: the VSYNC `ftn=` field is the hardware vsync EDGE timestamp.
    slot(m) = floor((m - ftn0)/P), phase = (m-ftn0) mod P in [0,P).
    Exact and unambiguous for any present phase. P = median 1-slot ftn
    delta (hardware truth).
  * FALLBACK (no VSYNC lines, e.g. emulator sanity runs): slot =
    round((m - anchor)/P) with anchor = circular mean of present phases.
    Ambiguous for presents >P/2 from the anchor (documented limitation);
    LATE detection is unreliable in this mode.

Windows:  --windows "startMs,endMs name; ..." (ms from first present)
or auto-detected (present rate > 60/s active, < 10/s idle).

Exit: 0 = PASS, 1 = FAIL, 2 = no usable data.
Self-test: --selftest  (synthesizes log FILES and runs the same main()
path end-to-end; the tool must classify each case or exit non-zero).
"""
import argparse, bisect, math, os, re, sys, tempfile
from collections import defaultdict

PRES_RE = re.compile(r"^PRES f=(\d+)\w* mon=(\d+)")
VSYNC_RE = re.compile(r"^VSYNC ftn=(\d+) mon=(\d+)")

def parse_log(lines):
    pres, ftns = [], []
    for l in lines:
        m = PRES_RE.match(l)
        if m:
            pres.append((int(m.group(1)), int(m.group(2))))
            continue
        m = VSYNC_RE.match(l)
        if m:
            ftns.append(int(m.group(1)))
    return pres, ftns

def vsync_period(ftns):
    """Hardware period: median of the 1-slot ftn deltas."""
    if len(ftns) < 30: return None
    d = sorted(b - a for a, b in zip(ftns, ftns[1:]) if 4_000_000 < b - a < 25_000_000)
    return d[len(d)//2] if d else None

def fallback_period(monots):
    """Median of consecutive PRES deltas, snapped to the nearest standard
    refresh period. (The telescoping MEAN is useless here: it spans idle
    gaps too — 391 presents over a 23 s run fit a '59.6 ms' period on a 60 Hz
    display.) The display is always a standard rate, so snap. Do NOT use a
    fitted period for the verdict on the phone; VSYNC edges are the truth.
    """
    if len(monots) < 2:
        return 8_333_333
    d = sorted(monots[i+1] - monots[i] for i in range(len(monots) - 1))
    med = d[len(d)//2]
    STD_MS = [1/240, 1/144, 1/120, 1/100, 1/90, 1/85, 1/72, 1/60, 1/50,
              1/48, 1/40, 1/30, 1/24]
    best = min(STD_MS, key=lambda r: abs(1.0/r - med/1e6))
    return int(1_000_000_000 / (1.0/best))

def anchor_phase(monots, period):
    """Circular mean of present phases (handles 0/P wrap and bimodal
    wake-source phases; the anchor stays <P/2 from every present as long
    as the modes are <P apart)."""
    sx = sy = 0.0
    for m in monots[:2000]:
        q = (m % period)/period * 2*math.pi
        sx += math.cos(q); sy += math.sin(q)
    ang = math.atan2(sy, sx)
    if ang < 0: ang += 2*math.pi
    return int(ang/(2*math.pi)*period)

def slot_of_ftn(m, ftns, P):
    """Assign a present to the slot of the NEAREST OBSERVED hardware edge,
    in a TIME-CONTINUOUS slot domain. (2026-10-30: the cumulative grid
    anchored at ftns[0] is WRONG after a display phase-slip — the 120 Hz
    panel clock re-locks at an arbitrary phase after idle blanking, so every
    present after the slip measures ~5-6 ms 'late' against the stale grid
    while landing perfectly against the real edges. Measured on the phone:
    an entire drag window (w1) read late=230/dropped=24 while per-cycle
    tracing showed presents at edge +2.0 ms.)

    The slot of present m is: round((ftn_i − ftns[0])/P) where ftn_i is the
    nearest observed edge before m (plus one when m is in that edge's second
    half). A phase-slip Δ<P shifts EVERY ftn by the same Δ, so it only adds
    a constant to all slots — drops/doubles (slot DIFFERENCES) are immune,
    and the phase reported is m − ftn_i (true distance from the hardware
    edge). Stalls are time-accurate because ftn times are CLOCK_MONOTONIC.
    When no edge was observed within the last P (chain off), extrapolate on
    the plain time grid from ftns[0] (same slot domain: round-based)."""
    i = bisect.bisect_right(ftns, m) - 1
    if i < 0:
        i = 0
    d = m - ftns[i]
    if d > P:
        x = (m - ftns[0]) / P
        s = int(x + 0.5) if x >= 0 else -int(-x + 0.5)
        return s, (m - ftns[0]) - s * P
    base = (ftns[i] - ftns[0]) / P
    s = int(base + 0.5) if base >= 0 else -int(-base + 0.5)
    if d >= P / 2:
        s += 1
    return s, d

def slot_of(m, P, grid):
    """grid = ('edge', ftn0) or ('anchor', anchor). Returns (slot, phase_ns).
    edge:   phase in [0, P) from the vsync edge (hardware truth)
    anchor: phase in (-P/2, P/2] centered on the anchor (fallback)"""
    mode, ref = grid
    if mode == "edge":
        s = (m - ref)//P
        return s, m - ref - s*P
    x = (m - ref)/P
    s = int(x + 0.5) if x >= 0 else -int(-x + 0.5)
    return s, (m - ref) - s*P

def auto_windows(pres, period):
    """NOTE: the active threshold must NOT be high — it is circular (a JANKY
    app has a LOW present rate and would then not be detected as active at
    all: emulated 60 Hz scroll presenting at 30-40 Hz missed a 0.75×
    threshold). 0.4×rate catches a scroll running at a quarter of refresh;
    idle (0-2 presents/s) stays far below it."""
    if not pres: return []
    t0 = pres[0][1]
    rate = 1000.0/(period/1e6)       # display Hz (120 phone, 60 emulator)
    hi, lo = 0.4*rate, 0.1*rate
    BUCK = 250_000_000
    dens = defaultdict(int)
    for f, m in pres:
        dens[(m - t0)//BUCK] += 1
    last = int((pres[-1][1] - t0)//BUCK)
    wins, start = [], None
    for b in range(last + 1):
        r = dens.get(b, 0) * 4   # per-second rate (BUCK = 250 ms)
        if r > hi and start is None:
            start = b
        elif r < lo and start is not None:
            wins.append((start, b)); start = None
    if start is not None:
        wins.append((start, last + 1))
    out = []
    for i, (b0, b1) in enumerate(wins):
        idx0 = next((k for k, (f, m) in enumerate(pres) if (m - t0)//BUCK >= b0), None)
        idx1 = len(pres)
        for k in range(len(pres) - 1, -1, -1):
            if (pres[k][1] - t0)//BUCK < b1:
                idx1 = k + 1
                break
        if idx0 is not None:
            out.append((f"w{i}", idx0, idx1))
    return out

def explicit_windows(pres, t0, spec):
    out = []
    for part in spec.split(";"):
        part = part.strip()
        if not part: continue
        rng, _, name = part.partition(" ")
        a, b = rng.split(",")
        a, b = int(a)*1_000_000, int(b)*1_000_000
        idx0 = next((k for k, (f, m) in enumerate(pres) if m >= t0 + a), None)
        idx1 = len(pres)
        for k in range(len(pres) - 1, -1, -1):
            if pres[k][1] <= t0 + b:
                idx1 = k + 1
                break
        if idx0 is not None:
            out.append((name or f"w{len(out)}", idx0, idx1))
    return out

def evaluate(pres, P, grid, windows, suspect=False, ftns=None):
    """suspect=True (fitted grid, no VSYNC edges): drops/doubles are reported
    for information but excluded from the verdict — a fitted grid can't prove
    them. Stalls and late-phase remain meaningful.
    ftns: observed hardware vsync edges — when given, slots/phases are
    assigned relative to the NEAREST observed edge (phase-slip robust); see
    slot_of_ftn."""
    t0 = pres[0][1]
    use_ftn = ftns is not None and len(ftns) > 10
    any_fail = False
    total = 0
    for name, i0, i1 in windows:
        seg = pres[i0:i1]
        total += len(seg)
        if len(seg) < 10:
            print(f"\n=== window '{name}' n={len(seg)} -> SKIP (short)")
            continue
        drops = doubles = lates = stalls = jit = 0
        drop_at, stall_at, late_at, double_at = [], [], [], []
        prev_s = prev_p = None
        phases = []
        for k in range(len(seg)):
            f, m = seg[k]
            if use_ftn:
                s, p = slot_of_ftn(m, ftns, P)
            else:
                s, p = slot_of(m, P, grid)
            phases.append(p)
            if prev_s is not None:
                gap = s - prev_s
                if gap >= 2:
                    drops += gap - 1
                    drop_at.append((k, gap, (m - seg[k-1][1])/1e6))
                elif gap == 0:
                    doubles += 1
                    double_at.append(k)
            if p > 5_000_000:
                lates += 1
                late_at.append((k, p/1e6))
            if k and (m - seg[k-1][1]) > 20_000_000:
                stalls += 1
                stall_at.append((k, (m - seg[k-1][1])/1e6))
            if prev_p is not None and abs(p - prev_p) > 500_000:
                jit += 1
            prev_s, prev_p = s, p
        ps = sorted(phases)
        spread = (ps[int(len(ps)*9)//10] - ps[len(ps)//10])/1e6
        ok = (lates == 0 and stalls == 0 and spread <= 0.5)
        if not suspect:
            ok = ok and drops == 0 and doubles == 0
        any_fail = any_fail or not ok
        print(f"\n=== window '{name}'  n={len(seg)}  "
              f"[{(seg[0][1]-t0)/1e6:.0f}..{(seg[-1][1]-t0)/1e6:.0f} ms]  ->  "
              f"{'PASS' if ok else 'FAIL'}")
        print(f"    dropped={drops} doubles={doubles} late={lates} stalls={stalls}  "
              f"phase p10-p90={spread:.2f} ms  phase-jumps>0.5ms={jit}")
        for k, gap, ms in drop_at[:5]:
            print(f"      DROP   #{k}: skipped {gap-1} slot(s), gap {ms:.1f} ms")
        for k, ms in stall_at[:5]:
            print(f"      STALL  #{k}: {ms:.1f} ms")
        for k, ms in late_at[:5]:
            print(f"      LATE   #{k}: phase +{ms:.2f} ms")
        for k in double_at[:5]:
            print(f"      DOUBLE #{k}")
    print(f"\n=== SUMMARY  windows={len(windows)} frames={total}  ->  "
          + ("PASS" if not any_fail and windows else "FAIL"))
    return (not any_fail) and bool(windows)


# ---------------------------------------------------------------- self-test
def selftest():
    """Synthesize LOG FILES (VSYNC ftn + PRES lines) and run the same
    run_file() path as production. Cases must be classified exactly."""
    P = 8_333_333
    T0 = 100_000_000
    def gen(name, off, skips=(), doubles=(), stall_after=None, n=600):
        lines, i = [], 0
        t = T0
        for i in range(n):
            if i in skips:
                t += P
                continue
            lines.append(f"VSYNC ftn={t} mon={t+200_000}")
            m = t + off(i)
            lines.append(f"PRES f={i} mon={m}")
            if i in doubles:
                lines.append(f"PRES f={i}b mon={m+300_000}")
            t += P
            if stall_after is not None and i == stall_after:
                t += 120_000_000
        path = tempfile.NamedTemporaryFile("w", suffix=".log", delete=False)
        path.write("\n".join(lines))
        path.close()
        return path.name

    cases = [
        ("perfect 120Hz +1ms", lambda i: 1_000_000, {}, set(), None, True,
         dict(drops=0, doubles=0, lates=0, stalls=0, max_spread=0.2)),
        ("3-slot drop at 200", lambda i: 1_000_000, {200, 201, 202}, set(), None, False,
         dict(drops=3, doubles=0, lates=0, stalls=1, max_spread=99)),
        ("double at 300", lambda i: 1_000_000, {}, {300}, None, False,
         dict(drops=0, doubles=1, lates=0, stalls=0, max_spread=0.2)),
        ("late i%7==0 +6ms (past latch)", lambda i: 6_000_000 if i % 7 == 0 else 1_000_000,
         {}, set(), None, False,
         dict(drops=0, doubles=0, lates=1, stalls=0, max_spread=99)),
        ("120ms stall at 400", lambda i: 1_000_000, {}, set(), 400, False,
         dict(drops=14, doubles=0, lates=0, stalls=1, max_spread=99)),
        ("phase swing +/-.6ms", lambda i: 1_000_000 + (i % 5 - 2)*300_000,
         {}, set(), None, False,
         dict(drops=0, doubles=0, lates=0, stalls=0, max_spread=1.25)),
        ("tiny jitter 0.3ms", lambda i: 1_000_000 + (i % 3 - 1)*150_000,
         {}, set(), None, True,
         dict(drops=0, doubles=0, lates=0, stalls=0, max_spread=0.31)),
    ]
    ok_all = True
    for name, off, skips, dbl, stall_after, expect_pass, exp in cases:
        path = gen(name, off, skips, dbl, stall_after)
        import io, contextlib
        buf = io.StringIO()
        with contextlib.redirect_stdout(buf):
            res, _sus = run_file(path, period_ms=None, windows_spec="")
        os.unlink(path)
        txt = buf.getvalue()
        m = re.search(r"dropped=(\d+) doubles=(\d+) late=(\d+) stalls=(\d+).*p10-p90=([\d.]+)", txt)
        if not m:
            print(f"  selftest {name:32s} *** could not parse report ***\n{txt}")
            ok_all = False
            continue
        d, db, la, st, sp = int(m.group(1)), int(m.group(2)), int(m.group(3)), int(m.group(4)), float(m.group(5))
        good = (res == expect_pass) and d >= exp["drops"] and db >= exp["doubles"] \
               and la >= exp["lates"] and st >= exp["stalls"] and sp <= exp["max_spread"]
        ok_all = ok_all and good
        print(f"  selftest {name:32s} expect={'PASS' if expect_pass else 'FAIL':4s} "
              f"got={'PASS' if res else 'FAIL':4s} (drops={d} dbl={db} late={la} stalls={st} "
              f"spread={sp:.2f})  {'OK' if good else '*** MISMATCH ***'}")
    print("self-test:", "PASS" if ok_all else "FAIL")
    return 0 if ok_all else 1

def run_file(path, period_ms, windows_spec):
    lines = open(path, errors="replace").read().splitlines()
    pres, ftns = parse_log(lines)
    if len(pres) < 50:
        print("no usable PRES data"); return False, False
    monots = [m for _, m in pres]
    src, P = None, None
    if ftns:
        P = vsync_period(ftns)
        if P: src = "vsync-edges"
    if period_ms:
        P, src = int(period_ms*1e6), "override"
    if not P:
        P, src = fallback_period(monots), "fallback-fit"
    grid = ("edge", ftns[0]) if (src == "vsync-edges" or (period_ms and ftns)) else \
           ("anchor", anchor_phase(monots, P))
    t0 = monots[0]
    suspect = (grid[0] == "anchor")
    print(f"display period: {P/1e6:.4f} ms  (rate {1e6/(P/1e6):.1f} Hz, {src}; "
          f"grid={'vsync-edge' if grid[0]=='edge' else 'anchor-fallback'}"
          + ("  *** SUSPECT: fitted grid — only stalls are meaningful here "
             "(drop/double need the hardware edge grid)" if suspect else ""))
    wins = explicit_windows(pres, t0, windows_spec) if windows_spec else auto_windows(pres, P)
    if not wins:
        print("no active windows found"); return False, suspect
    return evaluate(pres, P, grid, wins, suspect=suspect, ftns=ftns or None), suspect

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("file", nargs="?")
    ap.add_argument("--windows", default="", help='"startMs,endMs name; ..." (ms from first present)')
    ap.add_argument("--period-ms", type=float, default=None,
                    help="force the display period (phone: 8.3333)")
    ap.add_argument("--selftest", action="store_true")
    args = ap.parse_args()
    if args.selftest:
        sys.exit(selftest())
    if not args.file:
        ap.print_usage(); sys.exit(2)
    ok, suspect = run_file(args.file, args.period_ms, args.windows)
    if not ok:
        sys.exit(1)
    sys.exit(3 if suspect else 0)   # 3 = PASS-but-SUSPECT (fitted grid)

if __name__ == "__main__":
    main()