"""Image -> conditioning dicts, the MLX side of upstream's `encode_image_proj`.

Produces exactly the two-key contract the DiT blocks expect:

    cond   = {'global': [B, 1+regs, D], 'proj': [B, R^3, C]}
    uncond = zeros of the same shapes        (classifier-free guidance negative)

`global` is cross-attended; `proj` is projected and ADDED, one token per voxel. The
token counts differ (5 vs R^3) and that asymmetry IS the architecture — see proj.py.
"""

from __future__ import annotations

from typing import Dict, Tuple

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

from .dino import DinoV3Encoder
from .proj import ProjGrid, distance_from_fov

# The four extractor configurations upstream builds in inference.py. Only `ss` is wired
# so far; the other three additionally need the NAF upsampler, which doubles their
# proj_channels to embed_dim*2 by concatenating a high-res branch.
CONFIGS = {
    "ss":          dict(image_size=512,  grid_resolution=16, use_naf_upsample=False),
    "shape_512":   dict(image_size=512,  grid_resolution=32, use_naf_upsample=True,  naf_target_size=512),
    "shape_1024":  dict(image_size=1024, grid_resolution=64, use_naf_upsample=True,  naf_target_size=512),
    "tex_1024":    dict(image_size=1024, grid_resolution=64, use_naf_upsample=True,  naf_target_size=1024),
}


def load_image(path: str, image_size: int) -> np.ndarray:
    """Path -> [1,3,S,S] float32 in [0,1], LANCZOS-resized like upstream."""
    img = Image.open(path).convert("RGB").resize((image_size, image_size), Image.LANCZOS)
    arr = np.asarray(img, dtype=np.float32) / 255.0
    return arr.transpose(2, 0, 1)[None]


class ProjConditioner:
    """One extractor stage: DINOv3 + a ProjGrid at that stage's voxel resolution."""

    def __init__(self, stage: str = "ss", device: str | None = None):
        cfg = CONFIGS[stage]
        if cfg.get("use_naf_upsample"):
            raise NotImplementedError(
                f"stage '{stage}' needs the NAF upsampler (proj_channels = embed_dim*2); "
                "only 'ss' is wired so far"
            )
        self.stage = stage
        self.image_size = cfg["image_size"]
        self.encoder = DinoV3Encoder(image_size=self.image_size, device=device)
        self.proj_grid = ProjGrid(grid_resolution=cfg["grid_resolution"],
                                  image_resolution=self.image_size)

    def __call__(self, image_chw01: np.ndarray, camera_angle_x: float,
                 distance: float | None = None, mesh_scale: float = 1.0
                 ) -> Tuple[Dict[str, mx.array], Dict[str, mx.array]]:
        if distance is None:
            distance = distance_from_fov(camera_angle_x, mesh_scale, self.image_size)

        z_global, z_patch = self.encoder(image_chw01)
        z_proj = self.proj_grid(z_patch, camera_angle_x, distance, mesh_scale)

        cond = {"global": z_global, "proj": z_proj}
        uncond = {"global": mx.zeros_like(z_global), "proj": mx.zeros_like(z_proj)}
        return cond, uncond