#!/usr/bin/env python3
"""Render the SH3 ZSCHED fixture with the RTL's fixed-point perspective path.

The existing dump-derived reference is intentionally float-ish: it uses clipped
float vertices and computes S/Q at a reference sample point. This tool answers a
different question: if we mirror the RTL's rounded XYZ2 vertices, Q16.16
gradient setup, top-left edge rule, gs_persp_uv reciprocal, and GEQUAL Z RMW,
does the framebuffer line up with the board/zint dump?
"""
import os
import sys
import glob

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

FBW = 256
FBH = 210
TW = 512
TH = 512
PERSP_FRAC = 12
PSCALE = 4096
S24_MAX = (1 << 23) - 1


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


def load_mem64(path):
    out = []
    with open(path) as f:
        for ln in f:
            s = ln.strip()
            if s and not s.startswith("//"):
                out.append(int(s, 16) & 0xFFFFFFFFFFFFFFFF)
    return out


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


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


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

    def clip(poly, inside, cross):
        out = []
        for i, a in enumerate(poly):
            b = poly[(i + 1) % len(poly)]
            ina, inb = inside(a), inside(b)
            if ina:
                out.append(a)
            if ina != inb:
                out.append(cross(a, b))
        return out

    poly = [dict(x=v["x"], y=v["y"], z=v["z"], s=v["s"], t=v["t"], q=v["q"]) for v in tri]
    poly = clip(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(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(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(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[i], poly[i + 1]) for i in range(1, len(poly) - 1)]


def pack_vertex(v):
    s = round(v["s"] * TW * (1 << PERSP_FRAC) * PSCALE)
    t = round(v["t"] * TH * (1 << PERSP_FRAC) * PSCALE)
    q = round(v["q"] * (1 << PERSP_FRAC) * PSCALE)
    if abs(s) > S24_MAX or abs(t) > S24_MAX:
        raise SystemExit(f"[persp] ST overflow s={s} t={t}")
    if q < 0 or q > 0xFFFFFF:
        raise SystemExit(f"[persp] Q out of RTL 24-bit range q={q}")
    return {
        "x": max(0, min(FBW - 1, int(round(v["x"])))),
        "y": max(0, min(FBH - 1, int(round(v["y"])))),
        "z": max(0, min(0xFFFFFF, int(round(v["z"])))),
        "s": s & 0xFFFFFF,
        "t": t & 0xFFFFFF,
        "q": q & 0xFFFFFF,
    }


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


def top_or_left(ax, ay, bx, by):
    dx = bx - ax
    dy = by - ay
    return (dy > 0) or (dy == 0 and dx > 0)


def trunc_div(num, den):
    if den == 0:
        return 0
    sign = -1 if (num < 0) ^ (den < 0) else 1
    return sign * (abs(num) // abs(den))


def grad_num_dadx(a0, a1, a2, x0, y0, x1, y1, x2, y2):
    return ((a1 - a0) * (y2 - y0) - (a2 - a0) * (y1 - y0)) << 16


def grad_num_dady(a0, a1, a2, x0, y0, x1, y1, x2, y2):
    return ((a2 - a0) * (x1 - x0) - (a1 - a0) * (x2 - x0)) << 16


def interp_wide(base, dadx, dady, x, y, x0, y0):
    step = dadx * (x - x0) + dady * (y - y0)
    return (base + (step >> 16)) & 0xFFFFFF


def interp_z(base, dadx, dady, x, y, x0, y0):
    step = dadx * (x - x0) + dady * (y - y0)
    return max(0, min(0xFFFFFFFF, base + (step >> 16)))


def recip_lut(q, idx_bits=11, scale=24):
    out_max = (1 << (scale + 1)) - 1
    if q <= 0:
        return out_max
    top = idx_bits - 1
    e = q.bit_length() - 1
    m = (q >> (e - top)) if e >= top else (q << (top - e))
    m &= (1 << idx_bits) - 1
    if m == 0:
        return out_max
    r = ((1 << (scale + top)) // m) >> e
    return min(r, out_max)


def persp_uv(s, t, q):
    r = recip_lut(q)
    u = (s * r) >> 24
    v = (t * r) >> 24
    if u > 2047:
        u = 2047
    if v > 2047:
        v = 2047
    return u, v


def texel(idx_words, pal, u, v):
    lin = v * TW + u
    if lin < 0 or lin >= len(idx_words) * 4:
        return 0
    ci = (idx_words[lin // 4] >> (8 * (lin % 4))) & 0xFF
    return pal[ci] & 0xFFFFFF


def prep_tri(tri):
    verts = [pack_vertex(v) for v in tri]
    x0, y0 = verts[0]["x"], verts[0]["y"]
    x1, y1 = verts[1]["x"], verts[1]["y"]
    x2, y2 = verts[2]["x"], verts[2]["y"]
    det = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
    if det == 0:
        return None
    if det < 0:
        verts = [verts[0], verts[2], verts[1]]
        det = -det
    v0, v1, v2 = verts
    x0, y0 = v0["x"], v0["y"]
    x1, y1 = v1["x"], v1["y"]
    x2, y2 = v2["x"], v2["y"]
    bias = [
        0 if top_or_left(x0, y0, x1, y1) else 1,
        0 if top_or_left(x1, y1, x2, y2) else 1,
        0 if top_or_left(x2, y2, x0, y0) else 1,
    ]

    def grad(attr):
        a0, a1, a2 = v0[attr], v1[attr], v2[attr]
        return (
            trunc_div(grad_num_dadx(a0, a1, a2, x0, y0, x1, y1, x2, y2), det),
            trunc_div(grad_num_dady(a0, a1, a2, x0, y0, x1, y1, x2, y2), det),
        )

    return {
        "v": verts,
        "det": det,
        "bias": bias,
        "ds": grad("s"),
        "dt": grad("t"),
        "dq": grad("q"),
        "dz": grad("z"),
    }


def prep_tri_packed(verts):
    verts = [dict(v) for v in verts]
    x0, y0 = verts[0]["x"], verts[0]["y"]
    x1, y1 = verts[1]["x"], verts[1]["y"]
    x2, y2 = verts[2]["x"], verts[2]["y"]
    det = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
    if det == 0:
        return None
    if det < 0:
        verts = [verts[0], verts[2], verts[1]]
        det = -det
    v0, v1, v2 = verts
    x0, y0 = v0["x"], v0["y"]
    x1, y1 = v1["x"], v1["y"]
    x2, y2 = v2["x"], v2["y"]
    bias = [
        0 if top_or_left(x0, y0, x1, y1) else 1,
        0 if top_or_left(x1, y1, x2, y2) else 1,
        0 if top_or_left(x2, y2, x0, y0) else 1,
    ]

    def grad(attr):
        a0, a1, a2 = v0[attr], v1[attr], v2[attr]
        return (
            trunc_div(grad_num_dadx(a0, a1, a2, x0, y0, x1, y1, x2, y2), det),
            trunc_div(grad_num_dady(a0, a1, a2, x0, y0, x1, y1, x2, y2), det),
        )

    return {
        "v": verts,
        "det": det,
        "bias": bias,
        "ds": grad("s"),
        "dt": grad("t"),
        "dq": grad("q"),
        "dz": grad("z"),
    }


def build_tris(draw_idxs):
    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:
        rawv = [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 = [(rawv[i - 2], rawv[i - 1], rawv[i]) for i in range(2, len(rawv))]
        tris = []
        for tri in raw:
            for ct in clip_rect_z(tri, FBW, FBH):
                qs = [v["q"] for v in ct]
                if all(q < 0 for q in qs):
                    ct = tuple(dict(x=v["x"], y=v["y"], z=v["z"], s=-v["s"], t=-v["t"], q=-v["q"]) for v in ct)
                elif not all(q > 0 for q in qs):
                    raise SystemExit(f"[persp] mixed/zero Q in idx{e['first_idx']}: {qs}")
                pt = prep_tri(ct)
                if pt is not None:
                    tris.append(pt)
        out.append(tris)
    return out


def build_tris_from_feeder(tag="zsched"):
    out = []
    ep = 0
    while True:
        path = os.path.join(DATA, f"feeder_sh3_{tag}{ep}.mem")
        if not os.path.exists(path):
            break
        words = load_mem64(path)
        ntris = words[0] & 0xFFFF
        tris = []
        off = 7
        for _ in range(ntris):
            verts = []
            for _v in range(3):
                rgbaq = words[off]
                st = words[off + 1]
                xyz = words[off + 2]
                verts.append({
                    "x": (xyz >> 4) & 0xFFF,
                    "y": (xyz >> 20) & 0xFFF,
                    "z": (xyz >> 32) & 0xFFFFFFFF,
                    "s": st & 0xFFFFFF,
                    "t": (st >> 32) & 0xFFFFFF,
                    "q": (rgbaq >> 32) & 0xFFFFFF,
                })
                off += 3
            pt = prep_tri_packed(verts)
            if pt is not None:
                tris.append(pt)
        out.append(tris)
        ep += 1
    return out


def render(tag="zsched"):
    draw_idxs = load_epochs(tag)
    idx = [load_mem(os.path.join(DATA, f"sh3_{tag}{e}_idx.mem")) for e in range(len(draw_idxs))]
    pal = [load_mem(os.path.join(DATA, f"sh3_{tag}{e}_pal.mem")) for e in range(len(draw_idxs))]
    tris_by_ep = build_tris_from_feeder(tag)
    if len(tris_by_ep) != len(draw_idxs):
        tris_by_ep = build_tris(draw_idxs)

    fb = [0] * (FBW * FBH)
    zbuf = [-1] * (FBW * FBH)
    owner = [-1] * (FBW * FBH)
    emitted = [0] * len(draw_idxs)
    accepted = [0] * len(draw_idxs)
    frags = []
    for ep, tris in enumerate(tris_by_ep):
        for tri in tris:
            v0, v1, v2 = tri["v"]
            minx = max(0, min(v0["x"], v1["x"], v2["x"]))
            maxx = min(FBW - 1, max(v0["x"], v1["x"], v2["x"]))
            miny = max(0, min(v0["y"], v1["y"], v2["y"]))
            maxy = min(FBH - 1, max(v0["y"], v1["y"], v2["y"]))
            for y in range(miny, maxy + 1):
                for x in range(minx, maxx + 1):
                    e0 = edge(x, y, v0["x"], v0["y"], v1["x"], v1["y"]) + tri["bias"][0]
                    e1 = edge(x, y, v1["x"], v1["y"], v2["x"], v2["y"]) + tri["bias"][1]
                    e2 = edge(x, y, v2["x"], v2["y"], v0["x"], v0["y"]) + tri["bias"][2]
                    if e0 > 0 or e1 > 0 or e2 > 0:
                        continue
                    s = interp_wide(v0["s"], tri["ds"][0], tri["ds"][1], x, y, v0["x"], v0["y"])
                    t = interp_wide(v0["t"], tri["dt"][0], tri["dt"][1], x, y, v0["x"], v0["y"])
                    q = interp_wide(v0["q"], tri["dq"][0], tri["dq"][1], x, y, v0["x"], v0["y"])
                    z = interp_z(v0["z"], tri["dz"][0], tri["dz"][1], x, y, v0["x"], v0["y"])
                    u, v = persp_uv(s, t, q)
                    col = texel(idx[ep], pal[ep], u, v)
                    o = y * FBW + x
                    emitted[ep] += 1
                    frags.append((ep, x, y, z, col, u, v, s, t, q))
                    zq = min(z, 0xFFFF)
                    if zq >= zbuf[o]:
                        zbuf[o] = zq
                        fb[o] = col
                        owner[o] = ep
                        accepted[ep] += 1
    return draw_idxs, tris_by_ep, fb, owner, emitted, accepted, frags


def nearest_color(idx_words, pal, u, v, color, max_radius=64):
    for rad in range(max_radius + 1):
        for dv in range(-rad, rad + 1):
            for du in range(-rad, rad + 1):
                if max(abs(du), abs(dv)) != rad:
                    continue
                uu = u + du
                vv = v + dv
                if 0 <= uu < TW and 0 <= vv < TH and texel(idx_words, pal, uu, vv) == color:
                    return rad, du, dv
    return None


def compare_trace(frags, path, idx, pal):
    if not os.path.exists(path):
        return
    got = []
    with open(path) as f:
        for ln in f:
            p = ln.split()
            if len(p) >= 5:
                got.append((int(p[0]), int(p[1]), int(p[2]), int(p[3]), int(p[4], 16) & 0xFFFFFF))
    n = min(len(frags), len(got))
    coord_miss = color_miss = z_miss256 = 0
    near_bins = {"r1": 0, "r8": 0, "r32": 0, "r64": 0, "miss": 0}
    by_ep = {}
    examples = []
    for i in range(n):
        me = frags[i]
        hw = got[i]
        ep = hw[0]
        st = by_ep.setdefault(ep, {"n": 0, "coord": 0, "color": 0, "z256": 0})
        st["n"] += 1
        if me[:3] != hw[:3]:
            coord_miss += 1
            st["coord"] += 1
        if (me[4] & 0xFFFFFF) != hw[4]:
            color_miss += 1
            st["color"] += 1
            near = nearest_color(idx[ep], pal[ep], me[5], me[6], hw[4])
            if near is None:
                near_bins["miss"] += 1
            else:
                rad, _du, _dv = near
                if rad <= 1:
                    near_bins["r1"] += 1
                if rad <= 8:
                    near_bins["r8"] += 1
                if rad <= 32:
                    near_bins["r32"] += 1
                if rad <= 64:
                    near_bins["r64"] += 1
        if abs(me[3] - hw[3]) > 256:
            z_miss256 += 1
            st["z256"] += 1
        if len(examples) < 12 and (me[:3] != hw[:3] or (me[4] & 0xFFFFFF) != hw[4] or abs(me[3] - hw[3]) > 256):
            near = nearest_color(idx[ep], pal[ep], me[5], me[6], hw[4]) if (me[4] & 0xFFFFFF) != hw[4] else None
            examples.append((i, me, hw, near))
    print(f"[persp] trace={path} model_frags={len(frags)} trace_frags={len(got)} compared={n}")
    print(f"[persp] trace_miss coord={coord_miss} color={color_miss} z_gt256={z_miss256}")
    print(f"[persp] trace color-near model_uv: <=1 {near_bins['r1']}  <=8 {near_bins['r8']}  <=32 {near_bins['r32']}  <=64 {near_bins['r64']}  >64/notfound {near_bins['miss']}")
    for ep in sorted(by_ep):
        st = by_ep[ep]
        print(f"[persp] trace e{ep}: n={st['n']} coord={st['coord']} color={st['color']} z_gt256={st['z256']}")
    for i, me, hw, near in examples:
        ns = "near=none" if near is None else f"near r={near[0]} du={near[1]} dv={near[2]}"
        print(f"[persp] trace miss#{i}: model ep{me[0]} ({me[1]},{me[2]}) z={me[3]} uv=({me[5]},{me[6]}) col={me[4]&0xFFFFFF:06x}  "
              f"trace ep{hw[0]} ({hw[1]},{hw[2]}) z={hw[3]} col={hw[4]:06x} {ns}")


def compare_issue(frags, path):
    if not os.path.exists(path):
        return
    got = []
    with open(path) as f:
        for ln in f:
            p = ln.split()
            if len(p) >= 7:
                got.append((int(p[0]), int(p[1]), int(p[2]), int(p[3]), int(p[4]), int(p[5]), int(p[6])))
    n = min(len(frags), len(got))
    by_ep = {}
    s_miss = t_miss = u_miss = v_miss = recip_miss = 0
    examples = []
    for i in range(n):
        me = frags[i]
        hw = got[i]
        ep = hw[0]
        st = by_ep.setdefault(ep, {"n": 0, "s": 0, "t": 0, "u": 0, "v": 0, "recip": 0})
        st["n"] += 1
        mr = recip_lut(me[9])
        if me[7] != hw[1]:
            s_miss += 1
            st["s"] += 1
        if me[8] != hw[2]:
            t_miss += 1
            st["t"] += 1
        if mr != hw[3]:
            recip_miss += 1
            st["recip"] += 1
        if me[5] != hw[4]:
            u_miss += 1
            st["u"] += 1
        if me[6] != hw[5]:
            v_miss += 1
            st["v"] += 1
        if len(examples) < 12 and (me[7] != hw[1] or me[8] != hw[2] or mr != hw[3] or me[5] != hw[4] or me[6] != hw[5]):
            examples.append((i, me, hw, mr))
    print(f"[persp] issue={path} model_frags={len(frags)} issue_rows={len(got)} compared={n}")
    print(f"[persp] issue_miss s={s_miss} t={t_miss} recip={recip_miss} u={u_miss} v={v_miss}")
    for ep in sorted(by_ep):
        st = by_ep[ep]
        print(f"[persp] issue e{ep}: n={st['n']} s={st['s']} t={st['t']} recip={st['recip']} u={st['u']} v={st['v']}")
    for i, me, hw, mr in examples:
        print(f"[persp] issue miss#{i}: model ep{me[0]} ({me[1]},{me[2]}) s={me[7]} t={me[8]} q={me[9]} recip={mr} uv=({me[5]},{me[6]})  "
              f"rtl ep{hw[0]} s={hw[1]} t={hw[2]} recip={hw[3]} uv=({hw[4]},{hw[5]}) valid={hw[6]}")


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")
    board = load_mem(fb_path)
    trace_path = args[args.index("--trace") + 1] if "--trace" in args else os.path.join(ROOT, "sim", "traces", "rtl", "zsched_frags.txt")
    draw_idxs, tris_by_ep, fb, owner, emitted, accepted, frags = render()
    if len(board) != len(fb):
        raise SystemExit(f"[persp] {fb_path}: {len(board)} words != expected {len(fb)}")
    mism = []
    cov = sum(1 for x in owner if x >= 0)
    for i, (a, b) in enumerate(zip(board, fb)):
        if (a & 0xFFFFFF) != (b & 0xFFFFFF):
            mism.append(i)
    by_owner = [0] * len(draw_idxs)
    by_miss = [0] * len(draw_idxs)
    for o in range(len(fb)):
        ep = owner[o]
        if ep >= 0:
            by_owner[ep] += 1
    miss_set = set(mism)
    for o in miss_set:
        ep = owner[o]
        if ep >= 0:
            by_miss[ep] += 1

    print(f"[persp] fb={fb_path}")
    print(f"[persp] draw_idxs={draw_idxs} FB={FBW}x{FBH} covered={cov} mismatches={len(mism)}/{len(fb)}")
    for ep, idx in enumerate(draw_idxs):
        print(f"[persp] e{ep} idx{idx}: tris={len(tris_by_ep[ep])} emitted={emitted[ep]} accepted={accepted[ep]} final_owner={by_owner[ep]} miss={by_miss[ep]}")
    for o in mism[:12]:
        print(f"[persp] miss x={o % FBW} y={o // FBW} owner=e{owner[o]} model={fb[o] & 0xFFFFFF:06x} board={board[o] & 0xFFFFFF:06x}")
    trace_idxs = [load_mem(os.path.join(DATA, f"sh3_zsched{e}_idx.mem")) for e in range(len(draw_idxs))]
    trace_pals = [load_mem(os.path.join(DATA, f"sh3_zsched{e}_pal.mem")) for e in range(len(draw_idxs))]
    compare_trace(frags, trace_path, trace_idxs, trace_pals)
    compare_issue(frags, os.path.join(ROOT, "sim", "traces", "rtl", "zsched_issue.txt"))
    return 0 if not mism else 1


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