#!/usr/bin/env python
"""PASS 2 (TRAINING REPO) — trained SSFlowModel seed+imgcond -> stage-1 coords.

Three-pass SSFlow eval driver, pass 2 of 3 (see dump_seed.py for the design).

Replaces ONLY step 5 of ``tools/val_stage1_ref.py::stage1_and_seed`` (the
pretrained SAM3D ``ss_generator`` sampling) with the newly trained
``SSFlowModel`` (seed + single-view image condition).  Everything before it
(the oracle-visibility 64^3 seed -> ``z1``) arrives as a .npz from pass 1;
everything after it (stage-2 base noise + appearance seed) happens in pass 3.

MODEL
  Built exactly as ``mvsam3d/train/validate.py`` / ``.debug_ssflow_launch/
  batched_verify.py`` build it, from the ``config`` payload STORED INSIDE the
  checkpoint (authoritative: it is the config the run actually trained with),
  then the trained weights are loaded the way
  ``mvsam3d/train/val_appforce.py::_load_model`` does it:
      payload["trainable_state_dict"], strict=False, assert no unexpected keys.
  The frozen parts (DINO towers, PointPatchEmbed, d_embedder) come from the
  warm-start ``pretrained_ckpt`` (ss_generator.ckpt) applied in the constructor.
  ``ema.enable=false`` for this run, so there is NO ema sidecar and none is
  required (``--use-ema`` opts in if a sidecar ever appears).

CONDITION (ONE view: the PRIMARY/front tag)
  The training config pins ``data.num_views_min == num_views_max == 1``, so the
  condition is a single view, V=1, anchor == itself.  The tensors are built with
  the DATASET's OWN helpers (never re-implemented):
      SSImgCondDataset._full_frame_rgba(png, bbox, res_h, res_w)  -> full-frame RGBA
      ss_preprocess.preprocess_ss_view(rgba, pointmap=None)       -> image/rgb_image/
                                                                     mask/rgb_image_mask
      ss_preprocess.ss_crop_box(alpha, box_size_factor=1.2)       -> crop rect
      ss_imgcond_dataset.pad_square_rect(H, W)                    -> full rect
      mv_ss_condition.plucker_rays_for_view(...)                  -> the 4 RAY_KEYS grids
  This is exactly the ``_rgb_assets`` TOYS4K1K branch of the dataset, which
  already targets THIS eval layout: ``EXP/inputs/<obj>_<tag>.png`` (the tight
  CROP) + ``EXP/renders/<obj>_<tag>.npz`` carrying ``bbox`` in (y0, y1, x0, x1)
  order, ``res``, ``fx/fy/cx/cy`` and ``c2w_cv``.  The eval renders store
  fx/fy/cx/cy as scalars rather than a (V,3,3) ``K``; the intrinsics are read
  from those scalars, which is the same convention.

POINTMAP: ``cond_pointmap_mode="off"``.  The eval inputs are depth-free, so the
  two pointmap token blocks MUST be the official zero force-drop.  In
  ``MultiViewSSConditionEmbedder``, mode "off" makes ``_sample_pm_drop`` return
  all-True -> ``all_dropped`` -> ``torch.zeros(...)`` for both pointmap groups:
  BYTE-IDENTICAL to the mode-"auto" path the model trained under whenever
  ``has_pointmap == 0`` (the code comment states this and the branch confirms
  it).  Mode is not a parameter of the module, so switching it does not perturb
  any weight.

SEED LAYOUT: ``baf.grid_to_shape_latent(ss_enc, occ)`` returns
  ``z.view(1, 8, 4096).permute(0, 2, 1)``; ``mvsam3d.model.ss_vae.
  grid_latent_to_tokens`` returns ``z.reshape(B, 8, 4096).permute(0, 2, 1)`` of
  ``ss_encode(occ)`` == ``encoder(occ, sample_posterior=False, return_raw=False)``.
  Same encoder class, same architecture kwargs, same checkpoint file
  (md5 08536395562d0301911ed5c46ef0a539 in both trees) and the same defaults, so
  the two are THE SAME TENSOR -- no conversion is applied, z1 is fed to
  ``sample()`` as-is.

Output ``<out>/<obj>.npz``:
    coords          (N,4) int32  = out["coords"]           (PRIMARY)
    coords_original (M,4) int32  = out["coords_original"]  (debug)
    occ64_packed    uint8        np.packbits(logits[0,0] > 0)  -> unpack to 64^3
    downsample_factor, n_coords, ckpt_step, num_steps, view_tag

CLI:
  python ssflow_coords.py --exp EXPDIR --views 2 --seeddir <pass1 out> \
      --out <coords dir> --ckpt <step_XXXXXXX.pt> --gpu 7 [--shard i --nshards n]
NOTE ``--views`` selects the SAME object/tag table as pass 1/3, but the
CONDITION is always the single primary view (tag "front"); --views only affects
which seed dump dir / tag list the run is paired with.
"""
import os
import sys

REPO = "/lp-dev/jonghoon/mv-sam3d-6d-code/migrator/code/mv-sam3d-for-6d-v2"

# --- env preamble: must be set BEFORE torch / sam3d_objects import -----------
if not os.environ.get("_SSFLOW_COORDS_INENV"):
    import argparse as _ap
    _p = _ap.ArgumentParser(add_help=False)
    _p.add_argument("--gpu", default=None)
    _known, _ = _p.parse_known_args()
    env = dict(os.environ)
    env["_SSFLOW_COORDS_INENV"] = "1"
    env["PYTHONPATH"] = REPO
    env["SPARSE_ATTN_BACKEND"] = "sdpa"
    env["SPCONV_ALGO"] = "native"
    env["MIGRATOR_CHECKPOINTS"] = os.path.join(REPO, "migrator", "checkpoints")
    env["MIGRATOR_CACHE"] = "/lp-dev/jonghoon/mv-sam3d-6d-code/migrator/cache"
    env["OMP_NUM_THREADS"] = "4"
    env["MKL_NUM_THREADS"] = "4"
    gpu = _known.gpu if _known.gpu is not None else env.get("CUDA_VISIBLE_DEVICES")
    assert gpu, "pass --gpu <id> (or set CUDA_VISIBLE_DEVICES)"
    env["CUDA_VISIBLE_DEVICES"] = str(gpu)
    PY = "/lp-dev/jonghoon/mv-mesh/envs/mv-sam3d/bin/python"
    os.execve(PY, [PY, os.path.abspath(__file__)] + sys.argv[1:], env)

sys.path.insert(0, REPO)

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

import numpy as np  # noqa: E402
import torch        # noqa: E402

DEVICE = "cuda"
SEED = 42

# The tag table is shared with pass 1 / pass 3 (tools/val_stage1_ref.VIEW_TAGS).
VIEW_TAGS = ["front", "side", "back", "oside", "top", "bottom", "top2", "bottom2"]


def views_tags_for(n: int):
    assert 1 <= n <= len(VIEW_TAGS), f"--views {n} unsupported (max {len(VIEW_TAGS)})"
    return VIEW_TAGS[:n]


def _objects(args):
    """Verbatim from tools/val_stage1_ref.py::_objects."""
    if args.objects_file:
        objs = [l.strip() for l in open(args.objects_file) if l.strip()]
    elif args.objects:
        objs = list(args.objects)
    else:
        data = json.loads(Path(os.path.join(args.exp, "selection.json")).read_text())
        sels = data["selections"] if isinstance(data, dict) else data
        objs = [s["object"] for s in sels]
    if args.limit:
        objs = objs[: args.limit]
    if args.nshards > 1:
        objs = objs[args.shard::args.nshards]
    return objs


# --------------------------------------------------------------------------- #
# model
# --------------------------------------------------------------------------- #
def load_model(ckpt: str, device: str = DEVICE, use_ema: bool = False):
    """SSFlowModel from the checkpoint's own ``config`` + its trained weights.

    Mirrors validate.py / batched_verify.py for construction and
    val_appforce.py::_load_model for the weight load.
    """
    from mvsam3d.model.ssflow_model import SSFlowModel
    from mvsam3d import paths as _paths
    _paths.register_env()

    payload = torch.load(ckpt, map_location="cpu", weights_only=False)
    assert "trainable_state_dict" in payload, list(payload.keys())
    mc = dict(payload["config"]["model"])
    step = int(payload.get("step", -1))

    # Eval overrides (see the module docstring): the eval inputs are depth-free,
    # so the pointmap blocks are the official zero force-drop.
    mc["ss_mode"] = "imgcond"
    mc["cond_pointmap_mode"] = "off"
    mc["with_condition_embedder"] = True
    mc["use_checkpoint"] = False

    def _x(p):
        return os.path.expandvars(os.path.expanduser(p)) if p else p

    model = SSFlowModel(
        model_channels=mc["model_channels"], num_blocks=mc["num_blocks"],
        num_heads=mc["num_heads"], cond_channels=mc["cond_channels"],
        with_condition_embedder=True, sigma_min=mc["sigma_min"],
        time_scale=mc["time_scale"], timestep_sampler=mc["timestep_sampler"],
        timestep_lognorm_mean=mc["timestep_lognorm_mean"],
        timestep_lognorm_std=mc["timestep_lognorm_std"],
        pretrained_ckpt=_x(mc["pretrained_ckpt"]) or None,
        use_checkpoint=False,
        ss_mode="imgcond",
        cond_pointmap_mode="off",
        cond_dropout_prob=mc.get("cond_dropout_prob", 0.1),
        cond_enable_plucker=mc.get("cond_enable_plucker", True),
        cond_batch_towers=mc.get("cond_batch_towers", True),
    ).to(device)

    missing, unexpected = model.load_state_dict(payload["trainable_state_dict"],
                                                strict=False)
    assert not unexpected, unexpected[:10]
    print(f"[ssflow] loaded {ckpt} (step {step}, "
          f"{len(payload['trainable_state_dict'])} trained tensors, "
          f"{len(missing)} frozen/pretrained tensors left untouched)", flush=True)
    if use_ema:
        from mvsam3d.train.ema import EmaShadow, load_ema_sidecar
        ema = EmaShadow(model, rate=1.0)
        if load_ema_sidecar(ckpt, ema, model):
            with torch.no_grad():
                for n, p in model.named_parameters():
                    if n in ema.shadow:
                        p.copy_(ema.shadow[n])
            print("[ssflow] EMA sidecar applied", flush=True)
        else:
            print("[ssflow] no EMA sidecar (ema.enable=false for this run); "
                  "using raw weights", flush=True)
    model.eval()

    from mvsam3d.train.validate import _get_decoder
    decoder = _get_decoder(torch.device(device), _x(mc["ss_decoder_ckpt"]))
    return model, decoder, step


# --------------------------------------------------------------------------- #
# condition (V=1, the PRIMARY view) — built with the DATASET's own helpers
# --------------------------------------------------------------------------- #
def build_cond(model, exp: str, obj: str, tag: str, device: str = DEVICE):
    from mvsam3d.data.ss_imgcond_dataset import (
        SSImgCondDataset, pad_square_rect, COND_KEYS_BASE)
    from mvsam3d.data.ss_preprocess import preprocess_ss_view, ss_crop_box
    from mvsam3d.model.mv_ss_condition import plucker_rays_for_view, RAY_KEYS

    png = os.path.join(exp, "inputs", f"{obj}_{tag}.png")
    rz = np.load(os.path.join(exp, "renders", f"{obj}_{tag}.npz"))
    if not os.path.isfile(png):
        raise FileNotFoundError(png)
    bbox = rz["bbox"].astype(np.float64)        # (y0, y1, x0, x1) — dataset order
    res = int(rz["res"])
    fx, fy = float(rz["fx"]), float(rz["fy"])
    cx, cy = float(rz["cx"]), float(rz["cy"])
    c2w = torch.from_numpy(np.asarray(rz["c2w_cv"], dtype=np.float64))

    # crop -> full frame (dataset staticmethod; square store => res_h == res_w)
    rgba = SSImgCondDataset._full_frame_rgba(png, bbox, res, res)
    out = preprocess_ss_view(rgba, pointmap=None)          # 4 base keys, no batch dim

    crop_rect = ss_crop_box(rgba[..., 3], box_size_factor=1.2)
    full_rect = pad_square_rect(res, res)
    # V == 1 -> the anchor IS this view, so T = I (dataset: anchor_c2w = c2w[sub[0]])
    r = plucker_rays_for_view(c2w, c2w, fx, fy, cx, cy, crop_rect, full_rect)

    cond_inputs = {k: out[k].unsqueeze(0).unsqueeze(0).to(device)   # [1,1,C,518,518]
                   for k in COND_KEYS_BASE}
    rays = {k: r[k].unsqueeze(0).unsqueeze(0).to(device)            # [1,1,n*n,6]
            for k in RAY_KEYS}
    # pointmap keys deliberately ABSENT: cond_pointmap_mode='off' emits the
    # official zero force-drop blocks for both pointmap groups.
    cond = model.encode_condition_mv(cond_inputs, rays=rays,
                                     has_pm=torch.zeros(1, device=device))
    return cond


# --------------------------------------------------------------------------- #
def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--exp", required=True)
    ap.add_argument("--views", type=int, default=2)
    ap.add_argument("--seeddir", required=True, help="pass-1 dump dir (<obj>.npz with z1)")
    ap.add_argument("--out", required=True, help="coords dir (<obj>.npz each)")
    ap.add_argument("--ckpt", required=True, help="step_XXXXXXX.pt")
    ap.add_argument("--gpu", default=None, help="single GPU id (consumed by the preamble)")
    ap.add_argument("--num-steps", type=int, default=25)
    ap.add_argument("--seed", type=int, default=SEED)
    ap.add_argument("--use-ema", action="store_true",
                    help="apply an EMA sidecar if one exists (this run has none)")
    ap.add_argument("--objects", nargs="+", default=None)
    ap.add_argument("--objects-file", default=None)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--nshards", type=int, default=1)
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()

    exp = os.path.abspath(args.exp)
    views_tags = views_tags_for(args.views)
    tag = views_tags[0]                       # the PRIMARY view (V=1 condition)
    seeddir = Path(args.seeddir)
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    objs = _objects(args)
    print(f"[ssflow] exp={exp} views={args.views} cond_tag={tag} n_obj={len(objs)} "
          f"shard={args.shard}/{args.nshards} steps={args.num_steps} "
          f"CUDA={os.environ.get('CUDA_VISIBLE_DEVICES')}", flush=True)

    t0 = time.time()
    model, decoder, step = load_model(args.ckpt, DEVICE, use_ema=args.use_ema)
    print(f"[ssflow] model ready {time.time() - t0:.1f}s", flush=True)

    n_ok = n_fail = n_skip = 0
    for i, obj in enumerate(objs, 1):
        f = out / f"{obj}.npz"
        if f.is_file() and not args.force:
            n_skip += 1
            continue
        t1 = time.time()
        try:
            sz = np.load(seeddir / f"{obj}.npz")
            z1 = torch.from_numpy(sz["z1"].astype(np.float32)).to(DEVICE)
            assert z1.shape == (1, 4096, 8), z1.shape
            with torch.no_grad():
                cond = build_cond(model, exp, obj, tag, DEVICE)
                torch.manual_seed(args.seed)
                o = model.sample(z1, num_steps=args.num_steps, cond=cond,
                                 ss_decoder=decoder, use_fp32=True)
            coords = o["coords"].detach().cpu().numpy().astype(np.int32)
            coords_orig = o["coords_original"].detach().cpu().numpy().astype(np.int32)
            occ = (o["logits"][0, 0] > 0).detach().cpu().numpy()
            tmp = str(f) + ".tmp.npz"
            np.savez_compressed(
                tmp, coords=coords, coords_original=coords_orig,
                occ64_packed=np.packbits(occ.reshape(-1)),
                occ64_shape=np.array(occ.shape, dtype=np.int32),
                downsample_factor=int(o["downsample_factor"]),
                n_coords=int(coords.shape[0]), ckpt_step=int(step),
                num_steps=int(args.num_steps), view_tag=np.array(tag),
                views=np.array(views_tags), seed=int(args.seed))
            os.replace(tmp, f)
            n_ok += 1
            print(f"[ssflow {i}/{len(objs)}] OK {obj} N={coords.shape[0]} "
                  f"Norig={coords_orig.shape[0]} occ={int(occ.sum())} "
                  f"{time.time() - t1:.1f}s", flush=True)
        except Exception:
            n_fail += 1
            traceback.print_exc()
            print(f"[ssflow {i}/{len(objs)}] FAIL {obj}", flush=True)
        torch.cuda.empty_cache()
    print(f"SSFLOW_COORDS DONE ok={n_ok} fail={n_fail} skip={n_skip} "
          f"total={len(objs)}", flush=True)


if __name__ == "__main__":
    main()