#!/usr/bin/env python3
"""Materialise HELD-OUT ground-truth views for FORGE3DBench rest-658 (the 808-set objects NOT in FB150).

Copy of the FB150 branch of eval_heldout_nv/prepare_heldout.py (_fb_work / prep_fb), unchanged logic;
only the source selection / views json and the output root differ:
  selection = forgebench_eval/rest658/exp_4v/selection.json   (built by build_rest658.py = build_fb150.work)
  views4    = forgebench_eval/rest658/views_4v.json           (all_cam_visibility = modal/amodal area, vis_all)
ALL-VIEWS build (coordinator request 2026-09-26): EVERY scene cam (of 24) NOT in views4 gets an index record
with its visibility vis (modal/amodal area, views_4v.json all_cam_visibility) and modal px. Files are materialised
with build_fb150.build_view (identical crop/mask/camera code as the 4 input views; alpha = modal mask & depth>0)
whenever the modal mask is non-empty (px<400 views INCLUDED, flagged px_ok=False). Cams with an empty modal mask
(object fully hidden / out of frame) cannot be cropped: record only, files=False, px=0. Tag = h<cam:02d>.
The FB150 protocol set (vis>=0.20 & px>=400 & files) is applied at eval time by eval_rest658_wrapper.py
(--min-vis 0.2, --min-px 400); with those filters the NV set equals the FB150 rule exactly.
Outputs: <OUT>/rest658_heldout/{inputs,renders}, <OUT>/heldout_rest658.json (schema of heldout_fb150.json),
         <OUT>/inventory_rest658.json
"""
import argparse, json, sys
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor
import numpy as np

HERE = Path("/lp-dev/jonghoon/mv-mesh/.debug/hub_rest658_heldout")
FB150 = Path("/lp-dev/jonghoon/mv-mesh/.debug/forgebench_eval/fb150")        # build_fb150.py lives here
REST = Path("/lp-dev/jonghoon/mv-mesh/.debug/forgebench_eval/rest658")
VIEWS4 = REST / "views_4v.json"
MIN_PX = 400


def _fb_work(a):   # == prepare_heldout._fb_work
    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:
            continue
        tag = f"h{ci:02d}"
        try:
            BF.build_view(s, ci, tag, Path(out))
        except RuntimeError:          # empty modal mask after depth>0 -> record only
            recs.append(dict(tag=tag, cam=ci, vis=float(vv), px=0, px_ok=False, files=False, root=str(out)))
            continue
        png = Path(out) / "inputs" / f"{s['object']}_{tag}.png"
        px = int((np.asarray(Image.open(png))[..., 3] > 0).sum())
        recs.append(dict(tag=tag, cam=ci, vis=float(vv), px=px, px_ok=px >= MIN_PX, files=True, root=str(out)))
    return s["object"], recs


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=64)
    ap.add_argument("--limit", type=int, default=0)
    a = ap.parse_args()
    out = HERE / "rest658_heldout"
    for d in ("inputs", "renders"):
        (out / d).mkdir(parents=True, exist_ok=True)
    sel = json.load(open(REST / "exp_4v/selection.json"))["selections"]
    v4 = json.load(open(VIEWS4))["views"]
    assert len(sel) == 658 and set(v4) == {s["object"] for s in sel}
    for s in sel:   # the 4 input cams in the exp == the cams excluded here
        assert list(s["views4"]) == list(v4[s["object"]]["cams"]), s["object"]
    if a.limit:
        sel = sel[:a.limit]
    jobs = [(s, v4[s["object"]], str(out)) for s in sel]
    idx = {}
    with ProcessPoolExecutor(a.workers) as ex:
        for obj, recs in ex.map(_fb_work, jobs):
            idx[obj] = recs
    json.dump(idx, open(HERE / "heldout_rest658.json", "w"), indent=1)
    inv = {"rest658": {"rule": f"ALL cams not in views4 (files if modal mask non-empty); protocol filter vis>=0.20 & px>={MIN_PX} at eval",
                       "all_records": summ(idx), "with_files": summ(idx, lambda r: r["files"]),
                       "fb150_rule_vis>0_px>=400": summ(idx, lambda r: r["files"] and r["px_ok"] and r["vis"] > 0),
                       "PROTOCOL_vis>=0.20_px>=400": summ(idx, lambda r: r["files"] and r["px_ok"] and r["vis"] >= 0.20),
                       "vis>=0.20_any_px": summ(idx, lambda r: r["files"] and r["vis"] >= 0.20),
                       "empty_mask_no_files": summ(idx, lambda r: not r["files"]),
                       "vis>=0.40_px>=400": summ(idx, lambda r: r["files"] and r["px_ok"] and r["vis"] >= 0.40)}}
    json.dump(inv, open(HERE / "inventory_rest658.json", "w"), indent=1)
    print(json.dumps(inv, indent=1))


if __name__ == "__main__":
    main()