"""Pixal3D SparseStructureDecoder in MLX — latent -> occupancy grid.

Fully dense: 3D convs on a small voxel grid, no sparse ops at all. Pairs with `ss_flow`
to complete the structure stage (image -> latent -> occupancy), which is what decides
which voxels exist before the SLAT stage runs.

Config (from the sibling JSON):
    out_channels 1, latent_channels 8, channels [512, 128, 32],
    num_res_blocks 2, num_res_blocks_middle 2, norm_type "layer"

MLX's Conv3d is channels-LAST (NDHWC) where torch is NCDHW, so tensors are carried in
channels-last form throughout and only transposed at the boundaries. The weight
converter already emits `[O, kz, ky, kx, I]` to match.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import List, Optional

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


def _cln(x: mx.array, weight: mx.array, bias: mx.array, eps: float = 1e-5) -> mx.array:
    """ChannelLayerNorm32 — LayerNorm over channels, computed in fp32.

    Upstream permutes channels to last, applies LayerNorm, permutes back. Data here is
    already channels-last, so it is a plain LayerNorm over the final axis.
    """
    dt = x.dtype
    f = x.astype(mx.float32)
    f = (f - mx.mean(f, -1, keepdims=True)) * mx.rsqrt(mx.var(f, -1, keepdims=True) + eps)
    return (f * weight + bias).astype(dt)


def pixel_shuffle_3d(x: mx.array, scale: int) -> mx.array:
    """Channels-last 3D pixel shuffle: [B,H,W,D,C*s^3] -> [B,H*s,W*s,D*s,C].

    Upstream works in NCDHW and permutes (0,1,5,2,6,3,7,4) after splitting the channel
    axis into (C_, s, s, s). Here the channel axis is last, so the split factors sit
    next to it and the interleave order must still be (H,s)(W,s)(D,s) — getting that
    order wrong scrambles the grid while preserving its shape.
    """
    b, h, w, d, c = x.shape
    c_ = c // scale**3
    x = x.reshape(b, h, w, d, c_, scale, scale, scale)
    x = x.transpose(0, 1, 5, 2, 6, 3, 7, 4)  # b, h, s, w, s, d, s, c_
    return x.reshape(b, h * scale, w * scale, d * scale, c_)


class ResBlock3d(nn.Module):
    def __init__(self, channels: int, out_channels: Optional[int] = None):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels
        self.norm1_w = mx.ones((channels,))
        self.norm1_b = mx.zeros((channels,))
        self.norm2_w = mx.ones((self.out_channels,))
        self.norm2_b = mx.zeros((self.out_channels,))
        self.conv1 = nn.Conv3d(channels, self.out_channels, 3, padding=1)
        self.conv2 = nn.Conv3d(self.out_channels, self.out_channels, 3, padding=1)
        self.skip_connection = (
            nn.Conv3d(channels, self.out_channels, 1)
            if channels != self.out_channels
            else None
        )

    def __call__(self, x: mx.array) -> mx.array:
        h = _cln(x, self.norm1_w, self.norm1_b)
        h = self.conv1(nn.silu(h))
        h = _cln(h, self.norm2_w, self.norm2_b)
        h = self.conv2(nn.silu(h))
        return h + (self.skip_connection(x) if self.skip_connection else x)


class UpsampleBlock3d(nn.Module):
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        self.conv = nn.Conv3d(in_channels, out_channels * 8, 3, padding=1)

    def __call__(self, x: mx.array) -> mx.array:
        return pixel_shuffle_3d(self.conv(x), 2)


class SparseStructureDecoder(nn.Module):
    def __init__(
        self,
        out_channels: int = 1,
        latent_channels: int = 8,
        num_res_blocks: int = 2,
        channels: List[int] = (512, 128, 32),
        num_res_blocks_middle: int = 2,
        **_ignored,
    ):
        super().__init__()
        channels = list(channels)
        self.channels = channels
        self.input_layer = nn.Conv3d(latent_channels, channels[0], 3, padding=1)
        self.middle_block = [
            ResBlock3d(channels[0], channels[0]) for _ in range(num_res_blocks_middle)
        ]
        blocks = []
        for i, ch in enumerate(channels):
            blocks += [ResBlock3d(ch, ch) for _ in range(num_res_blocks)]
            if i < len(channels) - 1:
                blocks.append(UpsampleBlock3d(ch, channels[i + 1]))
        self.blocks = blocks
        self.out_norm_w = mx.ones((channels[-1],))
        self.out_norm_b = mx.zeros((channels[-1],))
        self.out_conv = nn.Conv3d(channels[-1], out_channels, 3, padding=1)

    def __call__(self, z: mx.array) -> mx.array:
        """z: [B, C, D, H, W] (torch order, for a familiar interface)."""
        x = z.transpose(0, 2, 3, 4, 1)  # -> channels-last
        x = self.input_layer(x)
        for b in self.middle_block:
            x = b(x)
        for b in self.blocks:
            x = b(x)
        x = _cln(x, self.out_norm_w, self.out_norm_b)
        x = self.out_conv(nn.silu(x))
        return x.transpose(0, 4, 1, 2, 3)  # back to [B, C, D, H, W]


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

def _remap(k: str) -> str:
    """Upstream nests norms/convs inside Sequential and nn.Module attrs."""
    k = k.replace("out_layer.0.weight", "out_norm_w").replace("out_layer.0.bias", "out_norm_b")
    k = k.replace("out_layer.2.", "out_conv.")
    k = k.replace(".norm1.weight", ".norm1_w").replace(".norm1.bias", ".norm1_b")
    k = k.replace(".norm2.weight", ".norm2_w").replace(".norm2.bias", ".norm2_b")
    return k


def load(weights_path: str | Path, config_path: str | Path | None = None):
    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("use_fp16", None)
    model = SparseStructureDecoder(**args)

    w = mx.load(str(wp))
    flat = dict(_flatten(model.parameters()))
    mapped, unmapped = {}, []
    for k, v in w.items():
        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} {tuple(v.shape)}")
    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,
    }


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