"""Per-image camera estimation with MoGe-2.

Everything downstream assumes the object exactly fills the frame at a known FOV:
`distance_from_fov` places the camera so a unit mesh spans the image, and the proj
grid back-projects through that camera. A fixed default FOV therefore does not just
change framing — it reconstructs the subject at the WRONG DEPTH SCALE, and the failure
is quiet: you get a clean, plausible mesh that is subtly the wrong shape.

Upstream estimates FOV per image with MoGe-2 and only falls back to a constant when
`--fov` is passed explicitly. This does the same. MoGe runs once per image in torch on
MPS, like DINOv3 and NAF — outside any denoising loop, so parity beats a port.
"""

from __future__ import annotations

import math

import numpy as np

MODEL_NAME = "Ruicheng/moge-2-vitl"
_model = None


def load_moge(device: str = "mps"):
    """Load MoGe-2 once per process."""
    global _model
    if _model is None:
        import torch
        from moge.model.v2 import MoGeModel

        _model = MoGeModel.from_pretrained(MODEL_NAME).eval().to(device)
        _model.requires_grad_(False)
    elif str(next(_model.parameters()).device).split(":")[0] != device:
        _model = _model.to(device)
    return _model


def estimate_fov(image_path: str, device: str = "mps") -> float:
    """Horizontal FOV in radians, from MoGe-2's predicted intrinsics.

    `intrinsics[0,0]` is fx NORMALISED by image width, so it must be multiplied back
    up before the arctan — the normalisation is easy to miss and yields a plausible
    but wrong angle if skipped.
    """
    import torch
    from PIL import Image

    img = Image.open(image_path).convert("RGB")
    width, _ = img.size
    arr = np.asarray(img, dtype=np.float32) / 255.0
    tensor = torch.from_numpy(arr).permute(2, 0, 1).to(device)

    model = load_moge(device)
    with torch.no_grad():
        out = model.infer(tensor)
    fx_norm = float(np.asarray(out["intrinsics"].squeeze().cpu())[0, 0])
    return 2.0 * math.atan(width / (2.0 * fx_norm * width))


def camera_for(image_path: str, fov: float | None = None, mesh_scale: float = 1.0,
               image_resolution: int = 512, device: str = "mps") -> dict:
    """(camera_angle_x, distance) for an image. Estimates FOV when `fov` is None."""
    from .proj import distance_from_fov

    estimated = fov is None
    if estimated:
        fov = estimate_fov(image_path, device)
    return {"camera_angle_x": float(fov),
            "distance": distance_from_fov(fov, mesh_scale, image_resolution),
            "estimated": estimated}