"""SparseTensor for MLX — a coords/feats pair, mirroring LATO.2's upstream container.

Upstream wraps either `spconv.SparseConvTensor` or `torchsparse.SparseTensor`; both are
CUDA-only, which is what blocks LATO.2 on Apple Silicon. Nothing about the *data* needs
CUDA — it is just coordinates plus features — so this is a plain MLX reimplementation.

Layout matches upstream exactly so weight conversion stays a straight mapping:
    coords : int32 [N, 4]  -> (batch, z, y, x)
    feats  : float  [N, C]
and rows belonging to one batch item are contiguous (upstream asserts this).
"""

from __future__ import annotations

from typing import List, Optional, Tuple

import mlx.core as mx
import numpy as np


class SparseTensor:
    """N non-empty voxels, each with a coordinate and a feature vector."""

    __slots__ = ("feats", "coords", "_scale", "_spatial_cache", "_layout")

    def __init__(
        self,
        feats: mx.array,
        coords: mx.array,
        scale: Tuple[int, int, int] = (1, 1, 1),
        spatial_cache: Optional[dict] = None,
        layout: Optional[List[slice]] = None,
    ):
        if feats.shape[0] != coords.shape[0]:
            raise ValueError(
                f"feats/coords length mismatch: {feats.shape[0]} vs {coords.shape[0]}"
            )
        if coords.ndim != 2 or coords.shape[1] != 4:
            raise ValueError(f"coords must be [N, 4] (batch,z,y,x), got {coords.shape}")
        self.feats = feats
        self.coords = coords
        self._scale = tuple(scale)
        # Indice maps are expensive to build and identical for every conv that shares a
        # coordinate set — upstream exploits this via spconv's `indice_key`. Same idea.
        self._spatial_cache = spatial_cache if spatial_cache is not None else {}
        self._layout = layout

    # -- basics ---------------------------------------------------------------
    @property
    def shape(self) -> Tuple[int, int]:
        return (self.batch_size, self.feats.shape[1])

    @property
    def batch_size(self) -> int:
        if self.coords.shape[0] == 0:
            return 0
        return int(mx.max(self.coords[:, 0]).item()) + 1

    @property
    def layout(self) -> List[slice]:
        """One slice per batch item. Relies on batch-contiguity, as upstream does."""
        if self._layout is None:
            b = np.asarray(self.coords[:, 0], dtype=np.int64)
            counts = np.bincount(b, minlength=self.batch_size)
            offs = np.cumsum(counts)
            self._layout = [
                slice(int(offs[i] - counts[i]), int(offs[i])) for i in range(len(counts))
            ]
        return self._layout

    def replace(self, feats: mx.array) -> "SparseTensor":
        """New tensor, same coordinates — so the indice-map cache stays valid."""
        return SparseTensor(
            feats,
            self.coords,
            scale=self._scale,
            spatial_cache=self._spatial_cache,
            layout=self._layout,
        )

    # -- cache ----------------------------------------------------------------
    def cache_get(self, key: str):
        return self._spatial_cache.get(key)

    def cache_put(self, key: str, value) -> None:
        self._spatial_cache[key] = value

    def __repr__(self) -> str:
        return (
            f"SparseTensor(N={self.coords.shape[0]}, C={self.feats.shape[1]}, "
            f"batch={self.batch_size}, scale={self._scale})"
        )


def subdivide(x: SparseTensor) -> SparseTensor:
    """Upsample ×2 by splitting each voxel into its 8 children (nearest-neighbour).

    Mirrors upstream `SparseSubdivide`: coords are doubled then offset by the unit
    cube, features are replicated. Child order is the C-order of nonzero(ones(2,2,2)),
    matching upstream's `torch.nonzero`, so replicated features line up identically.
    """
    n = x.coords.shape[0]
    offsets = np.stack(np.meshgrid(*[np.arange(2)] * 3, indexing="ij"), -1).reshape(-1, 3)
    offsets = np.concatenate([np.zeros((8, 1), dtype=offsets.dtype), offsets], axis=1)

    base = np.asarray(x.coords, dtype=np.int32).copy()
    base[:, 1:] *= 2
    new_coords = (base[:, None, :] + offsets[None, :, :]).reshape(n * 8, 4)

    new_feats = mx.repeat(x.feats, 8, axis=0)
    return SparseTensor(
        new_feats,
        mx.array(new_coords, dtype=mx.int32),
        scale=tuple(s * 2 for s in x._scale),
    )