"""Pixal3D SparseStructureFlowModel in MLX.

The smallest of Pixal3D's four flow checkpoints (5.0 GB) and the one that exercises
every DiT feature the others use — RoPE, per-head q/k RMS norm, shared AdaLN modulation,
and `image_attn_mode="proj"` conditioning — so getting this one loading and running
validates the shared DiT core against real trained weights.

Despite the name it is fully DENSE: it operates on a 16³ voxel grid flattened to 4096
tokens. No sparse convolution anywhere (the checkpoint contains zero rank-5 tensors).

Config comes from the sibling JSON, not from constructor defaults:
    resolution 16, in/out 8ch, model 1536, cond 1024, 30 blocks, 12 heads,
    mlp_ratio 5.3334, pe_mode rope, share_mod, qk_rms_norm (+cross), image_attn_mode proj
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Optional

import mlx.core as mx
import mlx.nn as nn

from trellis_sparse_mlx.dit import ModulatedTransformerCrossBlock, TimestepEmbedder


class SparseStructureFlowModel(nn.Module):
    def __init__(
        self,
        resolution: int = 16,
        in_channels: int = 8,
        out_channels: int = 8,
        model_channels: int = 1536,
        cond_channels: int = 1024,
        num_blocks: int = 30,
        num_heads: int = 12,
        mlp_ratio: float = 5.3334,
        share_mod: bool = True,
        qk_rms_norm: bool = True,
        qk_rms_norm_cross: bool = True,
        image_attn_mode: str = "proj",
        proj_in_channels: Optional[int] = None,
        pe_mode: str = "rope",
        **_ignored,
    ):
        super().__init__()
        self.resolution = resolution
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.model_channels = model_channels
        self.share_mod = share_mod
        self.pe_mode = pe_mode

        self.t_embedder = TimestepEmbedder(model_channels)
        if share_mod:
            # upstream is Sequential(SiLU, Linear) -> checkpoint key is adaLN_modulation.1
            self.adaLN_modulation = nn.Linear(model_channels, 6 * model_channels)
        self.input_layer = nn.Linear(in_channels, model_channels)
        self.blocks = [
            ModulatedTransformerCrossBlock(
                model_channels,
                cond_channels,
                num_heads=num_heads,
                mlp_ratio=mlp_ratio,
                share_mod=share_mod,
                use_rope=(pe_mode == "rope"),
                qk_rms_norm=qk_rms_norm,
                qk_rms_norm_cross=qk_rms_norm_cross,
                image_attn_mode=image_attn_mode,
                proj_in_channels=proj_in_channels,
            )
            for _ in range(num_blocks)
        ]
        self.out_layer = nn.Linear(model_channels, out_channels)
        self.rope_phases: Optional[mx.array] = None  # loaded from the checkpoint

    def __call__(self, x: mx.array, t: mx.array, cond) -> mx.array:
        """x: [B, C, D, H, W]; t: [B]; cond: dict/tuple of (global, proj)."""
        b, c = x.shape[0], x.shape[1]
        if c != self.in_channels or list(x.shape[2:]) != [self.resolution] * 3:
            raise ValueError(
                f"expected [B,{self.in_channels},{self.resolution}^3], got {x.shape}"
            )
        h = x.reshape(b, c, -1).transpose(0, 2, 1)  # [B, N, C]
        h = self.input_layer(h)

        t_emb = self.t_embedder(t)
        if self.share_mod:
            t_emb = self.adaLN_modulation(nn.silu(t_emb))

        phases = self.rope_phases
        for blk in self.blocks:
            h = blk(h, t_emb, cond, phases=phases)

        # Parameterless final LayerNorm — `F.layer_norm(h, h.shape[-1:])` upstream.
        # Easy to miss and impossible to catch by weight-key matching: it has no
        # parameters, so a loader can report a perfect 700/700 with 0 missing and 0
        # unmapped while the model is still wrong. The residual stream leaves the last
        # block at std ~230; without this the output is ~200x too large.
        mu = mx.mean(h, axis=-1, keepdims=True)
        var = mx.var(h, axis=-1, keepdims=True)
        h = (h - mu) * mx.rsqrt(var + 1e-5)

        h = self.out_layer(h)
        return h.transpose(0, 2, 1).reshape(b, self.out_channels, *[self.resolution] * 3)


# ------------------------------------------------------------------ loading

def _remap(k: str) -> str:
    """Checkpoint key -> module path.

    Upstream wraps a couple of things in nn.Sequential, so its keys carry numeric indices
    where this implementation uses named attributes:
        t_embedder.mlp.{0,2}    -> t_embedder.mlp_{0,2}      (Linear, SiLU, Linear)
        blocks.N.mlp.mlp.{0,2}  -> blocks.N.mlp_{0,2}
        adaLN_modulation.1      -> adaLN_modulation          (index 0 is the SiLU)
    """
    k = k.replace("t_embedder.mlp.0.", "t_embedder.mlp_0.")
    k = k.replace("t_embedder.mlp.2.", "t_embedder.mlp_2.")
    k = k.replace(".mlp.mlp.0.", ".mlp_0.")
    k = k.replace(".mlp.mlp.2.", ".mlp_2.")
    k = k.replace("adaLN_modulation.1.", "adaLN_modulation.")
    return k


def load(weights_path: str | Path, config_path: str | Path | None = None):
    """Build from the sibling JSON config and load weights. Returns (model, report)."""
    wp = Path(weights_path)
    cp = Path(config_path) if config_path else wp.with_suffix(".json")
    cfg = json.loads(cp.read_text())
    args = dict(cfg.get("args", {}))
    args.pop("dtype", None)
    args.pop("initialization", None)

    model = SparseStructureFlowModel(**args)
    w = mx.load(str(wp))

    flat = dict(_flatten(model.parameters()))
    mapped, unmapped = {}, []
    for k, v in w.items():
        if k == "rope_phases":
            model.rope_phases = v
            continue
        m = _remap(k)
        if m in flat:
            if flat[m].shape != v.shape:
                raise ValueError(f"shape mismatch {k} -> {m}: {flat[m].shape} vs {v.shape}")
            mapped[m] = v
        else:
            unmapped.append(f"{k} -> {m}")
    missing = [k for k in flat if k not in mapped]
    if mapped:
        model.update(_unflatten(mapped))
    return model, {
        "config": cfg.get("name"),
        "loaded": len(mapped),
        "params": len(flat),
        "missing": missing,
        "unmapped": unmapped,
        "rope_phases": None if model.rope_phases is None else model.rope_phases.shape,
    }


def _flatten(tree, prefix=""):
    if isinstance(tree, dict):
        for k, v in tree.items():
            yield from _flatten(v, f"{prefix}{k}.")
    elif isinstance(tree, list):
        for i, v in enumerate(tree):
            yield from _flatten(v, f"{prefix}{i}.")
    elif isinstance(tree, mx.array):
        yield prefix[:-1], tree


def _unflatten(flat: dict):
    root: dict = {}
    for key, val in flat.items():
        parts = key.split(".")
        node = root
        for i, p in enumerate(parts[:-1]):
            nxt = parts[i + 1]
            default = [] if nxt.isdigit() else {}
            if isinstance(node, list):
                idx = int(p)
                while len(node) <= idx:
                    node.append({})
                if isinstance(default, list) and not isinstance(node[idx], list):
                    node[idx] = default
                node = node[idx]
            else:
                if p not in node or not isinstance(node[p], (dict, list)):
                    node[p] = default
                node = node[p]
        if isinstance(node, list):
            idx = int(parts[-1])
            while len(node) <= idx:
                node.append(None)
            node[idx] = val
        else:
            node[parts[-1]] = val
    return root