#!/usr/bin/env python3
"""Inventory + materialise HELD-OUT ground-truth views (never fed to any model).

Output layout mirrors an exp dir so the evaluator reads held-out views with the
same code path as input views (inputs/<obj>_<tag>.png RGBA crop whose alpha is
the GT mask; renders/<obj>_<tag>.npz with fx,fy,cx,cy,c2w_cv,bbox,res).

FB150  (reuses build_fb150.build_view -> identical crop/mask/camera code as the
        4 input views):
  held-out = every scene cam (of 24) NOT in views4 whose modal visibility
  vis>0 and modal pixels >= MIN_PX.  The alpha written by build_view is the
  modal (visible) mask & depth>0 -> occluders and occluded object parts are
  excluded.  Tag = h<cam:02d>.
Omni300:
  saved held-out = top, top2, bottom, bottom2 (already in the exp, used in place).
  pool held-out  = every other official OmniObject3D render (100/obj) with
  >=400 valid px (same threshold as build_omni8v_rand) and not one of the 8
  saved views.  Materialised with omni_adapter.load_transforms/read_view/
  crop_bbox + the selection's (center_r, scale_r), exactly as
  build_omni8v_rand.build.  Tag = p<idx:03d>.  (depth_mm is not stored: the
  evaluator never reads it.)
"""
import argparse, json, os, sys
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor
import numpy as np

HERE = Path("/lp-dev/jonghoon/mv-mesh/.debug/eval_heldout_nv")
FB150 = Path("/lp-dev/jonghoon/mv-mesh/.debug/forgebench_eval/fb150")
VIEWS4 = Path("/home/nvidia/jonghoon/mv-mesh/.debug/forgebench150/views_4v.json")
OMNI = Path("/lp-dev/jonghoon/mv-mesh/exp_faithfulness/omni3d300_rand")
MIN_PX = 400
OMNI_IN4 = ["front", "side", "back", "oside"]
OMNI_SAVED_HO = ["top", "top2", "bottom", "bottom2"]


# ------------------------------------------------------------------ FB150
def _fb_work(a):
    s, v4, out = a
    sys.path.insert(0, str(FB150))
    import build_fb150 as BF  # build_view (same code path as input views)
    from PIL import Image
    vis = v4["all_cam_visibility"]; ins = set(v4["cams"])
    recs = []
    for ci, vv in enumerate(vis):
        if ci in ins or vv <= 0:
            continue
        tag = f"h{ci:02d}"
        try:
            BF.build_view(s, ci, tag, Path(out))
        except RuntimeError:          # empty modal mask after depth>0
            continue
        png = Path(out) / "inputs" / f"{s['object']}_{tag}.png"
        px = int((np.asarray(Image.open(png))[..., 3] > 0).sum())
        if px < MIN_PX:
            png.unlink(); (Path(out) / "renders" / f"{s['object']}_{tag}.npz").unlink()
            continue
        recs.append(dict(tag=tag, cam=ci, vis=float(vv), px=px, root=str(out)))
    return s["object"], recs


def prep_fb(workers):
    out = HERE / "fb150_heldout"
    for d in ("inputs", "renders"):
        (out / d).mkdir(parents=True, exist_ok=True)
    sel = json.load(open(FB150 / "exp_4v/selection.json"))["selections"]
    v4 = json.load(open(VIEWS4))["views"]
    jobs = [(s, v4[s["object"]], str(out)) for s in sel]
    idx = {}
    with ProcessPoolExecutor(workers) as ex:
        for obj, recs in ex.map(_fb_work, jobs):
            idx[obj] = recs
    return idx


# ------------------------------------------------------------------ Omni300
def _omni_work(a):
    rec, out, pool = a
    os.chdir("/lp-dev/jonghoon/mv-mesh")          # omni_adapter.RAW is relative
    sys.path.insert(0, "/lp-dev/jonghoon/mv-mesh/scratchpad_1k")
    import imageio.v3 as iio
    from omni_adapter import load_transforms, read_view, crop_bbox
    obj = rec["object"]
    recs = []
    for t in OMNI_SAVED_HO:
        png = OMNI / "inputs" / f"{obj}_{t}.png"; npz = OMNI / "renders" / f"{obj}_{t}.npz"
        if png.exists() and npz.exists():
            recs.append(dict(tag=t, cam=int(rec["view_idx"][t]), px=int(rec["valid_px"][t]),
                             root=str(OMNI), kind="saved"))
    if not pool:
        return obj, recs, None
    B, K, frames, c2w_raw = load_transforms(rec.get("category"), obj)
    c2w = c2w_raw.copy()
    c2w[:, :3, 3] = (c2w_raw[:, :3, 3] - np.array(rec["center_r"])) / rec["scale_r"]
    # frame check: recomputed canonical c2w of the 8 saved views == stored npz
    err = max(float(np.abs(np.load(OMNI / "renders" / f"{obj}_{t}.npz")["c2w_cv"] - c2w[i]).max())
              for t, i in rec["view_idx"].items())
    used = set(int(i) for i in rec["view_idx"].values())
    for i in range(len(frames)):
        if i in used:
            continue
        rgba, dep, valid = read_view(B, i)
        if valid.sum() < MIN_PX:
            continue
        bb = crop_bbox(valid); y0, y1, x0, x1 = bb
        tag = f"p{i:03d}"
        iio.imwrite(Path(out) / "inputs" / f"{obj}_{tag}.png", rgba[y0:y1, x0:x1])
        np.savez_compressed(Path(out) / "renders" / f"{obj}_{tag}.npz", fx=K["fx"], fy=K["fy"],
                            cx=K["cx"], cy=K["cy"], c2w_cv=c2w[i], bbox=np.array(bb), res=800)
        recs.append(dict(tag=tag, cam=i, px=int(valid.sum()), root=str(out), kind="pool"))
    return obj, recs, err


def prep_omni(workers, pool):
    out = HERE / "omni300_heldout"
    for d in ("inputs", "renders"):
        (out / d).mkdir(parents=True, exist_ok=True)
    sel = json.load(open(OMNI / "selection.json"))["selections"]
    idx, errs = {}, {}
    with ProcessPoolExecutor(workers) as ex:
        for obj, recs, err in ex.map(_omni_work, [(r, str(out), pool) for r in sel]):
            idx[obj] = recs; errs[obj] = err
    return idx, errs


def summ(idx, key=None):
    n = [len([r for r in v if key is None or key(r)]) for v in idx.values()]
    return dict(n_obj=len(n), total=int(sum(n)), min=int(min(n)), median=float(np.median(n)),
                max=int(max(n)), mean=round(float(np.mean(n)), 2), n_obj_zero=int(sum(x == 0 for x in n)))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--workers", type=int, default=32)
    ap.add_argument("--no-omni-pool", dest="pool", action="store_false", default=True)
    a = ap.parse_args()
    fb = prep_fb(a.workers)
    json.dump(fb, open(HERE / "heldout_fb150.json", "w"), indent=1)
    om, errs = prep_omni(a.workers, a.pool)
    json.dump(om, open(HERE / "heldout_omni300.json", "w"), indent=1)
    inv = {
        "fb150": {"rule": f"cams not in views4, vis>0, modal px>={MIN_PX}; mask=modal&depth>0 (build_fb150.build_view)",
                  "all": summ(fb), "vis>=0.40": summ(fb, lambda r: r["vis"] >= 0.40),
                  "vis>=0.10": summ(fb, lambda r: r["vis"] >= 0.10)},
        "omni300": {"rule": "saved top/top2/bottom/bottom2 + pool (official 100-view renders, >=400 valid px, not in 8 saved)",
                    "saved": summ(om, lambda r: r["kind"] == "saved"),
                    "pool": summ(om, lambda r: r["kind"] == "pool"), "all": summ(om),
                    "max_c2w_err_saved_vs_recomputed": max((e for e in errs.values() if e is not None), default=None)},
    }
    json.dump(inv, open(HERE / "inventory.json", "w"), indent=1)
    print(json.dumps(inv, indent=1))


if __name__ == "__main__":
    main()