#!/usr/bin/env python
"""Benchmark driver: Amodal3R (multi-view image, pose-free) on the mv-mesh dataset.

Amodal3R reconstructs a full 3D object from N OCCLUDED views + a 3-valued
visibility mask per view. We adapt it to our dataset layout WITHOUT touching the
official pipeline (amodal3r/repo/amodal3r/pipelines/image_to_3d.py): we IMPORT
`Amodal3RImageTo3DPipeline.run_multi_image` and feed it lists we assemble here.

Mask convention (official preprocess_image_w_mask):
    L grayscale:  <127 = OCCLUDED,  >230 = BACKGROUND,  ~188 = VISIBLE.
Our inputs are RGBA crops (alpha = VISIBLE/modal object mask on black).
    visible (alpha > 0.5 after the 518 LANCZOS resize)       -> 188
    OCCLUDED = amodal silhouette & ~visible, inside the crop -> 0
    everything else                                          -> 255
The occluded region needs the benchmark's real occlusion data (FORGE3DBench:
<src_scene>/mask_amodal/cam<ci:02d>/<obj_id>.png + renders/<obj>_<tag>.npz bbox,
with src_scene/obj_id/views4 from selection.json); see --mask-mode:
    auto (default) : occ2 when the selection entry carries src_scene+obj_id
                     (FORGE3DBench), else prod (clean renders, nothing occluded)
    occ2           : the 3-value mask above (occluded computed at native crop
                     resolution, resampled with the same LANCZOS/>0.5 as alpha)
    occ            : v1 ablation (amodal NEAREST-resized -> spurious ~1px ring)
    prod           : alpha only (188/255, no occluded pixels) = pre-2026-09-30
                     behaviour, which on FORGE3DBench IGNORED the occlusion masks.
The RGB image fed to the pipeline is the RGBA composited on black in every mode
(the pipeline re-masks bg/occluded internally anyway).
FIX 2026-09-30: before this date load_rgb_and_mask was used unconditionally, so
no pixel was ever marked occluded even on FORGE3DBench (see AMODAL3R_FIXES.md).

View-count generic: --views N selects the first N of [front, side, back, oside]
and calls run_multi_image with a list of that length (N=1 -> single-view list of
length 1; the same code path). To extend to MORE views, add tags to VIEW_TAGS.

Output GLB (textured, via the official postprocessing_utils.to_glb) is in the
model's own TRELLIS canonical frame -> it is a post-hoc baseline that must be run
through metrics/align_baselines.py before metrics/appeval/evaluate_appforce.py
(exactly like ReconViaGen / Pixal3D).

CLI (matches the existing batch drivers)
----------------------------------------
python batch_amodal3r.py --selection SEL.json --inputs INPUTDIR --exp EXPDIR
    --out OUTDIR --views {1|2|4} [--seed 42] [--limit N]
    [--gpu G] [--shard i --nshards n]

Idempotent (skips existing OUTDIR/<obj>.glb), per-object try/except, atomic write,
round-robin object sharding.  Re-execs itself into the shared reconviagen env
(== amodal3r/env) with the repo on PYTHONPATH.
"""
import os
import sys

AMODAL_ROOT = "/lp-dev/jonghoon/mv-mesh/amodal3r"
REPO = os.path.join(AMODAL_ROOT, "repo")
ENV = "/lp-dev/jonghoon/mv-mesh/envs/reconviagen"   # amodal3r/env -> this
ENV_PY = os.path.join(ENV, "bin", "python")
HF = "/lp-dev/jonghoon/mv-mesh/hf_cache"

_ENV_VARS = {
    "PYTHONUNBUFFERED": "1",
    "OMP_NUM_THREADS": "4",
    "MKL_NUM_THREADS": "4",
    "SPCONV_ALGO": "native",
    "ATTN_BACKEND": "flash-attn",
    "PYOPENGL_PLATFORM": "egl",
    "MKL_THREADING_LAYER": "GNU",
    "MKL_SERVICE_FORCE_INTEL": "0",
    "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
    "HF_HOME": HF,
    "HUGGINGFACE_HUB_CACHE": HF,
    "HF_HUB_CACHE": HF,
    "TORCH_HOME": os.path.join(HF, "torch"),
}


def _parse_gpu_from_argv():
    for i, a in enumerate(sys.argv):
        if a == "--gpu" and i + 1 < len(sys.argv):
            return sys.argv[i + 1]
        if a.startswith("--gpu="):
            return a.split("=", 1)[1]
    return None


def _reexec_in_env():
    env = dict(os.environ)
    for k, v in _ENV_VARS.items():
        env[k] = v
    env["CONDA_PREFIX"] = ENV
    env["PATH"] = os.path.join(ENV, "bin") + os.pathsep + env.get("PATH", "")
    env["PYTHONPATH"] = REPO + os.pathsep + env.get("PYTHONPATH", "")
    if not env.get("CUDA_VISIBLE_DEVICES"):
        g = _parse_gpu_from_argv()
        if g:
            env["CUDA_VISIBLE_DEVICES"] = g
    env["_AMODAL3R_BATCH_INENV"] = "1"
    print(f"[a3r] re-exec in {ENV_PY} (CUDA={env.get('CUDA_VISIBLE_DEVICES')})",
          flush=True)
    os.execve(ENV_PY, [ENV_PY, os.path.abspath(__file__)] + sys.argv[1:], env)


if not os.environ.get("_AMODAL3R_BATCH_INENV"):
    if not os.path.isfile(ENV_PY):
        sys.exit(f"ERROR: env python not found: {ENV_PY}")
    _reexec_in_env()

for _k, _v in _ENV_VARS.items():
    os.environ.setdefault(_k, _v)

import argparse   # noqa: E402
import json       # noqa: E402
import time       # noqa: E402
import traceback  # noqa: E402
from pathlib import Path  # noqa: E402

os.chdir(AMODAL_ROOT)
sys.path.insert(0, REPO)

import numpy as np           # noqa: E402
from PIL import Image        # noqa: E402
import torch                 # noqa: E402

VIEW_TAGS = {1: ["front"], 2: ["front", "side"],
             4: ["front", "side", "back", "oside"],
             8: ["front", "side", "back", "oside", "top", "bottom", "top2", "bottom2"]}
WORK_RES = 518   # feed all views at the model's working resolution (square)


def load_rgb_and_mask(png_path, res=WORK_RES):
    """RGBA crop -> (RGB PIL on black, L 3-value mask PIL) at res x res.

    alpha>0 -> visible(188); alpha==0 -> background(255). No occluded pixels
    on clean data (so the pipeline's masks_occ comes out all-zero)."""
    im = Image.open(png_path).convert("RGBA").resize((res, res), Image.LANCZOS)
    a = np.asarray(im).astype(np.float32) / 255.0
    alpha = a[:, :, 3:4]
    rgb = (a[:, :, :3] * alpha * 255.0).astype(np.uint8)   # composite on black
    rgb_img = Image.fromarray(rgb, mode="RGB")
    vis = (a[:, :, 3] > 0.5)
    mask = np.where(vis, 188, 255).astype(np.uint8)
    mask_img = Image.fromarray(mask, mode="L")
    return rgb_img, mask_img


# --- occlusion-aware 3-value mask (FIX 2026-09-30; functions verbatim from the
# --- validated occmask.py used for the 808-object FORGE3DBench rerun) ---------
def prod_rgb_mask(png_path, res=WORK_RES):
    im = Image.open(png_path).convert("RGBA").resize((res, res), Image.LANCZOS)
    a = np.asarray(im).astype(np.float32) / 255.0
    alpha = a[:, :, 3:4]
    rgb = (a[:, :, :3] * alpha * 255.0).astype(np.uint8)
    L = np.where(a[:, :, 3] > 0.5, 188, 255).astype(np.uint8)
    return rgb, L


def amodal_crop(src_scene, obj_id, cam, bbox, res=WORK_RES):
    y0, y1, x0, x1 = [int(v) for v in bbox]
    am = np.array(Image.open(f"{src_scene}/mask_amodal/cam{int(cam):02d}/{obj_id}.png").convert("L")) > 0
    am = am[y0:y1, x0:x1]
    return np.array(Image.fromarray(am.astype(np.uint8) * 255).resize((res, res), Image.NEAREST)) > 127


def occ2_rgb_mask(png_path, npz_path, src_scene, obj_id, cam, res=WORK_RES):
    """FINAL fix (v2).  occluded = (amodal & ~visible) computed at the NATIVE crop resolution, then resampled with the
    SAME filter + threshold as the visible alpha (LANCZOS, >0.5).  v1 (occ_rgb_mask) resized the binary amodal with
    NEAREST but the alpha with LANCZOS, so every view got a spurious ~1-px 'occluded' ring around the visible
    silhouette."""
    rgb, L = prod_rgb_mask(png_path, res)
    y0, y1, x0, x1 = [int(v) for v in np.load(npz_path)["bbox"]]
    vis = np.asarray(Image.open(png_path).convert("RGBA"))[:, :, 3] > 0
    am = np.array(Image.open(f"{src_scene}/mask_amodal/cam{int(cam):02d}/{obj_id}.png").convert("L"))[y0:y1, x0:x1] > 0
    assert am.shape == vis.shape, (am.shape, vis.shape)
    occ = (am & ~vis).astype(np.uint8) * 255
    occr = np.asarray(Image.fromarray(occ, "L").resize((res, res), Image.LANCZOS)).astype(np.float32) / 255.0 > 0.5
    L[occr & (L == 255)] = 0
    return rgb, L, occr


def occ_rgb_mask(png_path, npz_path, src_scene, obj_id, cam, res=WORK_RES):
    rgb, L = prod_rgb_mask(png_path, res)
    bbox = np.load(npz_path)["bbox"]
    amr = amodal_crop(src_scene, obj_id, cam, bbox, res)
    L[amr & (L == 255)] = 0
    return rgb, L, amr


def main():
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--method", default="amodal3r",
                    help="(single method; kept for CLI parity with other drivers)")
    ap.add_argument("--selection", required=True)
    ap.add_argument("--inputs", required=True)
    ap.add_argument("--exp", required=True,
                    help="EXPDIR with inputs/ (unused directly here but kept for "
                         "CLI parity with the other baseline drivers)")
    ap.add_argument("--out", required=True)
    ap.add_argument("--views", type=int, choices=[1, 2, 4, 8], required=True)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--ss-steps", type=int, default=12)
    ap.add_argument("--slat-steps", type=int, default=12)
    ap.add_argument("--simplify", type=float, default=0.95)
    ap.add_argument("--texture-size", type=int, default=1024)
    ap.add_argument("--mode", default="stochastic",
                    choices=["stochastic", "multidiffusion"])
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--objects", nargs="+", default=None)
    ap.add_argument("--gpu", default=None)
    ap.add_argument("--mask-mode", choices=["auto", "occ2", "occ", "prod"], default="auto",
                    help="auto = occ2 if the selection entry has src_scene+obj_id (FORGE3DBench), else prod")
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--nshards", type=int, default=1)
    args = ap.parse_args()

    inputs_dir = os.path.abspath(args.inputs)
    out_dir = os.path.abspath(args.out)
    os.makedirs(out_dir, exist_ok=True)
    views_tags = VIEW_TAGS[args.views]

    data = json.loads(Path(args.selection).read_text())
    sels = data["selections"] if isinstance(data, dict) else data
    selmap = {s["object"]: s for s in sels}
    TAG_ORDER = ["front", "side", "back", "oside"]
    objects = [s["object"] for s in sels]
    if args.objects:
        want = set(args.objects)
        objects = [o for o in objects if o in want]
    if args.limit:
        objects = objects[: args.limit]
    if args.nshards > 1:
        objects = objects[args.shard::args.nshards]

    print(f"[a3r] method=amodal3r views={args.views} tags={views_tags} "
          f"mode={args.mode} mask_mode={args.mask_mode} seed={args.seed} shard={args.shard}/{args.nshards} "
          f"n_obj={len(objects)} CUDA={os.environ.get('CUDA_VISIBLE_DEVICES')}",
          flush=True)

    from amodal3r.pipelines import Amodal3RImageTo3DPipeline
    from amodal3r.utils import postprocessing_utils

    t0 = time.time()
    print("[a3r] loading Amodal3R pipeline (Sm0kyWu/Amodal3R) ...", flush=True)
    pipeline = Amodal3RImageTo3DPipeline.from_pretrained("Sm0kyWu/Amodal3R")
    pipeline.cuda()
    print(f"[a3r] pipeline ready in {time.time() - t0:.1f}s", flush=True)

    counts = {"ok": 0, "fail": 0, "skip": 0}
    for i, obj in enumerate(objects, 1):
        out_path = os.path.join(out_dir, f"{obj}.glb")
        if os.path.isfile(out_path):
            counts["skip"] += 1
            print(f"[a3r {i}/{len(objects)}] SKIP (exists) {obj}", flush=True)
            continue
        t1 = time.time()
        try:
            images, masks = [], []
            for tag in views_tags:
                p = os.path.join(inputs_dir, f"{obj}_{tag}.png")
                if not os.path.isfile(p):
                    raise FileNotFoundError(p)
                so = selmap[obj]
                mm = args.mask_mode
                if mm == "auto":
                    mm = "occ2" if ("src_scene" in so and "obj_id" in so) else "prod"
                if mm == "prod":
                    rgb, msk = load_rgb_and_mask(p)
                else:
                    ci = so["views4"][TAG_ORDER.index(tag)] if "views4" in so else so["cam"]
                    fn = occ2_rgb_mask if mm == "occ2" else occ_rgb_mask
                    rgb_a, L_a, _ = fn(
                        p, os.path.join(os.path.abspath(args.exp), "renders", f"{obj}_{tag}.npz"),
                        so["src_scene"], so["obj_id"], ci)
                    rgb, msk = Image.fromarray(rgb_a, mode="RGB"), Image.fromarray(L_a, mode="L")
                images.append(rgb)
                masks.append(msk)
            print(f"[a3r {i}/{len(objects)}] {obj} -> run_multi_image "
                  f"n_views={len(images)}", flush=True)
            outputs = pipeline.run_multi_image(
                images, masks,
                seed=args.seed,
                formats=["mesh", "gaussian"],
                mode=args.mode,
                sparse_structure_sampler_params={"steps": args.ss_steps,
                                                 "cfg_strength": 7.5},
                slat_sampler_params={"steps": args.slat_steps, "cfg_strength": 3},
            )
            glb = postprocessing_utils.to_glb(
                outputs["gaussian"][0], outputs["mesh"][0],
                simplify=args.simplify, texture_size=args.texture_size,
                verbose=False)
            nverts, nfaces = len(glb.vertices), len(glb.faces)
            if nverts == 0:
                raise RuntimeError("exported GLB has 0 vertices")
            tmp = out_path + ".tmp.glb"
            glb.export(tmp)
            os.replace(tmp, out_path)
            del outputs, glb, images, masks
            torch.cuda.empty_cache()
            counts["ok"] += 1
            sz = os.path.getsize(out_path)
            print(f"[a3r {i}/{len(objects)}] OK {obj} {time.time() - t1:.1f}s "
                  f"verts={nverts} faces={nfaces} -> {out_path} ({sz} bytes)",
                  flush=True)
        except Exception:
            counts["fail"] += 1
            traceback.print_exc()
            print(f"[a3r {i}/{len(objects)}] FAIL {obj} {time.time() - t1:.1f}s",
                  flush=True)
            try:
                torch.cuda.empty_cache()
            except Exception:
                pass

    print(f"AMODAL3R DONE ok={counts['ok']} fail={counts['fail']} "
          f"skip={counts['skip']} total={len(objects)}", flush=True)
    return 1 if counts["fail"] else 0


if __name__ == "__main__":
    sys.exit(main())