"""Image -> GLB through the full Pixal3D cascade, with a silhouette check.

Usage:  python scripts/image_to_mesh.py IMAGE [-o OUT.glb] [--fov RAD] [--seed N]

The silhouette IoU is the acceptance test: re-project the mesh through the same camera
and compare against the input matte. Pixal3D's entire claim is pixel alignment, so a
run that completes with a poor IoU has failed even though nothing raised.
"""

import argparse
import sys
import time
from pathlib import Path

import mlx.core as mx
import numpy as np
import trimesh

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

from pixal3d_mlx.mesh import to_camera_frame                       # noqa: E402
from pixal3d_mlx.models import load_all, normalization             # noqa: E402
from pixal3d_mlx.pipeline import DEFAULT_FOV, image_to_mesh        # noqa: E402

DEFAULT_IMAGE = REPO / "upstream" / "Pixal3D" / "assets" / "images" / "0_img.png"


def silhouette_iou(vertices, image_path, fov, res=512):
    """Re-project the mesh through the generating camera; IoU against the input matte.

    Vertices MUST be rotated into the camera frame first — o_voxel returns them in the
    voxel-grid frame while ProjGrid rotates its lattice before projecting.
    """
    from PIL import Image
    from scipy.ndimage import binary_dilation

    from pixal3d_mlx.cond import preprocess_image
    from pixal3d_mlx.proj import _FRONT_VIEW, distance_from_fov, project_points

    img = Image.open(image_path)
    if img.mode != "RGBA":
        return None
    matte = np.asarray(preprocess_image(img).convert("L").resize((res, res))) > 8

    tm = _FRONT_VIEW.copy()
    tm[1, 3] = -distance_from_fov(fov, 1.0, res)
    pts = to_camera_frame(vertices).astype(np.float32)[None]
    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
    proj = binary_dilation(proj, np.ones((3, 3), bool))
    return (proj & matte).sum() / (proj | matte).sum(), proj, matte


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("image", nargs="?", default=str(DEFAULT_IMAGE))
    ap.add_argument("-o", "--output", default=str(REPO / "output.glb"))
    ap.add_argument("--fov", type=float, default=DEFAULT_FOV)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--no-check", action="store_true")
    a = ap.parse_args()

    t = time.time()
    models = load_all()
    print(f"loaded models ({time.time() - t:.1f}s, lazy — weights fault in on first use)")

    v, f, info = image_to_mesh(a.image, models, camera_angle_x=a.fov, seed=a.seed,
                               normalization=normalization("shape"))

    mesh = trimesh.Trimesh(v.cpu().numpy(), f.cpu().numpy(), process=False)
    mesh.export(a.output)
    print(f"\nTOTAL {info['seconds']}s   peak {info['peak_gb']} GB   -> {a.output}")

    if not a.no_check:
        got = silhouette_iou(mesh.vertices, a.image, a.fov)
        if got is None:
            print("(input has no alpha matte — skipping silhouette check)")
        else:
            iou, _, _ = got
            print(f"silhouette IoU {iou:.3f}")
            if iou < 0.85:
                print("WARNING: low IoU — the reconstruction is not tracking the input")
                sys.exit(1)


if __name__ == "__main__":
    main()