# SPDX-License-Identifier: Apache-2.0
"""Host-side pre/post-processing for the MoGe-2 server.

Two jobs, both pure torch/numpy/PIL on the host:

* ``fit_to_canvas``: put ANY client image onto the single canonical canvas the
  decoder trace was captured for (default 1920x1080, the README's validated
  setting: ``num_tokens=1800`` -> 32x57 token grid). ``fit="pad"`` keeps the
  aspect ratio (letterbox, centred, neutral gray) so the geometry stays faithful;
  ``fit="stretch"`` squashes the image onto the canvas.
* ``canvas_to_camera`` + ``to_original``: the tail of the upstream
  ``MoGeModel.infer()`` (moge/model/v2.py) -- recover focal + z-shift from the
  affine point map, build normalized intrinsics, apply the metric scale -- then
  crop the letterbox region away, resize to the client's original resolution,
  re-project points from the resized depth (``force_projection``) and apply the
  validity mask.

Imports of the vendored ``moge`` / ``utils3d`` happen inside the functions:
importing ``tt_moge`` puts them on ``sys.path`` and their import chain pulls in
cv2/scipy, so nothing here runs at module import.
"""
from __future__ import annotations

import math
from typing import Dict, Optional, Tuple

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

Region = Tuple[int, int, int, int]  # x0, y0, x1, y1 on the canvas (pixels)

PAD_GRAY = (128, 128, 128)


def token_grid(width: int, height: int, num_tokens: int) -> Tuple[int, int]:
    """(base_h, base_w) MoGe uses for an image of this size -- mirrors MoGeModel.forward."""
    aspect_ratio = width / height
    return round((num_tokens / aspect_ratio) ** 0.5), round((num_tokens * aspect_ratio) ** 0.5)


def fit_to_canvas(im: Image.Image, canvas: Tuple[int, int], fit: str,
                  contiguous: bool = True) -> Tuple[torch.Tensor, Region]:
    """RGB PIL image -> (float32 [1,3,Hc,Wc] in [0,1], region the image occupies).

    ``canvas`` is (width, height). ``contiguous=False`` returns the channels-last view (same values,
    no NCHW copy; the C++ input prep reads any strides).
    """
    cw, ch = canvas
    if fit == "stretch":
        pasted = im.resize((cw, ch), Image.BILINEAR)
        region: Region = (0, 0, cw, ch)
    elif fit == "pad":
        s = min(cw / im.width, ch / im.height)
        rw, rh = max(1, round(im.width * s)), max(1, round(im.height * s))
        x0, y0 = (cw - rw) // 2, (ch - rh) // 2
        pasted = Image.new("RGB", (cw, ch), PAD_GRAY)
        pasted.paste(im.resize((rw, rh), Image.BILINEAR), (x0, y0))
        region = (x0, y0, x0 + rw, y0 + rh)
    else:
        raise ValueError(f"unknown fit mode {fit!r}")
    # one pass: float32(u8) / 255 in float32 (== np.asarray(pasted, np.float32) / 255.0, bit for bit)
    arr = np.true_divide(np.asarray(pasted), np.float32(255.0), dtype=np.float32)
    x = torch.from_numpy(arr).permute(2, 0, 1)[None]
    return (x.contiguous() if contiguous else x), region


def region_mask(canvas: Tuple[int, int], region: Region) -> torch.Tensor:
    """bool [Hc, Wc]: True inside the pasted image, False on the letterbox bars."""
    cw, ch = canvas
    m = torch.zeros(ch, cw, dtype=torch.bool)
    x0, y0, x1, y1 = region
    m[y0:y1, x0:x1] = True
    return m


@torch.inference_mode()
def canvas_to_camera(
    out: Dict[str, torch.Tensor],
    valid: Optional[torch.Tensor],
    fov_x: Optional[float] = None,
) -> Dict[str, torch.Tensor]:
    """``MoGeModel.infer()``'s recovery step on one canvas-sized TtMoGe output.

    ``out`` is the TtMoGe / MoGeModel.forward dict (batch 1): ``points`` [1,H,W,3]
    affine, ``normal`` [1,H,W,3], ``mask`` [1,H,W] sigmoid, ``metric_scale`` [1].
    ``valid`` (bool [H,W]) restricts the focal/shift solve to real (non-letterbox)
    pixels. Returns canvas-resolution tensors, metric-scaled, UNMASKED:
    ``depth`` [H,W], ``points`` [H,W,3] (affine map shifted into camera space --
    the exact upstream output when force_projection=False), ``normal`` [H,W,3],
    ``mask`` bool [H,W], ``intrinsics`` [3,3] normalized (cx=cy=0.5), plus scalars
    ``focal``, ``shift``, ``metric_scale``.
    """
    from moge.utils.geometry_torch import recover_focal_shift
    import utils3d

    points = out["points"][0].float().clone()      # (H, W, 3)
    normal = out["normal"][0].float()
    mask_prob = out["mask"][0].float()
    metric_scale = out["metric_scale"].float().reshape(-1)[0]
    height, width = points.shape[:2]
    aspect_ratio = width / height

    mask_binary = mask_prob > 0.5
    if valid is not None:
        mask_binary &= valid

    p, m = points[None], mask_binary[None]
    if fov_x is None:
        focal, shift = recover_focal_shift(p, m)
    else:
        focal = torch.tensor(
            [aspect_ratio / (1 + aspect_ratio ** 2) ** 0.5 / math.tan(math.radians(fov_x) / 2)],
            dtype=p.dtype,
        )
        _, shift = recover_focal_shift(p, m, focal=focal)
    focal = focal.reshape(-1)[0]
    shift = shift.reshape(-1)[0]

    fx = focal / 2 * (1 + aspect_ratio ** 2) ** 0.5 / aspect_ratio
    fy = focal / 2 * (1 + aspect_ratio ** 2) ** 0.5
    half = torch.tensor(0.5, dtype=p.dtype)
    intrinsics = utils3d.pt.intrinsics_from_focal_center(fx[None], fy[None], half, half)[0]

    points[..., 2] += shift
    mask_binary &= points[..., 2] > 0          # negative depth should never happen; drop it if it does
    depth = points[..., 2].clone()

    points *= metric_scale
    depth *= metric_scale
    return {
        "depth": depth, "points": points, "normal": normal, "mask": mask_binary,
        "intrinsics": intrinsics, "focal": focal, "shift": shift, "metric_scale": metric_scale,
    }


def fov_x_to_canvas(fov_x_deg: float, canvas: Tuple[int, int], region: Region) -> float:
    """A client's horizontal FoV (of ITS image) -> the horizontal FoV of the padded canvas.

    The image occupies ``region`` (width rw) on a canvas of width cw with the same focal
    length in pixels, so tan(fov_c/2) = tan(fov/2) * cw / rw. Identity for ``fit="stretch"``.
    """
    cw = canvas[0]
    rw = region[2] - region[0]
    return math.degrees(2 * math.atan(math.tan(math.radians(fov_x_deg) / 2) * cw / rw))


def crop_intrinsics(intrinsics: torch.Tensor, canvas: Tuple[int, int], region: Region) -> torch.Tensor:
    """Normalized canvas intrinsics -> normalized intrinsics of the cropped region.

    Normalized intrinsics are invariant to a uniform resize, so the result also holds
    for the client's original image (the region has its aspect ratio).
    """
    cw, ch = canvas
    x0, y0, x1, y1 = region
    rw, rh = x1 - x0, y1 - y0
    k = intrinsics.clone()
    k[0, 0] = intrinsics[0, 0] * cw / rw
    k[1, 1] = intrinsics[1, 1] * ch / rh
    k[0, 2] = (intrinsics[0, 2] * cw - x0) / rw
    k[1, 2] = (intrinsics[1, 2] * ch - y0) / rh
    return k


def _resize_map(t: torch.Tensor, size: Tuple[int, int], mode: str) -> torch.Tensor:
    """[H,W] or [H,W,C] -> same rank at ``size`` (height, width)."""
    squeeze = t.dim() == 2
    x = t[None, None] if squeeze else t.permute(2, 0, 1)[None]
    kw = {} if mode == "nearest" else {"align_corners": False}
    y = F.interpolate(x, size, mode=mode, **kw)
    return y[0, 0] if squeeze else y[0].permute(1, 2, 0)


@torch.inference_mode()
def to_original(
    cam: Dict[str, torch.Tensor],
    canvas: Tuple[int, int],
    region: Region,
    size: Tuple[int, int],
    force_projection: bool = True,
    apply_mask: bool = True,
) -> Dict[str, torch.Tensor]:
    """Crop the letterbox away, resize to the original (height, width), finish infer().

    Returns ``depth`` [H,W] f32 (metres), ``points`` [H,W,3] f32 (camera space, metres),
    ``normal`` [H,W,3] f32 (unit), ``mask`` bool [H,W], ``intrinsics`` [3,3] normalized.
    With ``apply_mask`` invalid pixels get depth/points = +inf and normal = 0, exactly
    like the upstream ``infer()``.
    """
    import utils3d

    x0, y0, x1, y1 = region
    sl = (slice(y0, y1), slice(x0, x1))
    depth = _resize_map(cam["depth"][sl], size, "bilinear")
    normal = _resize_map(cam["normal"][sl], size, "bilinear")
    normal = F.normalize(normal, dim=-1)
    mask = _resize_map(cam["mask"][sl].float(), size, "nearest") > 0.5
    intrinsics = crop_intrinsics(cam["intrinsics"], canvas, region)

    if force_projection:
        points = utils3d.pt.depth_map_to_point_map(depth, intrinsics)
    else:
        points = _resize_map(cam["points"][sl], size, "bilinear")

    mask &= torch.isfinite(depth) & (depth > 0)
    if apply_mask:
        inf = torch.tensor(float("inf"), dtype=depth.dtype)
        depth = torch.where(mask, depth, inf)
        points = torch.where(mask[..., None], points, inf)
        normal = torch.where(mask[..., None], normal, torch.zeros_like(normal))

    return {"depth": depth, "points": points, "normal": normal, "mask": mask, "intrinsics": intrinsics}


def fov_deg(intrinsics: torch.Tensor) -> Tuple[float, float]:
    """Horizontal / vertical FoV in degrees from normalized intrinsics."""
    fx, fy = float(intrinsics[0, 0]), float(intrinsics[1, 1])
    return math.degrees(2 * math.atan(0.5 / fx)), math.degrees(2 * math.atan(0.5 / fy))