#!/usr/bin/env python3
"""Per-frame quad geometry and vertex colour from a `log_ui_draws` capture.

`ui_draw_order.py` answers "in what order were these painted"; this answers
"what did this quad look like on frame N, and on N+1". That is what a keyframe
TIME unit has to be measured against: the bundle says an element ramps from
t=31 to t=34, and the only way to learn what a `t` is worth is to count the
rendered frames the same ramp takes in the running game.

    kf_time_probe.py <capture.log> [--csv out.csv]

Emits one row per quad per frame: frame, draw index, pixel rect, whether the
quad is axis-aligned, and the per-vertex colour word (whose alpha byte is the
element's fade). Frame numbers are the emulator's VdSwap count, so they are
SUBMITTED FRAMES, not wall-clock — which is the point: an emulator that runs
at half speed does not move them.
"""
import re
import sys

W, H = 1280, 720
VERT = re.compile(r"\[(-?\d+\.\d+),(-?\d+\.\d+),z=(-?\d+\.\d+)(?:,col=([0-9A-F]{8}))?\]")


def quads(log):
    """Yield (frame, draw, x0, y0, w, h, rot, col) for every quad in the log."""
    frame, pending = None, None
    for line in open(log):
        if line.startswith("# every draw"):
            frame = int(line.rstrip().split()[-1].split("..")[0])
            continue
        if line.startswith("--- frame"):
            frame = int(line.split()[2])
            continue
        m = re.match(r"\s*(\d+) (prim=.*)", line)
        if m:
            pending = (int(m.group(1)), m.group(2))
            continue
        if "vb=0x" in line and pending:
            hits = VERT.findall(line)
            verts = [(float(a), float(b), c) for a, b, _z, c in hits]
            if verts:
                # Two conventions in one log: the UI sprite shader emits NDC,
                # the full-screen pass emits pixels already.
                if max(abs(v) for x, y, _c in verts for v in (x, y)) > 4.0:
                    pts = [(x, y, c) for x, y, c in verts]
                else:
                    pts = [((x + 1) / 2 * W, (1 - y) / 2 * H, c) for x, y, c in verts]
                per = 4 if "prim=13" in pending[1] else len(pts)
                for q in range(0, len(pts), per):
                    chunk = pts[q:q + per]
                    if not chunk:
                        continue
                    xs = [p[0] for p in chunk]
                    ys = [p[1] for p in chunk]
                    x0, x1, y0, y1 = min(xs), max(xs), min(ys), max(ys)
                    rot = all(
                        (abs(x - x0) < 1.0 or abs(x - x1) < 1.0)
                        and (abs(y - y0) < 1.0 or abs(y - y1) < 1.0)
                        for x, y in zip(xs, ys)
                    )
                    cols = {p[2] for p in chunk if p[2]}
                    col = sorted(cols)[0] if len(cols) == 1 else (
                        "/".join(sorted(cols)) if cols else "")
                    yield (frame, pending[0], round(x0), round(y0),
                           round(x1 - x0), round(y1 - y0), "" if rot else "ROT",
                           col)
            pending = None


def main():
    log = sys.argv[1]
    out = None
    if "--csv" in sys.argv:
        out = open(sys.argv[sys.argv.index("--csv") + 1], "w")
        out.write("frame,draw,x,y,w,h,rot,col\n")
    for row in quads(log):
        line = "%d,%d,%d,%d,%d,%d,%s,%s" % row
        if out:
            out.write(line + "\n")
        else:
            print(line)
    if out:
        out.close()
    return 0


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