#!/usr/bin/env python
"""Camera-only placement of Pixal3D FB150 outputs in the GT canonical frame (NO fit / ICP / scale fit),
plus input-view silhouette IoU (rendered mapped mesh alpha vs input RGBA alpha, GT camera, stored bbox crop).

  MV (batch_pixal3d_mv_cropK.py):   G (glb) -> W = A @ G ;  W = F(d) @ inv(Ca) @ X_gt
      => X_gt = Ca @ inv(F(d)) @ A @ G,   Ca = anchor_c2w_cv @ diag(1,-1,-1,1)  (<obj>.cams.json;
         = c2w_cv(view0) for --gauge upstream, its look-at rotation for --gauge lookat), d = |Ca t|.
  SV (batch_pixal3d.py, MoGe fov):  cam_map_sv(mode='sim') from .debug/pixal3d_frame/cam_align.py (copied
      verbatim below): similarity whose scale/offset are fixed by the cameras only (GT K, GT c2w, MoGe
      fov, preprocess crop box, depth z0 of the GT origin in the input camera).
Usage: cammap_fb150.py {mv|sv} PRED_DIR OUT_DIR [--objs ...] [--svcams json]
"""
import argparse, json, math, os, sys
import numpy as np, torch, trimesh
from PIL import Image

sys.path.insert(0, "/lp-dev/jonghoon/mv-mesh/.debug/qual_gen/scripts")
import render_panels as RP  # noqa: E402  (load, raster; imports metrics/appeval/render as R)
R = RP.R

FB = "/lp-dev/jonghoon/mv-mesh/.debug/forgebench_eval/fb150"
A = np.array([[-1, 0, 0], [0, 0, 1], [0, 1, 0]], float)
FLIP = np.diag([1., -1., -1., 1.])


def F_mat(d):
    return np.array([[1, 0, 0, 0], [0, 0, -1, -d], [0, 1, 0, 0], [0, 0, 0, 1]], float)


def cam_map_mv(G, anchor_c2w_cv):
    Ca = np.asarray(anchor_c2w_cv, float) @ FLIP
    d = float(np.linalg.norm(Ca[:3, 3]))
    M = Ca @ np.linalg.inv(F_mat(d))
    return (G @ A.T) @ M[:3, :3].T + M[:3, 3]


def cam_map_sv(G, z, cam):   # == cam_align.cam_map_sv(..., mode='sim')
    c2w = np.asarray(z["c2w_cv"], float)
    fx, fy = float(z["fx"]), float(z["fy"])
    y0, y1, x0, x1 = [int(v) for v in z["bbox"]]
    cxp, cyp = float(z["cx"]) - x0, float(z["cy"]) - y0
    fov, d = cam["camera_angle_x"], cam["distance"]
    ox, oy = cam["box"][0], cam["box"][1]; S = cam["side"]
    W = G @ A.T
    p = np.c_[W[:, 0], -W[:, 2], W[:, 1] + d]
    z0 = (np.linalg.inv(c2w) @ np.r_[0, 0, 0, 1.0])[2]
    ray = np.array([(ox + S / 2 - cxp) / fx, (oy + S / 2 - cyp) / fy, 1.0])
    rn = ray / np.linalg.norm(ray)
    zc = np.array([0, 0, 1.0]); v = np.cross(zc, rn); c = zc @ rn
    vx = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
    Rr = np.eye(3) + vx + vx @ vx / (1 + c)
    k = z0 * (S / fx) / (d * 2 * math.tan(fov / 2))
    P = k * ((p - np.array([0, 0, d])) @ Rr.T) + ray * z0
    return P @ c2w[:3, :3].T + c2w[:3, 3], dict(k=float(k), z0=float(z0))


def iv_cam(z, res=512):
    y0, y1, x0, x1 = [int(v) for v in z["bbox"]]
    ch, cw = y1 - y0, x1 - x0
    s = res / max(ch, cw)
    H, W = int(round(ch * s)), int(round(cw * s))
    P = R.gl_proj_from_K(float(z["fx"]) * s, float(z["fy"]) * s, (float(z["cx"]) - x0) * s,
                         (float(z["cy"]) - y0) * s, W, H, 0.05, 50.0)
    Vgl = R.cv_extrinsic_to_gl(np.asarray(z["c2w_cv"], np.float64))
    return P @ Vgl, Vgl, H, W


def sil_iou(glb, exp, obj, tags, ctx):
    g = RP.load(glb)
    out = {}
    for t in tags:
        z = np.load(f"{exp}/renders/{obj}_{t}.npz")
        mvp, view, H, W = iv_cam(z)
        rgba, _ = RP.raster(g, mvp, view, H, W, ctx)
        pr = rgba[..., 3] > 0.5
        a = np.array(Image.open(f"{exp}/inputs/{obj}_{t}.png").getchannel(3).resize((W, H), Image.BILINEAR)) > 127
        out[t] = float((pr & a).sum() / max((pr | a).sum(), 1))
    del g; torch.cuda.empty_cache()
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("mode", choices=["mv", "sv"])
    ap.add_argument("pred"); ap.add_argument("out")
    ap.add_argument("--objs", nargs="+")
    ap.add_argument("--svcams", default="/lp-dev/jonghoon/mv-mesh/.debug/pixal3d_fb150_fix/fb150_sv_cams_all.json")
    ap.add_argument("--no-iou", action="store_true")
    a = ap.parse_args()
    os.makedirs(a.out, exist_ok=True)
    exp = f"{FB}/exp_4v" if a.mode == "mv" else f"{FB}/exp_1v"
    tags = ["front", "side", "back", "oside"] if a.mode == "mv" else ["front"]
    objs = a.objs or [s["object"] for s in json.load(open(f"{exp}/selection.json"))["selections"]]
    svc = json.load(open(a.svcams)) if a.mode == "sv" else None
    ctx = R.get_ctx()
    res = {}
    for i, o in enumerate(objs, 1):
        src = f"{a.pred}/{o}.glb"
        if not os.path.isfile(src):
            res[o] = dict(status="missing"); print(f"[{i}] {o} MISSING", flush=True); continue
        m = trimesh.load(src, force="mesh", process=False)
        G = np.asarray(m.vertices, float)
        if a.mode == "mv":
            cj = json.load(open(f"{a.pred}/{o}.cams.json")) if os.path.isfile(f"{a.pred}/{o}.cams.json") else None
            anchor = cj["anchor_c2w_cv"] if cj else np.load(f"{exp}/renders/{o}_front.npz")["c2w_cv"]
            X = cam_map_mv(G, anchor); extra = dict(gauge=cj["gauge"] if cj else "upstream(no cams.json)")
        else:
            X, extra = cam_map_sv(G, np.load(f"{exp}/renders/{o}_front.npz"), svc[o])
        m.vertices = X
        dst = f"{a.out}/{o}.glb"
        m.export(dst + ".tmp.glb"); os.replace(dst + ".tmp.glb", dst)
        r = dict(status="ok", nverts=int(len(X)), extent=float((X.max(0) - X.min(0)).max()), **extra)
        if not a.no_iou:
            r["iou"] = sil_iou(dst, exp, o, tags, ctx); r["iou_mean"] = float(np.mean(list(r["iou"].values())))
        res[o] = r
        print(f"[{i}/{len(objs)}] {o} {json.dumps(r)}", flush=True)
    jp = f"{a.out}/_cammap.json"
    old = json.load(open(jp)) if os.path.isfile(jp) else {}
    old.update(res); json.dump(old, open(jp, "w"), indent=1)


if __name__ == "__main__":
    main()