"""Pixal3D stage wiring for MLX.

What runs today: the STRUCTURE stage, end to end on Metal —
    noise -> ss_flow (Euler sampling) -> latent -> ss_dec -> occupancy grid

What is still missing for a real image->3D run is the conditioning, not the models.
Pixal3D conditions on DINOv3 features in two forms:

    global : [B, M, 1024]  pooled image tokens, cross-attended
    proj   : [B, N, 1024]  VIEW-ALIGNED features, one token per spatial position,
                           projected and ADDED rather than attended

The `proj` half is the actual novelty — it back-projects pixel features into 3D through
the camera, which is what keeps silhouettes exact. Producing it needs the DINOv3 encoder
plus Pixal3D's `DinoV3ProjFeatureExtractor` camera logic.

Deliberate design note: DINOv3 is a stock ViT-L/16 run ONCE per image, not inside the
25-step denoising loop, so it is not worth porting to MLX — torch on MPS runs it fine and
using `transformers` gives exact parity with upstream for free. The effort belongs in the
30-block DiT loops, which are already here and verified. Note facebook/dinov3 is gated;
upstream itself falls back to the ungated `camenduru` mirror.
"""

from __future__ import annotations

from typing import Callable, Optional

import mlx.core as mx

from .sampler import FlowEulerSampler


def run_structure_stage(
    flow_model,
    decoder,
    cond,
    neg_cond=None,
    resolution: int = 16,
    latent_channels: int = 8,
    steps: int = 25,
    rescale_t: float = 3.0,
    guidance_strength: float = 1.0,
    guidance_interval=None,
    seed: int = 0,
    progress: Optional[Callable[[int, int], None]] = None,
):
    """noise -> ss_flow -> ss_dec. Returns (occupancy_logits, latent).

    `occupancy_logits` is [B, 1, R*4, R*4, R*4]; threshold at 0 for the occupied set that
    seeds the SLAT stage.
    """
    mx.random.seed(seed)
    noise = mx.random.normal((1, latent_channels, resolution, resolution, resolution))
    sampler = FlowEulerSampler()
    latent = sampler.sample(
        flow_model,
        noise,
        cond=cond,
        neg_cond=neg_cond,
        steps=steps,
        rescale_t=rescale_t,
        guidance_strength=guidance_strength,
        guidance_interval=guidance_interval,
        progress=progress,
    )
    occ = decoder(latent)
    mx.eval(occ)
    return occ, latent


def occupied_coords(occ: mx.array, threshold: float = 0.0) -> mx.array:
    """Occupancy logits -> int32 [N, 4] (batch, z, y, x) coords for the SLAT stage."""
    import numpy as np

    o = np.asarray(occ)[:, 0]  # [B, D, H, W]
    idx = np.argwhere(o > threshold).astype(np.int32)
    return mx.array(idx)