#!/usr/bin/env python3
"""FORGE3DBench-150: select + mark 150 objects, pick 4v views, build 1v/4v exps.

Selection (seed 0):
  * candidates = the 674 objects with an existing ReconViaGen 1v mesh
    (.debug/forgebench_eval/gen/rvg/*.glb) + for scenes with NO such object, the
    remaining forgebench_1v (808-set) objects of that scene.
  * each asset used at most once; every one of the 120 scenes gets >=1 object
    (max bipartite matching scene->asset, RVG-having candidates preferred),
    then 30 extra objects (distinct unused assets, distinct scenes, RVG-having
    only, stratified round-robin over env/anchor).
  * each object keeps its 1v view (cam) from the 808-set manifest.
4v views: [1v cam] + 3 random other cams of the same scene with vis>=0.40
  (random.Random(f"{uid}/4v/{SEED}")); fallback = highest-vis remaining cams (flagged).
  Tags: companions assigned to Hunyuan3D-2mv slots by relative azimuth
  (min-cost assignment over the 6 permutations): side=right, back=back, oside=left,
  using the image-space criterion of .debug_hy2mv/LOG.md A4 (F.r>0 -> right).
Exps: fb150/exp_1v (files copied from exp_faithfulness/forgebench_1v),
      fb150/exp_4v (front copied from 1v; side/back/oside built with
      build_forgebench_1v.py conventions).
"""
import os, sys, json, random, itertools, shutil, numpy as np, trimesh
from pathlib import Path
from PIL import Image
from concurrent.futures import ProcessPoolExecutor

sys.path.insert(0, "/lp-dev/jonghoon/mv-mesh/.debug/forgebench_eval")
import build_forgebench_1v as B  # crop_bbox, pointmap_from_depth, R_FIX, ROOT

SEED = 0
VIS_THR = 0.40
FB = Path("/lp-dev/jonghoon/mv-mesh/.debug/forgebench_eval")
FB150 = FB / "fb150"
EXP1 = Path("/lp-dev/jonghoon/mv-mesh/exp_faithfulness/forgebench_1v")
HOME = Path("/home/nvidia/jonghoon/mv-mesh/.debug/forgebench150")
TAGS4 = ["front", "side", "back", "oside"]
SLOT = {"side": "right", "back": "back", "oside": "left"}


def scene_key(s):
    return (s["env"], s["anchor"], s["scene"])


def select():
    man = json.load(open(FB / "manifest.json"))["selections"]
    rvg = {f[:-4] for f in os.listdir(FB / "gen/rvg") if f.endswith(".glb") and ".tmp" not in f}
    for s in man:
        s["has_rvg"] = s["object"] in rvg
    scenes = sorted({scene_key(s) for s in man})
    rvg_scenes = {scene_key(s) for s in man if s["has_rvg"]}
    rng = random.Random(SEED)
    # candidate lists per scene
    cand = {}
    for sk in scenes:
        objs = [s for s in man if scene_key(s) == sk]
        pool = [s for s in objs if s["has_rvg"]] if sk in rvg_scenes else objs
        pool = sorted(pool, key=lambda s: s["object"])
        rng.shuffle(pool)
        cand[sk] = pool
    # max bipartite matching scene -> asset (Kuhn), scenes in random order
    order = scenes[:]
    rng.shuffle(order)
    match_asset = {}  # asset -> scene

    def try_scene(sk, seen):
        for s in cand[sk]:
            a = s["asset"]
            if a in seen:
                continue
            seen.add(a)
            if a not in match_asset or try_scene(match_asset[a], seen):
                match_asset[a] = sk
                return True
        return False

    for sk in order:
        if not try_scene(sk, set()):
            raise SystemExit(f"cannot cover scene {sk} with distinct assets")
    scene_asset = {sk: a for a, sk in match_asset.items()}
    chosen = []
    for sk in scenes:
        a = scene_asset[sk]
        s = [c for c in cand[sk] if c["asset"] == a][0]
        chosen.append(dict(s, role="coverage"))
    used = {c["asset"] for c in chosen}
    # 30 extras: RVG-having, unused assets, distinct scenes, round-robin over anchors
    groups = {}
    for sk in sorted(rvg_scenes):
        groups.setdefault((sk[0], sk[1]), []).append(sk)
    for g in groups.values():
        rng.shuffle(g)
    gkeys = sorted(groups)
    extras, used_sc = [], set()
    while len(extras) < 30:
        progressed = False
        for gk in gkeys:
            if len(extras) >= 30:
                break
            while groups[gk]:
                sk = groups[gk].pop()
                if sk in used_sc:
                    continue
                pool = [s for s in cand[sk] if s["asset"] not in used and s["has_rvg"]]
                if not pool:
                    continue
                s = pool[0]
                extras.append(dict(s, role="extra")); used.add(s["asset"]); used_sc.add(sk)
                progressed = True
                break
        if not progressed:
            raise SystemExit("ran out of extras")
    sel = sorted(chosen + extras, key=lambda s: s["object"])
    assert len(sel) == 150 and len({s["asset"] for s in sel}) == 150
    assert len({scene_key(s) for s in sel}) == 120
    return sel, rvg


def vis_all(s):
    """per-cam visibility (modal/amodal), exactly as build_forgebench_1v."""
    p = s["src_scene"]; oid = s["obj_id"]; mc = np.array(s["mask_color"], np.int16)
    cams = json.load(open(f"{p}/cameras.json"))["cameras"]
    ncam = min(len(cams), B.NCAM)
    vis = []
    for ci in range(ncam):
        mp = f"{p}/mask/cam{ci:02d}.png"; amp = f"{p}/mask_amodal/cam{ci:02d}/{oid}.png"
        if not (os.path.exists(mp) and os.path.exists(amp)):
            vis.append(0.0); continue
        m = np.array(Image.open(mp).convert("RGB")).astype(np.int16)
        ma = int((np.abs(m - mc).sum(2) <= 8).sum())
        aa = int((np.array(Image.open(amp).convert("L")) > 0).sum())
        vis.append(round(ma / aa, 4) if aa > 0 else 0.0)
    return vis


def horiz(v):
    v = np.array(v, float).copy(); v[2] = 0
    return v / (np.linalg.norm(v) + 1e-12)


def assign_slots(s, front, comps):
    cams = json.load(open(f"{s['src_scene']}/cameras.json"))["cameras"]
    ctr = np.array(s["center"])
    def pos(ci): return np.array(cams[ci]["extrinsic_c2w"])[:3, 3]
    def right(ci): return np.array(cams[ci]["extrinsic_c2w"])[:3, 0]
    F = horiz(pos(front) - ctr)
    feats = {}
    for c in comps:
        d = horiz(pos(c) - ctr); r = horiz(right(c))
        a = float(F @ d); sr = float(F @ r)
        az = float(np.degrees(np.arctan2(np.cross(F, d)[2], F @ d)))
        feats[c] = dict(cos_to_front=round(a, 4), F_dot_right=round(sr, 4), rel_azimuth_deg=round(az, 2))
    ideal = {"side": (0.0, 1.0), "back": (-1.0, 0.0), "oside": (0.0, -1.0)}
    best = None
    for perm in itertools.permutations(comps):
        cost = sum((feats[c]["cos_to_front"] - ideal[t][0]) ** 2 + (feats[c]["F_dot_right"] - ideal[t][1]) ** 2
                   for t, c in zip(["side", "back", "oside"], perm))
        if best is None or cost < best[0]:
            best = (cost, perm)
    return dict(zip(["side", "back", "oside"], best[1])), feats, best[0]


def pick_4v(s, vis):
    front = s["cam"]
    qual = [c for c in range(len(vis)) if c != front and vis[c] >= VIS_THR]
    r = random.Random(f"{s['object']}/4v/{SEED}")
    flag = None
    if len(qual) >= 3:
        comps = r.sample(sorted(qual), 3)
    else:
        rest = sorted([c for c in range(len(vis)) if c != front and c not in qual], key=lambda c: -vis[c])
        comps = sorted(qual) + rest[:3 - len(qual)]
        flag = f"only {len(qual)} other cams with vis>=0.40; filled by highest-vis"
    tagmap, feats, cost = assign_slots(s, front, comps)
    return comps, tagmap, feats, cost, flag, len(qual)


def build_view(s, ci, tag, out):
    """side/back/oside view files; same code path as build_forgebench_1v.build_object."""
    p = s["src_scene"]; uid = s["object"]; mc = np.array(s["mask_color"], np.int16)
    scale_r = s["scale_r"]; center = np.array(s["center"])
    cam = json.load(open(f"{p}/cameras.json"))["cameras"][ci]; Kd = cam["intrinsics"]
    c2w = np.array(cam["extrinsic_c2w"]).copy(); c2w[:3, 3] = (c2w[:3, 3] - center) / scale_r
    modal = np.array(Image.open(f"{p}/mask/cam{ci:02d}.png").convert("RGB")).astype(np.int16)
    sel = (np.abs(modal - mc).sum(2) <= 8)
    raw = np.array(Image.open(f"{p}/depth/cam{ci:02d}.png")).astype(np.float64)
    sel = sel & (raw > 0)
    if sel.sum() == 0:
        raise RuntimeError(f"empty mask {uid} cam{ci}")
    dmm = np.zeros(raw.shape, np.float64); dmm[sel] = raw[sel] / scale_r
    dmm = np.clip(np.rint(dmm), 0, 65535).astype(np.uint16)
    bb = B.crop_bbox(sel); y0, y1, x0, x1 = bb
    rgb = np.array(Image.open(f"{p}/rgb/cam{ci:02d}.png").convert("RGB"))
    alpha = np.where(sel, 255, 0).astype(np.uint8)
    Image.fromarray(np.dstack([rgb, alpha])[y0:y1, x0:x1], "RGBA").save(out / "inputs" / f"{uid}_{tag}.png")
    np.savez_compressed(out / "renders" / f"{uid}_{tag}.npz", depth_mm=dmm, fx=float(Kd["fx"]), fy=float(Kd["fy"]),
                        cx=float(Kd["cx"]), cy=float(Kd["cy"]), c2w_cv=c2w, bbox=np.array(bb, np.int64),
                        res=int(raw.shape[0]), up_axis="y")


def pm_of(out, uid, tag):
    d = np.load(out / "renders" / f"{uid}_{tag}.npz")
    Kv = {k: float(d[k]) for k in ("fx", "fy", "cx", "cy")}
    return B.pointmap_from_depth(d["depth_mm"], Kv, tuple(d["bbox"]))


def work(s):
    vis = vis_all(s)
    assert abs(vis[s["cam"]] - s["visibility"]) < 1e-3, (s["object"], vis[s["cam"]], s["visibility"])
    comps, tagmap, feats, cost, flag, nq = pick_4v(s, vis)
    uid = s["object"]
    e1, e4 = FB150 / "exp_1v", FB150 / "exp_4v"
    for e in (e1, e4):
        for sub in ("inputs", "renders"):
            (e / sub).mkdir(parents=True, exist_ok=True)
    # 1v exp: copy existing files (identical inputs to the 808-set run)
    for src, dst in [(EXP1 / "inputs" / f"{uid}_front.png", "inputs"),
                     (EXP1 / "renders" / f"{uid}_front.npz", "renders"),
                     (EXP1 / "renders" / f"{uid}_canon.glb", "renders")]:
        for e in (e1, e4):
            shutil.copy2(src, e / dst / src.name)
    (e1 / "npz_1v" / uid).mkdir(parents=True, exist_ok=True)
    shutil.copy2(EXP1 / "npz_1v" / uid / "da3_output.npz", e1 / "npz_1v" / uid / "da3_output.npz")
    for tag in ("side", "back", "oside"):
        build_view(s, tagmap[tag], tag, e4)
    pms = np.stack([pm_of(e4, uid, t) for t in TAGS4])
    for v in (1, 2, 4):
        dd = e4 / f"npz_{v}v" / uid; dd.mkdir(parents=True, exist_ok=True)
        np.savez(dd / "da3_output.npz", pointmaps_sam3d=pms[:v], image_files=np.array([f"v{k+1}.png" for k in range(v)]))
    cams4 = [s["cam"]] + [tagmap[t] for t in ("side", "back", "oside")]
    return dict(object=uid, all_cam_visibility=vis, n_qual_other=nq,
                random_draw_order=comps, cams=cams4, tags=TAGS4,
                visibility=[vis[c] for c in cams4],
                hy3d2mv_slots={"front": s["cam"], **{SLOT[t]: tagmap[t] for t in ("side", "back", "oside")}},
                tag_to_hy3d2mv_slot={"front": "front", **SLOT},
                companion_geometry={str(c): feats[c] for c in comps}, slot_assign_cost=round(cost, 4),
                flag=flag, seed=SEED)


def main():
    sel, rvg = select()
    with ProcessPoolExecutor(32) as ex:
        v4 = list(ex.map(work, sel))
    v4d = {r["object"]: r for r in v4}
    keep = ["object", "env", "anchor", "scene", "asset", "obj_id", "cam", "visibility", "n_qual",
            "role", "has_rvg", "mask_color", "scale_r", "center", "extent", "src_scene", "front", "modal_area"]
    recs = [{k: s[k] for k in keep} for s in sel]
    for r in recs:
        r["episode"] = f"{r['env']}/{r['anchor']}/{r['scene']}"
    # scene coverage
    scenes = sorted({(r["env"], r["anchor"], r["scene"]) for r in recs})
    per_anchor = {}
    for r in recs:
        per_anchor.setdefault(f"{r['env']}/{r['anchor']}", set()).add(r["scene"])
    cov = {"n_objects": len(recs), "n_scenes_covered": len(scenes), "n_scenes_total": 120,
           "n_unique_assets": len({r["asset"] for r in recs}),
           "per_anchor_scenes": {k: len(v) for k, v in sorted(per_anchor.items())},
           "per_anchor_objects": {k: sum(1 for r in recs if f"{r['env']}/{r['anchor']}" == k) for k in sorted(per_anchor)},
           "n_rvg_reused": sum(r["has_rvg"] for r in recs),
           "rvg_to_generate": sorted(r["object"] for r in recs if not r["has_rvg"])}
    HOME.mkdir(parents=True, exist_ok=True)
    json.dump({"seed": SEED, "vis_thr": VIS_THR, "source_manifest": str(FB / "manifest.json"),
               "coverage": cov, "selections": recs}, open(HOME / "selection_150.json", "w"), indent=1)
    lines = ["# FORGE3DBench-150 scene coverage", "",
             f"objects={cov['n_objects']}  scenes covered={cov['n_scenes_covered']}/120  unique assets={cov['n_unique_assets']}",
             f"RVG 1v reused={cov['n_rvg_reused']}  RVG 1v to generate={len(cov['rvg_to_generate'])}", "",
             "| env/anchor | scenes covered | objects |", "|---|---:|---:|"]
    for k in cov["per_anchor_scenes"]:
        lines.append(f"| {k} | {cov['per_anchor_scenes'][k]}/20 | {cov['per_anchor_objects'][k]} |")
    lines += ["", "RVG 1v generated fresh (scenes without any existing RVG mesh):"] + [f"- {o}" for o in cov["rvg_to_generate"]]
    lines += ["", "| scene | objects |", "|---|---|"]
    for sk in scenes:
        lines.append(f"| {'/'.join(sk)} | " + ", ".join(r["object"] for r in recs if (r['env'], r['anchor'], r['scene']) == sk) + " |")
    (HOME / "scene_coverage.md").write_text("\n".join(lines) + "\n")
    json.dump({"seed": SEED, "vis_thr": VIS_THR,
               "rule": "cams[0]=1v anchor cam; 3 companions = random.Random(f'{object}/4v/0').sample(other cams with vis>=0.40, 3); "
                       "fallback highest-vis if <3 qualify (flag). Companions assigned to tags side/back/oside by min-cost "
                       "match to Hunyuan3D-2mv slots right/back/left (features: F.d cos to front, F.r image-right criterion, LOG.md A4).",
               "tag_to_hy3d2mv_slot": {"front": "front", **SLOT},
               "n_flagged": sum(1 for r in v4 if r["flag"]),
               "views": v4d}, open(HOME / "views_4v.json", "w"), indent=1)
    for e in ("exp_1v", "exp_4v"):
        json.dump({"source": f"forgebench150_{e[-2:]}", "seed": SEED, "vis_thr": VIS_THR, "n": len(recs),
                   "selections": [dict(r, **({"views4": v4d[r['object']]['cams']} if e == 'exp_4v' else {})) for r in recs]},
                  open(FB150 / e / "selection.json", "w"), indent=1)
    print(json.dumps(cov, indent=1))
    print("flagged 4v:", [r["object"] for r in v4 if r["flag"]])


if __name__ == "__main__":
    main()