"""Dump OpenVR ground truth for spike S3 — run ONCE per unit, dev-time only.

Requires SteamVR RUNNING with the Beyond connected (it will grab direct
mode; that's fine, this is the one sanctioned SteamVR run). Dumps:

  - ComputeDistortion UV grid per eye per channel (the ground truth)
  - GetProjectionRaw tangents per eye        (cross-check vs intrinsics)
  - GetEyeToHeadTransform per eye            (cross-check vs eye_to_head + ipd)
  - RecommendedRenderTargetSize, HMD serial/model

Output: distortion_dump_<serial>.json in this directory.

Usage:  python dump_openvr.py [--grid N]   (default 65 -> 65x65 per eye)
"""

import argparse
import json
import sys

import openvr

GRID_DEFAULT = 65


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--grid", type=int, default=GRID_DEFAULT)
    args = ap.parse_args()
    n = args.grid

    try:
        openvr.init(openvr.VRApplication_Background)
    except openvr.error_code.InitError as e:
        print(f"OpenVR init failed: {e}")
        print("Start SteamVR with the Beyond connected, then re-run.")
        sys.exit(1)

    vrsys = openvr.VRSystem()
    hmd = openvr.k_unTrackedDeviceIndex_Hmd
    serial = vrsys.getStringTrackedDeviceProperty(hmd, openvr.Prop_SerialNumber_String)
    model = vrsys.getStringTrackedDeviceProperty(hmd, openvr.Prop_ModelNumber_String)
    rt_w, rt_h = vrsys.getRecommendedRenderTargetSize()
    print(f"HMD: {model} {serial}, recommended RT {rt_w}x{rt_h} per eye")

    out = {
        "serial": serial,
        "model": model,
        "recommended_render_target": [rt_w, rt_h],
        "grid": n,
        "eyes": [],
    }

    for eye_name, eye in (("left", openvr.Eye_Left), ("right", openvr.Eye_Right)):
        left, right, top, bottom = vrsys.getProjectionRaw(eye)
        m = vrsys.getEyeToHeadTransform(eye)
        eye_to_head = [[m[r][c] for c in range(4)] for r in range(3)]

        points = []
        bad = 0
        for j in range(n):
            v = j / (n - 1)
            for i in range(n):
                u = i / (n - 1)
                ok, dc = vrsys.computeDistortion(eye, u, v)
                if not ok:
                    bad += 1
                    points.append(None)
                    continue
                points.append({
                    "uv": [u, v],
                    "r": [dc.rfRed[0], dc.rfRed[1]],
                    "g": [dc.rfGreen[0], dc.rfGreen[1]],
                    "b": [dc.rfBlue[0], dc.rfBlue[1]],
                })
        print(f"  {eye_name}: {n*n - bad}/{n*n} points ok, "
              f"projection_raw L{left:+.4f} R{right:+.4f} T{top:+.4f} B{bottom:+.4f}")
        out["eyes"].append({
            "eye": eye_name,
            "projection_raw": [left, right, top, bottom],
            "eye_to_head": eye_to_head,
            "points": points,
        })

    openvr.shutdown()

    fname = f"distortion_dump_{serial.lower().replace('-', '_')}.json"
    with open(fname, "w") as f:
        json.dump(out, f)
    print(f"Wrote {fname}")


if __name__ == "__main__":
    main()