#!/usr/bin/env python3
"""Score a Ch357 ZSCHED framebuffer dump.

The z/paint/replay owner modes compare against the dump-derived texture oracle.
The replay-color mode is the hard board/RTL content check: it replays the emitted
fragment trace colors through the same clamp16 GEQUAL persistent-Z rule and
compares the framebuffer colors directly.

Usage:
  tools/analyze_zsched_fb.py sim/data/top_psmct32_raster_demo/sh3_zsched_board_fb.mem
  tools/analyze_zsched_fb.py <fb.mem> --owner paint
  tools/analyze_zsched_fb.py <fb.mem> --owner replay --frags sim/traces/rtl/zsched_frags.txt
  tools/analyze_zsched_fb.py <fb.mem> --owner replay-color --frags sim/traces/rtl/zsched_frags.txt
  tools/analyze_zsched_fb.py <fb.mem> --owner replay --radius 8
  tools/analyze_zsched_fb.py <fb.mem> --maps
"""
import os
import re
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.normpath(os.path.join(HERE, ".."))
DATA = os.path.join(ROOT, "sim", "data", "top_psmct32_raster_demo")


def load_mem(path):
    vals = []
    with open(path) as f:
        for ln in f:
            s = ln.strip()
            if not s or s.startswith("//"):
                continue
            vals.append(int(s, 16) & 0xFFFFFFFF)
    return vals


def parse_params(tag):
    path = os.path.join(DATA, f"sh3_{tag}_params.vh")
    txt = open(path).read()
    out = {}
    for key in ("FBPXW", "FBH", "TW", "TH", "N_EPOCHS"):
        m = re.search(rf"localparam\s+int\s+{key}\s*=\s*([0-9]+)\s*;", txt)
        if not m:
            raise SystemExit(f"[analyze] missing {key} in {path}")
        out[key] = int(m.group(1))
    return out


def cell(idx, pal, tw, u, v):
    lin = v * tw + u
    w = idx[lin // 4]
    ix = (w >> (8 * (lin % 4))) & 0xFF
    return pal[ix] & 0xFFFFFF


def match_radius(fbc, idx, pal, tw, th, tu, tv, max_radius):
    for rad in range(max_radius + 1):
        for du in range(-rad, rad + 1):
            for dv in range(-rad, rad + 1):
                if max(abs(du), abs(dv)) != rad:
                    continue
                uu = tu + du
                vv = tv + dv
                if 0 <= uu < tw and 0 <= vv < th and fbc == cell(idx, pal, tw, uu, vv):
                    return rad
    return None


def write_maps(prefix, w, h, fb, status):
    try:
        from PIL import Image
    except Exception:
        print("[analyze] PIL not available; skipping maps")
        return

    colors = {
        0: (0, 0, 0),       # uncovered
        1: (0, 180, 0),     # owner exact
        2: (170, 170, 0),   # owner <=1 texel
        3: (0, 80, 220),    # wrong owner/other covering epoch
        4: (220, 0, 0),     # no covering epoch match
        5: (255, 0, 255),   # no palette
    }
    img = Image.new("RGB", (w, h))
    for o, st in enumerate(status):
        img.putpixel((o % w, o // w), colors.get(st, (255, 255, 255)))
    out = f"{prefix}_class.png"
    img.resize((w * 3, h * 3), Image.NEAREST).save(out)
    print(f"[analyze] wrote {out}")

    diff = Image.new("RGB", (w, h))
    for o, px in enumerate(fb[:w * h]):
        if status[o] in (4, 5):
            diff.putpixel((o % w, o // w), (255, 0, 0))
        else:
            diff.putpixel((o % w, o // w), (px & 0xFF, (px >> 8) & 0xFF, (px >> 16) & 0xFF))
    out = f"{prefix}_bad_overlay.png"
    diff.resize((w * 3, h * 3), Image.NEAREST).save(out)
    print(f"[analyze] wrote {out}")


def replay_owner(frags, npx, w):
    zbuf = [-1] * npx
    z_owner = [-1] * npx
    paint_owner = [-1] * npx
    cov = [0] * npx
    with open(frags) as f:
        for ln in f:
            p = ln.split()
            if len(p) < 5:
                continue
            ep = int(p[0])
            x = int(p[1])
            y = int(p[2])
            z = int(p[3])
            o = y * w + x
            if o < 0 or o >= npx:
                continue
            zq = 0xFFFF if z > 0xFFFF else (0 if z < 0 else z)
            cov[o] = 1
            paint_owner[o] = ep
            if zq >= zbuf[o]:
                zbuf[o] = zq
                z_owner[o] = ep
    return cov, z_owner, paint_owner, zbuf


def replay_color(frags, npx, w):
    zbuf = [-1] * npx
    cov = [0] * npx
    col = [0] * npx
    frag_count = 0
    pass_count = 0
    with open(frags) as f:
        for ln in f:
            p = ln.split()
            if len(p) < 5:
                continue
            x = int(p[1])
            y = int(p[2])
            z = int(p[3])
            c = int(p[4], 16) & 0xFFFFFF
            o = y * w + x
            if o < 0 or o >= npx:
                continue
            frag_count += 1
            zq = 0xFFFF if z > 0xFFFF else (0 if z < 0 else z)
            if zq >= zbuf[o]:
                zbuf[o] = zq
                cov[o] = 1
                col[o] = c
                pass_count += 1
    return cov, col, frag_count, pass_count


def replay_color_check(fb_path, fb, frags, w, h):
    npx = w * h
    cov, exp, frag_count, pass_count = replay_color(frags, npx, w)
    covered = sum(cov)
    mismatches = []
    stray = 0
    for o in range(npx):
        got = fb[o] & 0xFFFFFF
        want = exp[o] & 0xFFFFFF
        if cov[o]:
            if got != want:
                if len(mismatches) < 12:
                    mismatches.append((o, got, want, "covered"))
        elif got != 0:
            stray += 1
            if len(mismatches) < 12:
                mismatches.append((o, got, want, "uncovered"))

    mismatch_count = 0
    for o in range(npx):
        got = fb[o] & 0xFFFFFF
        want = exp[o] & 0xFFFFFF
        if (cov[o] and got != want) or ((not cov[o]) and got != 0):
            mismatch_count += 1

    print(f"[analyze] fb={fb_path}")
    print(f"[analyze] owner=replay-color FB={w}x{h}")
    print(f"[analyze] replay_frags={frags}")
    print(f"[analyze] replay fragments={frag_count} zpass_updates={pass_count} covered={covered} stray_uncovered={stray}")
    print(f"[analyze] replay-color exact {npx - mismatch_count}/{npx} ({100.0 * (npx - mismatch_count) / npx:.2f}%) mismatches={mismatch_count}")
    for o, got, want, kind in mismatches:
        print(f"[analyze] bad {kind}: px={o} x={o % w} y={o // w} got={got:06x} exp={want:06x}")
    return 0 if mismatch_count == 0 else 1


def main(argv):
    args = argv[1:]
    fb_path = args[0] if args and not args[0].startswith("--") else os.path.join(DATA, "sh3_zsched_board_fb.mem")
    tag = args[args.index("--tag") + 1] if "--tag" in args else "zsched"
    owner_mode = args[args.index("--owner") + 1] if "--owner" in args else "z"
    frags = args[args.index("--frags") + 1] if "--frags" in args else os.path.join(ROOT, "sim", "traces", "rtl", "zsched_frags.txt")
    max_radius = int(args[args.index("--radius") + 1]) if "--radius" in args else 1
    emit_maps = "--maps" in args
    if owner_mode not in ("z", "paint", "replay", "replay-color"):
        raise SystemExit("[analyze] --owner must be z, paint, replay, or replay-color")

    p = parse_params(tag)
    w, h, tw, th, ne = p["FBPXW"], p["FBH"], p["TW"], p["TH"], p["N_EPOCHS"]
    npx = w * h

    fb = load_mem(fb_path)
    if len(fb) != npx:
        raise SystemExit(f"[analyze] {fb_path}: {len(fb)} words != expected {npx}")

    if owner_mode == "replay-color":
        return replay_color_check(fb_path, fb, frags, w, h)

    idx = [load_mem(os.path.join(DATA, f"sh3_{tag}{e}_idx.mem")) for e in range(ne)]
    pal = [load_mem(os.path.join(DATA, f"sh3_{tag}{e}_pal.mem")) for e in range(ne)]
    ref = load_mem(os.path.join(DATA, f"sh3_{tag}_refmap.mem"))
    rep = [load_mem(os.path.join(DATA, f"sh3_{tag}{e}_refmap.mem")) for e in range(ne)]
    zowner_path = os.path.join(DATA, f"sh3_{tag}_zowner.mem")
    zowner = load_mem(zowner_path) if os.path.exists(zowner_path) else None
    if owner_mode == "z" and zowner is None:
        raise SystemExit(f"[analyze] missing {zowner_path}; use --owner paint or regenerate z oracle")
    replay = replay_owner(frags, npx, w) if owner_mode == "replay" else None

    palette_set = set()
    for e in range(ne):
        palette_set.update(x & 0xFFFFFF for x in pal[e])

    totals = {
        "covered": 0, "interior": 0, "multi": 0, "reject": 0,
        "owner_r0": 0, "owner_r1": 0, "any_r1": 0,
        "interior_owner_r1": 0, "multi_owner_r1": 0, "reject_owner_r1": 0,
        "palette_bad": 0, "owner_missing_ref": 0,
    }
    status = [0] * npx
    owner_counts = [0] * ne
    by_owner = [{"cov": 0, "r0": 0, "r1": 0, "multi": 0, "multi_r1": 0, "reject": 0, "reject_r1": 0} for _ in range(ne)]
    radius_hist = [0] * (max_radius + 1)
    bad_examples = []

    for o in range(npx):
        if owner_mode == "z":
            zw = zowner[o]
            cov = (zw >> 31) & 1
            if not cov:
                continue
            owner = ((zw >> 24) & 0xF) - 1
            paint_owner = (zw & 0xF) - 1
            reject = owner != paint_owner
            multi = (ref[o] >> 28) & 1
            interior = (rep[owner][o] >> 30) & 1 if 0 <= owner < ne else 0
        elif owner_mode == "replay":
            covs, owners, paint_owners, _zq = replay
            if not covs[o]:
                continue
            owner = owners[o]
            paint_owner = paint_owners[o]
            reject = owner != paint_owner
            multi = (ref[o] >> 28) & 1
            interior = (rep[owner][o] >> 30) & 1 if 0 <= owner < ne else 0
        else:
            rw = ref[o]
            cov = (rw >> 31) & 1
            if not cov:
                continue
            owner = (rw >> 24) & 0x7
            reject = False
            multi = (rw >> 28) & 1
            interior = (rw >> 30) & 1

        totals["covered"] += 1
        totals["interior"] += int(interior)
        totals["multi"] += int(multi)
        totals["reject"] += int(reject)
        if 0 <= owner < ne:
            owner_counts[owner] += 1
            by_owner[owner]["cov"] += 1
            by_owner[owner]["multi"] += int(multi)
            by_owner[owner]["reject"] += int(reject)

        fbc = fb[o] & 0xFFFFFF
        if fbc not in palette_set:
            totals["palette_bad"] += 1
            status[o] = 5
            if len(bad_examples) < 12:
                bad_examples.append((o, owner, fbc, "palette"))
            continue

        owner_rad = None
        if 0 <= owner < ne and (rep[owner][o] >> 31):
            rm = rep[owner][o]
            tu = (rm >> 9) & 0x1FF
            tv = rm & 0x1FF
            owner_rad = match_radius(fbc, idx[owner], pal[owner], tw, th, tu, tv, max_radius)
        else:
            totals["owner_missing_ref"] += 1

        if owner_rad == 0:
            radius_hist[0] += 1
            totals["owner_r0"] += 1
            totals["owner_r1"] += 1
            totals["interior_owner_r1"] += int(interior)
            totals["multi_owner_r1"] += int(multi)
            totals["reject_owner_r1"] += int(reject)
            if 0 <= owner < ne:
                by_owner[owner]["r0"] += 1
                by_owner[owner]["r1"] += 1
                by_owner[owner]["multi_r1"] += int(multi)
                by_owner[owner]["reject_r1"] += int(reject)
            status[o] = 1
            continue
        if owner_rad == 1:
            radius_hist[1] += 1
            totals["owner_r1"] += 1
            totals["interior_owner_r1"] += int(interior)
            totals["multi_owner_r1"] += int(multi)
            totals["reject_owner_r1"] += int(reject)
            if 0 <= owner < ne:
                by_owner[owner]["r1"] += 1
                by_owner[owner]["multi_r1"] += int(multi)
                by_owner[owner]["reject_r1"] += int(reject)
            status[o] = 2
            continue

        if owner_rad is not None:
            radius_hist[owner_rad] += 1
            totals["owner_r1"] += 1
            totals["interior_owner_r1"] += int(interior)
            totals["multi_owner_r1"] += int(multi)
            totals["reject_owner_r1"] += int(reject)
            if 0 <= owner < ne:
                by_owner[owner]["r1"] += 1
                by_owner[owner]["multi_r1"] += int(multi)
                by_owner[owner]["reject_r1"] += int(reject)
            status[o] = 2
            continue

        any_ok = False
        for e in range(ne):
            if not (rep[e][o] >> 31):
                continue
            rm = rep[e][o]
            tu = (rm >> 9) & 0x1FF
            tv = rm & 0x1FF
            if match_radius(fbc, idx[e], pal[e], tw, th, tu, tv, max_radius) is not None:
                any_ok = True
                break
        if any_ok:
            totals["any_r1"] += 1
            status[o] = 3
        else:
            status[o] = 4
            if len(bad_examples) < 12:
                bad_examples.append((o, owner, fbc, "nomatch"))

    def pct(num, den):
        return 100.0 * num / den if den else 0.0

    print(f"[analyze] fb={fb_path}")
    print(f"[analyze] tag={tag} owner={owner_mode} FB={w}x{h} epochs={ne} radius={max_radius}")
    if owner_mode == "replay":
        print(f"[analyze] replay_frags={frags}")
    print(f"[analyze] oracle covered={totals['covered']} interior={totals['interior']} multi={totals['multi']} reject={totals['reject']}")
    print("[analyze] owner counts: " + " ".join(f"e{e}={owner_counts[e]}" for e in range(ne)))
    for e in range(ne):
        b = by_owner[e]
        rlbl = "<=1" if max_radius == 1 else f"<=R{max_radius}"
        print(f"[analyze] e{e}: {rlbl} {b['r1']}/{b['cov']} ({pct(b['r1'], b['cov']):.2f}%)  "
              f"multi{rlbl} {b['multi_r1']}/{b['multi']} ({pct(b['multi_r1'], b['multi']):.2f}%)  "
              f"reject{rlbl} {b['reject_r1']}/{b['reject']} ({pct(b['reject_r1'], b['reject']):.2f}%)")
    print(f"[analyze] owner exact      {totals['owner_r0']}/{totals['covered']} ({pct(totals['owner_r0'], totals['covered']):.2f}%)")
    if max_radius == 1:
        print(f"[analyze] owner <=1 texel  {totals['owner_r1']}/{totals['covered']} ({pct(totals['owner_r1'], totals['covered']):.2f}%)")
        print(f"[analyze] interior <=1     {totals['interior_owner_r1']}/{totals['interior']} ({pct(totals['interior_owner_r1'], totals['interior']):.2f}%)")
        print(f"[analyze] multi <=1        {totals['multi_owner_r1']}/{totals['multi']} ({pct(totals['multi_owner_r1'], totals['multi']):.2f}%)")
    else:
        print(f"[analyze] owner <=R texel  {totals['owner_r1']}/{totals['covered']} ({pct(totals['owner_r1'], totals['covered']):.2f}%)")
        print(f"[analyze] interior <=R     {totals['interior_owner_r1']}/{totals['interior']} ({pct(totals['interior_owner_r1'], totals['interior']):.2f}%)")
        print(f"[analyze] multi <=R        {totals['multi_owner_r1']}/{totals['multi']} ({pct(totals['multi_owner_r1'], totals['multi']):.2f}%)")
        print("[analyze] owner radius hist: " + " ".join(f"r{r}={radius_hist[r]}" for r in range(max_radius + 1)))
    if owner_mode == "z":
        print(f"[analyze] reject <=1       {totals['reject_owner_r1']}/{totals['reject']} ({pct(totals['reject_owner_r1'], totals['reject']):.2f}%)")
    print(f"[analyze] wrong-owner <=1  {totals['any_r1']}/{totals['covered']} ({pct(totals['any_r1'], totals['covered']):.2f}%)")
    print(f"[analyze] palette_bad={totals['palette_bad']} owner_missing_ref={totals['owner_missing_ref']}")
    if bad_examples:
        for o, owner, fbc, kind in bad_examples:
            print(f"[analyze] bad {kind}: px={o} x={o % w} y={o // w} owner=e{owner} fb={fbc:06x}")

    if emit_maps:
        prefix = os.path.splitext(fb_path)[0] + f"_{owner_mode}oracle"
        write_maps(prefix, w, h, fb, status)

    return 0 if totals["palette_bad"] == 0 else 1


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))