"""The rest of the sparse layer set, in MLX.

Nothing here is CUDA-locked upstream — these are ordinary linear/norm/attention layers
that merely take a SparseTensor instead of a dense one. They are reimplemented rather
than adapted because upstream's versions inherit from torch modules.

Two normalisation shapes are easy to conflate, and upstream uses both:

  LayerNorm32          applied to `x.feats` directly -> per-voxel over channels.
  SparseGroupNorm32    reshapes [N_b, C] -> [1, C, N_b] per batch item, so statistics
                       are over (channels-in-group x voxels) WITHIN one batch item.
                       Getting this wrong is silent: shapes match either way.

Attention runs in `attn_mode="full"`, which upstream defines as full attention *within*
each batch item (never across). Batch rows are contiguous, so each item is one slice.
"""

from __future__ import annotations

from typing import Optional

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

from .conv import SubMConv3d
from .tensor import SparseTensor

# ---------------------------------------------------------------- primitives


class SparseLinear(nn.Module):
    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features, bias=bias)

    def __call__(self, x: SparseTensor) -> SparseTensor:
        return x.replace(self.linear(x.feats))


class LayerNorm32(nn.Module):
    """Per-voxel LayerNorm over channels; computed in fp32 as upstream does."""

    def __init__(self, dim: int, affine: bool = False, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.affine = affine
        if affine:
            self.weight = mx.ones((dim,))
            self.bias = mx.zeros((dim,))

    def __call__(self, feats: mx.array) -> mx.array:
        dt = feats.dtype
        f = feats.astype(mx.float32)
        mu = mx.mean(f, axis=-1, keepdims=True)
        var = mx.var(f, axis=-1, keepdims=True)
        f = (f - mu) * mx.rsqrt(var + self.eps)
        if self.affine:
            f = f * self.weight + self.bias
        return f.astype(dt)


class SparseGroupNorm32(nn.Module):
    """GroupNorm over (channels-in-group x voxels), per batch item. fp32 internally."""

    def __init__(self, num_groups: int, num_channels: int, eps: float = 1e-5):
        super().__init__()
        if num_channels % num_groups != 0:
            raise ValueError(f"{num_channels} channels not divisible by {num_groups}")
        self.num_groups = num_groups
        self.num_channels = num_channels
        self.eps = eps
        self.weight = mx.ones((num_channels,))
        self.bias = mx.zeros((num_channels,))

    def __call__(self, x: SparseTensor) -> SparseTensor:
        dt = x.feats.dtype
        g, c = self.num_groups, self.num_channels
        parts = []
        for sl in x.layout:
            f = x.feats[sl].astype(mx.float32)  # [n_b, C]
            n_b = f.shape[0]
            if n_b == 0:
                parts.append(f)
                continue
            # -> [G, (C/G)*n_b] so mean/var cover channels *and* voxels in the group
            grouped = f.T.reshape(g, (c // g) * n_b)
            mu = mx.mean(grouped, axis=1, keepdims=True)
            var = mx.var(grouped, axis=1, keepdims=True)
            grouped = (grouped - mu) * mx.rsqrt(var + self.eps)
            f = grouped.reshape(c, n_b).T
            parts.append(f * self.weight + self.bias)
        out = parts[0] if len(parts) == 1 else mx.concatenate(parts, axis=0)
        return x.replace(out.astype(dt))


class SparseSiLU(nn.Module):
    def __call__(self, x: SparseTensor) -> SparseTensor:
        return x.replace(nn.silu(x.feats))


class SparseGELU(nn.Module):
    def __call__(self, x: SparseTensor) -> SparseTensor:
        return x.replace(nn.gelu(x.feats))


# ---------------------------------------------------------------- blocks


class SparseResBlock(nn.Module):
    """norm1(affine) -> silu -> conv1 -> norm2(no affine) -> silu -> conv2 + skip."""

    def __init__(self, channels: int, out_channels: Optional[int] = None):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels
        self.norm1 = LayerNorm32(channels, affine=True, eps=1e-6)
        self.norm2 = LayerNorm32(self.out_channels, affine=False, eps=1e-6)
        self.conv1 = SubMConv3d(channels, self.out_channels, 3)
        self.conv2 = SubMConv3d(self.out_channels, self.out_channels, 3)
        self.skip_connection = (
            SparseLinear(channels, self.out_channels)
            if channels != self.out_channels
            else None
        )

    def __call__(self, x: SparseTensor) -> SparseTensor:
        h = x.replace(self.norm1(x.feats))
        h = h.replace(nn.silu(h.feats))
        h = self.conv1(h)
        h = h.replace(self.norm2(h.feats))
        h = h.replace(nn.silu(h.feats))
        h = self.conv2(h)
        skip = self.skip_connection(x).feats if self.skip_connection else x.feats
        return h.replace(h.feats + skip)


class SparseFeedForwardNet(nn.Module):
    """Upstream is nn.Sequential(Linear, GELU, Linear), so its checkpoint keys are
    `mlp.mlp.0` and `mlp.mlp.2` — index 1 is the activation and carries no weights.
    Named `mlp_0`/`mlp_2` here because a Python list with a None hole does not survive
    MLX's parameter tree; the loader remaps the dotted indices onto these."""

    def __init__(self, channels: int, mlp_ratio: float = 4.0):
        super().__init__()
        hidden = int(channels * mlp_ratio)
        self.mlp_0 = nn.Linear(channels, hidden)
        self.mlp_2 = nn.Linear(hidden, channels)

    def __call__(self, x: SparseTensor) -> SparseTensor:
        return x.replace(self.mlp_2(nn.gelu_approx(self.mlp_0(x.feats))))


def _sdpa_per_batch(
    q: mx.array, k: mx.array, v: mx.array, layout_q, layout_kv, heads: int, scale: float
) -> mx.array:
    """Full attention inside each batch item. q/k/v are [N, H, D] flattened over batch."""
    outs = []
    for sq, skv in zip(layout_q, layout_kv):
        qi = q[sq].transpose(1, 0, 2)[None]  # [1, H, n, D]
        ki = k[skv].transpose(1, 0, 2)[None]
        vi = v[skv].transpose(1, 0, 2)[None]
        o = mx.fast.scaled_dot_product_attention(qi, ki, vi, scale=scale)
        outs.append(o[0].transpose(1, 0, 2))  # [n, H, D]
    return outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=0)


class SparseMultiHeadAttention(nn.Module):
    """attn_mode='full' only — the sole mode LATO.2's model code instantiates."""

    def __init__(
        self,
        channels: int,
        num_heads: int,
        ctx_channels: Optional[int] = None,
        attn_type: str = "self",
        qkv_bias: bool = True,
    ):
        super().__init__()
        if channels % num_heads != 0:
            raise ValueError(f"{channels} channels not divisible by {num_heads} heads")
        self.channels = channels
        self.num_heads = num_heads
        self.head_dim = channels // num_heads
        self.scale = self.head_dim**-0.5
        self._type = attn_type
        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
        if attn_type == "self":
            self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
        else:
            self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
            self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
        self.to_out = nn.Linear(channels, channels)

    def __call__(
        self, x: SparseTensor, context: Optional[SparseTensor] = None
    ) -> SparseTensor:
        n = x.feats.shape[0]
        h, d = self.num_heads, self.head_dim
        if self._type == "self":
            qkv = self.to_qkv(x.feats).reshape(n, 3, h, d)
            q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]
            lq = lkv = x.layout
        else:
            if context is None:
                raise ValueError("cross-attention needs a context")
            q = self.to_q(x.feats).reshape(n, h, d)
            m = context.feats.shape[0]
            kv = self.to_kv(context.feats).reshape(m, 2, h, d)
            k, v = kv[:, 0], kv[:, 1]
            lq, lkv = x.layout, context.layout
        o = _sdpa_per_batch(q, k, v, lq, lkv, h, self.scale)
        return x.replace(self.to_out(o.reshape(n, self.channels)))


class SparseTransformerBlock(nn.Module):
    """Pre-norm self-attention + FFN. Norms are non-affine (ln_affine=False upstream)."""

    def __init__(self, channels: int, num_heads: int, mlp_ratio: float = 4.0):
        super().__init__()
        self.norm1 = LayerNorm32(channels, affine=False, eps=1e-6)
        self.norm2 = LayerNorm32(channels, affine=False, eps=1e-6)
        self.attn = SparseMultiHeadAttention(channels, num_heads)
        self.mlp = SparseFeedForwardNet(channels, mlp_ratio)

    def __call__(self, x: SparseTensor) -> SparseTensor:
        h = self.attn(x.replace(self.norm1(x.feats)))
        x = x.replace(x.feats + h.feats)
        h = self.mlp(x.replace(self.norm2(x.feats)))
        return x.replace(x.feats + h.feats)


class SparseTransformerCrossBlock(nn.Module):
    """Pre-norm self-attn -> cross-attn -> FFN."""

    def __init__(
        self, channels: int, ctx_channels: int, num_heads: int, mlp_ratio: float = 4.0
    ):
        super().__init__()
        self.norm1 = LayerNorm32(channels, affine=False, eps=1e-6)
        self.norm2 = LayerNorm32(channels, affine=False, eps=1e-6)
        self.norm3 = LayerNorm32(channels, affine=False, eps=1e-6)
        self.context_norm = LayerNorm32(ctx_channels, affine=False, eps=1e-6)
        self.self_attn = SparseMultiHeadAttention(channels, num_heads)
        self.cross_attn = SparseMultiHeadAttention(
            channels, num_heads, ctx_channels=ctx_channels, attn_type="cross"
        )
        self.mlp = SparseFeedForwardNet(channels, mlp_ratio)

    def __call__(self, x: SparseTensor, context: SparseTensor) -> SparseTensor:
        h = self.self_attn(x.replace(self.norm1(x.feats)))
        x = x.replace(x.feats + h.feats)
        ctx = context.replace(self.context_norm(context.feats))
        h = self.cross_attn(x.replace(self.norm2(x.feats)), ctx)
        x = x.replace(x.feats + h.feats)
        h = self.mlp(x.replace(self.norm3(x.feats)))
        return x.replace(x.feats + h.feats)


class SparseConvNeXtBlock3d(nn.Module):
    """conv -> norm -> MLP, residual. The decoder workhorse in Pixal3D.

    Note the ordering: convolution comes FIRST, before the norm — unlike the pre-norm
    SparseResBlock. And the residual adds the block input, not the post-conv tensor.
    """

    def __init__(self, channels: int, mlp_ratio: float = 4.0):
        super().__init__()
        self.channels = channels
        self.norm = LayerNorm32(channels, affine=True, eps=1e-6)
        self.conv = SubMConv3d(channels, channels, 3)
        hidden = int(channels * mlp_ratio)
        self.mlp_0 = nn.Linear(channels, hidden)
        self.mlp_2 = nn.Linear(hidden, channels)

    def __call__(self, x: SparseTensor) -> SparseTensor:
        h = self.conv(x)
        h = h.replace(self.norm(h.feats))
        h = h.replace(self.mlp_2(nn.silu(self.mlp_0(h.feats))))
        return h.replace(h.feats + x.feats)


class SparseResBlockC2S3d(nn.Module):
    """Upsampling residual block: conv to 8x channels, then channel->spatial.

    `conv1` widens to `out_channels * 8` so that `channel2spatial` can redistribute those
    channels into the 8 children of each voxel. Which children exist is decided by the
    `subdiv` mask — predicted by `to_subdiv` when `pred_subdiv` is set — so the occupied
    set grows selectively rather than always x8.

    The skip path is a `repeat_interleave`, not a linear: the input is upsampled by the
    same channel->spatial step and its channels are repeated to match `out_channels`.
    """

    def __init__(
        self,
        channels: int,
        out_channels: Optional[int] = None,
        pred_subdiv: bool = False,
    ):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels
        self.pred_subdiv = pred_subdiv
        self.norm1 = LayerNorm32(channels, affine=True, eps=1e-6)
        self.norm2 = LayerNorm32(self.out_channels, affine=False, eps=1e-6)
        self.conv1 = SubMConv3d(channels, self.out_channels * 8, 3)
        self.conv2 = SubMConv3d(self.out_channels, self.out_channels, 3)
        if pred_subdiv:
            self.to_subdiv = SparseLinear(channels, 8)

    def __call__(self, x: SparseTensor, subdiv: Optional[SparseTensor] = None):
        from .tensor import channel2spatial

        if self.pred_subdiv:
            subdiv = self.to_subdiv(x)
        mask = subdiv.replace(subdiv.feats > 0) if subdiv is not None else None

        h = x.replace(self.norm1(x.feats))
        h = h.replace(nn.silu(h.feats))
        h = self.conv1(h)
        h = channel2spatial(h, mask, 2)
        xu = channel2spatial(x, mask, 2)

        h = h.replace(self.norm2(h.feats))
        h = h.replace(nn.silu(h.feats))
        h = self.conv2(h)

        reps = self.out_channels // max(self.channels // 8, 1)
        skip = mx.repeat(xu.feats, reps, axis=1) if reps > 1 else xu.feats
        out = h.replace(h.feats + skip)
        return (out, subdiv) if self.pred_subdiv else out