#!/usr/bin/env python3
"""Compare captured RTL Z fragments against the fixture's software Z planes.

This is intentionally per-fragment, not per-final-frame. It reconstructs the
same clipped triangles used by the ZSCHED feeder fixture, then asks whether each
captured RTL fragment's Z is close to the expected screen-linear interpolated Z
at that pixel.

Usage:
  tools/diagnose_zsched_frag_z.py [sim/traces/rtl/zsched_frags.txt]
"""
import os
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")
sys.path.insert(0, HERE)

import gs_make_sh3_multidraw_fixture as MD


def load_epochs(tag="zsched"):
    path = os.path.join(DATA, f"sh3_{tag}_epochs.txt")
    out = []
    with open(path) as f:
        for ln in f:
            s = ln.strip()
            if not s or s.startswith("#") or s.startswith("META"):
                continue
            p = s.split()
            out.append(int(p[1]))
    return out


def find_dump():
    import glob
    c = glob.glob(os.path.join(ROOT, "captures", "gs", "silenthill3", "*224139*.gs.zst"))
    if not c:
        raise SystemExit("[fragz] no SH3 224139 dump found")
    return c[0]


def edge(ax, ay, bx, by, px, py):
    return (px - ax) * (by - ay) - (py - ay) * (bx - ax)


def clip_rect_z(tri, w, h):
    def lerp(p1, p2, al):
        return {k: (p1[k] + al * (p2[k] - p1[k])) for k in ("x", "y", "z", "s", "t", "q")}

    def clip_edge(poly, inside, isect):
        out = []
        n = len(poly)
        for i in range(n):
            a = poly[i]
            b = poly[(i + 1) % n]
            ina = inside(a)
            inb = inside(b)
            if ina:
                out.append(a)
            if ina != inb:
                out.append(isect(a, b))
        return out

    poly = [dict(x=v["x"], y=v["y"], z=v.get("z", 0.0), s=v["s"], t=v["t"], q=v["q"]) for v in tri]
    poly = clip_edge(poly, lambda p: p["x"] >= 0.0, lambda a, b: lerp(a, b, (0.0 - a["x"]) / (b["x"] - a["x"])))
    if not poly:
        return []
    poly = clip_edge(poly, lambda p: p["x"] <= w, lambda a, b: lerp(a, b, (w - a["x"]) / (b["x"] - a["x"])))
    if not poly:
        return []
    poly = clip_edge(poly, lambda p: p["y"] >= 0.0, lambda a, b: lerp(a, b, (0.0 - a["y"]) / (b["y"] - a["y"])))
    if not poly:
        return []
    poly = clip_edge(poly, lambda p: p["y"] <= h, lambda a, b: lerp(a, b, (h - a["y"]) / (b["y"] - a["y"])))
    if len(poly) < 3:
        return []
    return [(poly[0], poly[k], poly[k + 1]) for k in range(1, len(poly) - 1)]


def tri_z_at(tri, x, y):
    v0, v1, v2 = tri
    ar = edge(v0["x"], v0["y"], v1["x"], v1["y"], v2["x"], v2["y"])
    if abs(ar) < 1e-9:
        return None
    inv = 1.0 / ar
    w0 = edge(v1["x"], v1["y"], v2["x"], v2["y"], x, y) * inv
    w1 = edge(v2["x"], v2["y"], v0["x"], v0["y"], x, y) * inv
    w2 = 1.0 - w0 - w1
    if w0 < -0.001 or w1 < -0.001 or w2 < -0.001:
        return None
    return w0 * v0["z"] + w1 * v1["z"] + w2 * v2["z"]


def build_epoch_tris(draw_idxs, fbw, fbh):
    got, _ = MD.load_draws(find_dump(), draw_idxs)
    eps = [got[i] for i in draw_idxs]
    ox = int(min(min(v["x"] for v in e["verts"]) for e in eps))
    oy = int(min(min(v["y"] for v in e["verts"]) for e in eps))
    out = []
    for e in eps:
        fv = [dict(x=v["x"] - ox, y=v["y"] - oy, z=v["z"], s=v["s"], t=v["t"], q=v["q"]) for v in e["verts"]]
        raw = [(fv[i - 2], fv[i - 1], fv[i]) for i in range(2, len(fv))]
        tris = []
        for tri in raw:
            tris.extend(clip_rect_z(tri, fbw, fbh))
        out.append(tris)
    return out


def best_expected(tris, x, y, sample_center):
    sx = x + 0.5 if sample_center else float(x)
    sy = y + 0.5 if sample_center else float(y)
    vals = []
    for tri in tris:
        z = tri_z_at(tri, sx, sy)
        if z is not None:
            vals.append(z)
    return vals


def main(argv):
    frags = argv[1] if len(argv) > 1 else os.path.join(ROOT, "sim", "traces", "rtl", "zsched_frags.txt")
    draw_idxs = load_epochs("zsched")
    fbw, fbh = 256, 210
    tris_by_ep = build_epoch_tris(draw_idxs, fbw, fbh)

    stats = {
        "n": 0, "nocov_center": 0, "nocov_corner": 0,
        "center_le0": 0, "center_le1": 0, "center_le16": 0, "center_le256": 0,
        "corner_le0": 0, "corner_le1": 0, "corner_le16": 0, "corner_le256": 0,
    }
    abs_sum_center = 0.0
    abs_sum_corner = 0.0
    ratios = []
    examples = []

    with open(frags) as f:
        for ln in f:
            p = ln.split()
            if len(p) < 5:
                continue
            ep, x, y, rz = int(p[0]), int(p[1]), int(p[2]), int(p[3])
            stats["n"] += 1
            cvals = best_expected(tris_by_ep[ep], x, y, True)
            kvals = best_expected(tris_by_ep[ep], x, y, False)
            if not cvals:
                stats["nocov_center"] += 1
                cdiff = None
            else:
                cz = min(cvals, key=lambda z: abs(rz - z))
                cdiff = abs(rz - cz)
                abs_sum_center += cdiff
                if cz:
                    ratios.append(rz / cz)
                if cdiff <= 0.5:
                    stats["center_le0"] += 1
                if cdiff <= 1:
                    stats["center_le1"] += 1
                if cdiff <= 16:
                    stats["center_le16"] += 1
                if cdiff <= 256:
                    stats["center_le256"] += 1
            if not kvals:
                stats["nocov_corner"] += 1
                kdiff = None
            else:
                kz = min(kvals, key=lambda z: abs(rz - z))
                kdiff = abs(rz - kz)
                abs_sum_corner += kdiff
                if kdiff <= 0.5:
                    stats["corner_le0"] += 1
                if kdiff <= 1:
                    stats["corner_le1"] += 1
                if kdiff <= 16:
                    stats["corner_le16"] += 1
                if kdiff <= 256:
                    stats["corner_le256"] += 1
            if len(examples) < 16 and ((cdiff is None or cdiff > 256) and (kdiff is None or kdiff > 256)):
                examples.append((ep, x, y, rz, cdiff, kdiff, cvals[:3], kvals[:3]))

    n = stats["n"]
    def pct(k):
        return 100.0 * stats[k] / n if n else 0.0

    print(f"[fragz] fragments={n} draw_idxs={draw_idxs}")
    print(f"[fragz] center coverage misses={stats['nocov_center']} corner coverage misses={stats['nocov_corner']}")
    print(f"[fragz] center absdiff: <=0.5 {stats['center_le0']} ({pct('center_le0'):.2f}%)  <=1 {stats['center_le1']} ({pct('center_le1'):.2f}%)  <=16 {stats['center_le16']} ({pct('center_le16'):.2f}%)  <=256 {stats['center_le256']} ({pct('center_le256'):.2f}%)")
    print(f"[fragz] corner absdiff: <=0.5 {stats['corner_le0']} ({pct('corner_le0'):.2f}%)  <=1 {stats['corner_le1']} ({pct('corner_le1'):.2f}%)  <=16 {stats['corner_le16']} ({pct('corner_le16'):.2f}%)  <=256 {stats['corner_le256']} ({pct('corner_le256'):.2f}%)")
    if n - stats["nocov_center"] > 0:
        print(f"[fragz] mean center absdiff={abs_sum_center / (n - stats['nocov_center']):.2f}")
    if n - stats["nocov_corner"] > 0:
        print(f"[fragz] mean corner absdiff={abs_sum_corner / (n - stats['nocov_corner']):.2f}")
    if ratios:
        ratios.sort()
        mid = ratios[len(ratios) // 2]
        print(f"[fragz] raw/expected center ratio: p10={ratios[len(ratios)//10]:.4f} median={mid:.4f} p90={ratios[(len(ratios)*9)//10]:.4f}")
    for ep, x, y, rz, cdiff, kdiff, cvals, kvals in examples:
        print(f"[fragz] bad ep{ep} ({x},{y}) rtl_z={rz} center_diff={cdiff} corner_diff={kdiff} center_z={[round(v,1) for v in cvals]} corner_z={[round(v,1) for v in kvals]}")
    return 0


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