#!/usr/bin/env python3
"""Held-out-GT-view novel-view (NV) appearance evaluator.

NV protocol: render the PREDICTED mesh at held-out dataset cameras (views never
fed to the model; see prepare_heldout.py) and compare with the dataset's own GT
image of that view, restricted to the GT mask (FB150: modal/visible mask, so
occluders and occluded object parts are never scored; Omni: GT alpha).

Everything except the held-out camera list and the mask restriction is
IMPORTED from metrics/appeval (unchanged):
  * IV + geometry (+ optional legacy 24-orbit NV): evaluate_appforce.eval_object
  * renderer: render.prepare_mesh / render.render_input_view (nvdiffrast)
  * metrics : appearance.appearance_metrics / average_views
  * crop    : held-out npz bbox (y0,y1,x0,x1) + RGBA crop png, identical format
              and code path as the input views (eval_object IV block).
Mask restriction (the only new metric logic): pred alpha is multiplied by the GT
mask M before appearance_metrics composites both over white, so pixels outside
M are white in both images (ignored) and M-pixels the pred does not cover are
scored as white-vs-GT (penalised).
"""
from __future__ import annotations
import argparse, json, sys
from pathlib import Path
import numpy as np

sys.path.insert(0, "/home/nvidia/jonghoon/mv-mesh/metrics/appeval")
import render as R                                             # noqa: E402
import evaluate_appforce as EA                                  # noqa: E402
from appearance import appearance_metrics, average_views, composite_white  # noqa: E402
from geometry import geometry_metrics                            # noqa: E402

HERE = Path("/lp-dev/jonghoon/mv-mesh/.debug/eval_heldout_nv")
DS = {
    "fb150": dict(exp=Path("/lp-dev/jonghoon/mv-mesh/.debug/forgebench_eval/fb150/exp_4v"),
                  index=HERE / "heldout_fb150.json"),
    "omni300": dict(exp=Path("/lp-dev/jonghoon/mv-mesh/exp_faithfulness/omni3d300_rand"),
                    index=HERE / "heldout_omni300.json"),
}
MKEYS = ("lpips", "ssim", "clip", "psnr")


def heldout_view(g, rec, obj, ctx, gt_g=None):
    """One held-out view: returns metrics dict (+ coverage diagnostics)."""
    root = Path(rec["root"])
    z = np.load(root / "renders" / f"{obj}_{rec['tag']}.npz")
    K = {k: float(z[k]) for k in ("fx", "fy", "cx", "cy")}
    res = int(z["res"]); y0, y1, x0, x1 = z["bbox"].tolist()
    ref = EA.read_input_png(root / "inputs" / f"{obj}_{rec['tag']}.png")
    pred = R.render_input_view(g, K, z["c2w_cv"], res, res, ctx=ctx).cpu().numpy()[y0:y1, x0:x1]
    h, w = min(pred.shape[0], ref.shape[0]), min(pred.shape[1], ref.shape[1])
    pred, ref = pred[:h, :w].copy(), ref[:h, :w]
    M = (ref[..., 3] > 0.5).astype(np.float32)
    pred[..., 3] *= M                                  # restrict to GT mask
    ref = ref.copy(); ref[..., 3] = M
    m = appearance_metrics(pred, ref)
    m["coverage"] = float((pred[..., 3] > 0.5).sum() / max(M.sum(), 1))  # frac of mask hit by pred
    m.update(view=rec["tag"], cam=rec["cam"], mask_px=int(M.sum()))
    if "vis" in rec:
        m["vis"] = rec["vis"]
    return m, pred, ref


def run_object(a, obj, views, ctx, save_debug):
    gt_path = EA.resolve_gt(a.exp / "renders", obj)
    pred_path = gt_path if a.pred_is_gt else a.meshes / f"{obj}.glb"
    if not pred_path.exists():
        return {"object": obj, "error": "missing pred glb"}
    # ---- IV + geometry (+legacy orbit NV) via the UNCHANGED existing code path
    if a.pred_is_gt:
        row = {"object": obj, "pred": "GT"}
    else:
        row = EA.eval_object(a.exp, a.meshes, a.exp / "renders", obj, EA.VIEW_TAGS[a.views],
                             a.legacy_novel, a.out, save_debug=False, ctx=ctx)
        if "error" in row:
            return row
        if not a.legacy_novel and not a.no_geometry:
            row["geometry"] = geometry_metrics(EA.load_mesh(pred_path), EA.load_mesh(gt_path))
    # ---- held-out NV
    g = R.prepare_mesh(EA.load_mesh(pred_path))
    per, sheet = [], []
    for rec in views:
        m, p, r = heldout_view(g, rec, obj, ctx)
        per.append(m)
        if save_debug and len(sheet) < 8:
            sheet += [composite_white(r), composite_white(p)]
        if a.pred_is_gt:   # pipeline identity: GT render vs itself in the mask
            ident = appearance_metrics(p, p)
            m["identity_lpips"], m["identity_psnr"] = ident["lpips"], ident["psnr"]
    if per:
        row["heldout_nv"] = {k: float(np.mean([d[k] for d in per])) for k in MKEYS + ("coverage",)}
        row["heldout_nv"]["n_views"] = len(per)
        row["heldout_nv"]["per_view"] = per
    if sheet:
        EA.save_sheet(sheet, a.out / "debug" / f"{obj}_heldout.png")
    return row


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dataset", choices=list(DS), required=True)
    ap.add_argument("--meshes", type=Path, help="DIR with {object}.glb (canonical GT frame)")
    ap.add_argument("--pred-is-gt", action="store_true", help="sanity: score the GT canon mesh")
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--views", choices=list(EA.VIEW_TAGS), default="quad", help="IV tags")
    ap.add_argument("--omni-set", choices=["saved", "pool", "all"], default="saved")
    ap.add_argument("--min-vis", type=float, default=0.0, help="FB150: min modal/amodal visibility")
    ap.add_argument("--legacy-novel", action="store_true", help="also run old 24-orbit NV (+geometry)")
    ap.add_argument("--no-geometry", action="store_true")
    ap.add_argument("--objects", nargs="*"); ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--shard", type=int, default=0); ap.add_argument("--nshards", type=int, default=1)
    a = ap.parse_args()
    a.exp = DS[a.dataset]["exp"]
    idx = json.load(open(DS[a.dataset]["index"]))
    objs = sorted(idx)
    if a.objects:
        objs = [o for o in objs if o in set(a.objects)]
    if a.limit:
        objs = objs[:a.limit]
    objs = objs[a.shard::a.nshards]
    a.out.mkdir(parents=True, exist_ok=True)
    ctx = R.get_ctx()
    res = []
    for i, obj in enumerate(objs):
        v = idx[obj]
        if a.dataset == "fb150":
            v = [r for r in v if r["vis"] >= a.min_vis]
        elif a.omni_set != "all":
            v = [r for r in v if r["kind"] == a.omni_set]
        try:
            r = run_object(a, obj, v, ctx, save_debug=i < 6)
        except Exception as e:
            import traceback; traceback.print_exc()
            r = {"object": obj, "error": f"{type(e).__name__}: {e}"}
        res.append(r)
        hn = r.get("heldout_nv", {})
        print(f"[{i+1}/{len(objs)}] {obj} n={hn.get('n_views')} lpips={hn.get('lpips', float('nan')):.4f} "
              f"psnr={hn.get('psnr', float('nan')):.2f} cov={hn.get('coverage', float('nan')):.3f}", flush=True)
    (a.out / "results.json").write_text(json.dumps(res, indent=1))
    ok = [r for r in res if "heldout_nv" in r]
    summ = {"n_objects": len(ok), "n_views_total": int(sum(r["heldout_nv"]["n_views"] for r in ok)),
            "heldout_nv": {k: float(np.mean([r["heldout_nv"][k] for r in ok])) for k in MKEYS + ("coverage",)}}
    iv = [r for r in ok if r.get("input_view")]
    if iv:
        summ["input_view"] = {k: float(np.mean([r["input_view"][k] for r in iv])) for k in MKEYS}
    ge = [r for r in ok if r.get("geometry")]
    if ge:
        summ["geometry"] = {k: float(np.mean([r["geometry"][k] for r in ge]))
                            for k in ("cd_l1", "f01", "f02", "f05", "normal_consistency") if k in ge[0]["geometry"]}
    nv = [r for r in ok if r.get("novel_view")]
    if nv:
        summ["legacy_orbit_nv"] = {k: float(np.mean([r["novel_view"][k] for r in nv])) for k in MKEYS}
    (a.out / "summary.json").write_text(json.dumps(summ, indent=1))
    print(json.dumps(summ, indent=1))


if __name__ == "__main__":
    main()