"""Real image -> occupancy grid, with a silhouette check.

Usage:  python scripts/run_structure.py [image_path] [--fov RAD] [--steps N]

The silhouette IoU at the end is the point. Pixal3D's whole claim is that
back-projected pixel features keep the reconstruction aligned to the source image, so
re-projecting the occupied voxels through the SAME camera should reproduce the input
matte. A high IoU means the proj conditioning is genuinely steering the flow; a low one
means we are producing a generic blob and the conditioning is not landing — which a
"it ran without crashing" check would happily miss.
"""

import argparse
import sys
import time
from pathlib import Path

import mlx.core as mx
import numpy as np
from PIL import Image

REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO))

from pixal3d_mlx import ss_dec, ss_flow                              # noqa: E402
from pixal3d_mlx.cond import ProjConditioner, preprocess_image       # noqa: E402
from pixal3d_mlx.pipeline import DEFAULT_FOV, image_to_occupancy, occupied_coords  # noqa: E402
from pixal3d_mlx.proj import ProjGrid, distance_from_fov             # noqa: E402

CK = Path("/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts")
W = REPO / "weights"
DEFAULT_IMAGE = REPO / "upstream" / "Pixal3D" / "assets" / "images" / "0_img.png"


def silhouette_iou(coords, image_path, fov, res=512, grid_res=64):
    """Re-project occupied voxels through the camera; IoU against the input matte."""
    img = Image.open(image_path)
    if img.mode != "RGBA":
        return None
    matte = np.asarray(preprocess_image(img.copy()).convert("L").resize((res, res))) > 8

    # occupancy voxel indices -> the same [-1,1]^3 lattice the conditioner projected
    lattice = ProjGrid(grid_resolution=grid_res, image_resolution=res)
    idx = np.asarray(coords)[:, 1:]                       # drop batch column
    flat = idx[:, 0] * grid_res**2 + idx[:, 1] * grid_res + idx[:, 2]
    pts = np.asarray(lattice.grid_points)[flat][None] / 2.0

    from pixal3d_mlx.proj import _FRONT_VIEW, project_points
    tm = _FRONT_VIEW.copy()
    tm[1, 3] = -distance_from_fov(fov, 1.0, res)
    px, _, _ = project_points(mx.array(pts), mx.array(tm[None]), fov, res)
    px = np.asarray(px)[0].astype(int)

    keep = (px[:, 0] >= 0) & (px[:, 0] < res) & (px[:, 1] >= 0) & (px[:, 1] < res)
    proj = np.zeros((res, res), bool)
    proj[px[keep, 1], px[keep, 0]] = True

    # voxels are coarse (64^3 -> ~8px), so dilate before comparing against a 512px matte
    from scipy.ndimage import binary_dilation
    proj = binary_dilation(proj, np.ones((9, 9), bool))
    inter = (proj & matte).sum()
    union = (proj | matte).sum()
    return inter / union if union else 0.0, proj, matte


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("image", nargs="?", default=str(DEFAULT_IMAGE))
    ap.add_argument("--fov", type=float, default=DEFAULT_FOV)
    ap.add_argument("--steps", type=int, default=None)
    ap.add_argument("--seed", type=int, default=0)
    a = ap.parse_args()

    t0 = time.time()
    flow, frep = ss_flow.load(CK / "ss_flow_img_dit_1_3B_64_bf16.safetensors",
                              CK / "ss_flow_img_dit_1_3B_64_bf16.json")
    dec, drep = ss_dec.load(W / "ss_dec_conv3d_16l8_fp16.safetensors",
                            CK / "ss_dec_conv3d_16l8_fp16.json")
    assert not frep["missing"] and not drep["missing"], (frep["missing"], drep["missing"])
    print(f"models      {time.time() - t0:6.1f}s   ss_flow {frep['params']}  ss_dec {drep['params']}")

    t0 = time.time()
    cond_model = ProjConditioner("ss")
    print(f"dinov3      {time.time() - t0:6.1f}s   {cond_model.encoder.device}")

    overrides = {"steps": a.steps} if a.steps else {}
    t0 = time.time()
    occ, latent, cond = image_to_occupancy(
        a.image, flow, dec, conditioner=cond_model,
        camera_angle_x=a.fov, seed=a.seed, **overrides,
    )
    t_gen = time.time() - t0

    coords = occupied_coords(occ)
    n_vox = coords.shape[0]
    print(f"structure   {t_gen:6.1f}s   occ {tuple(occ.shape)}  occupied {n_vox} "
          f"({100 * n_vox / 64 ** 3:.2f}% of 64^3)")
    print(f"peak mem    {mx.get_peak_memory() / 2 ** 30:6.1f} GB")
    print(f"cond        global {cond['global'].shape}  proj {cond['proj'].shape}")

    if n_vox == 0:
        print("\nEMPTY OCCUPANCY — the structure stage produced nothing")
        sys.exit(1)

    got = silhouette_iou(coords, a.image, a.fov)
    if got is None:
        print("\n(no alpha matte on input — skipping silhouette check)")
        return
    iou, proj, matte = got
    print(f"silhouette  IoU {iou:.3f}   (projected {proj.sum()} px vs matte {matte.sum()} px)")
    out = REPO / "silhouette_check.png"
    Image.fromarray(np.stack([proj * 255, matte * 255, np.zeros_like(proj, np.uint8)],
                             -1).astype(np.uint8)).save(out)
    print(f"            wrote {out}  (red=reconstruction, green=input, yellow=overlap)")


if __name__ == "__main__":
    main()