"""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 VarLenTensor:
    """Variable-length batched features: `feats` plus an explicit slice-per-batch layout.

    Pixal3D splits the abstraction this way — `VarLenTensor` is the sequence container and
    `SparseTensor(VarLenTensor)` adds voxel coordinates. LATO.2 has only the combined
    class. Modelling the split here lets one package serve both without either port
    having to adapt at every call site.
    """

    def __init__(self, feats: mx.array, layout: Optional[List[slice]] = None):
        self.feats = feats
        self._layout = layout
        self._spatial_cache: dict = {}

    @staticmethod
    def layout_from_seqlen(seqlen: List[int]) -> List[slice]:
        out, start = [], 0
        for n in seqlen:
            out.append(slice(start, start + n))
            start += n
        return out

    @staticmethod
    def from_tensor_list(tensors: List[mx.array]) -> "VarLenTensor":
        lay = VarLenTensor.layout_from_seqlen([t.shape[0] for t in tensors])
        return VarLenTensor(mx.concatenate(tensors, axis=0), lay)

    @property
    def layout(self) -> List[slice]:
        if self._layout is None:
            self._layout = [slice(0, self.feats.shape[0])]
        return self._layout

    # Pixal3D's spelling of the cache API; LATO.2 uses cache_get/cache_put. Both work.
    def get_spatial_cache(self, key: str):
        return self._spatial_cache.get(key)

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


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

    def __init__(
        self,
        feats: mx.array,
        coords: mx.array,
        scale: Tuple = (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, derived from coords. Relies on batch-contiguity.

        Overrides VarLenTensor.layout, which has no coordinates to derive from and so
        treats the whole tensor as a single sequence.
        """
        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 (LATO.2 spelling; see VarLenTensor for Pixal3D's) ---------------
    cache_get = VarLenTensor.get_spatial_cache
    cache_put = VarLenTensor.register_spatial_cache

    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),
    )


def downsample(x: "SparseTensor", factor: int = 2) -> "SparseTensor":
    """Downsample by `factor`, reducing colliding voxels with MAX.

    Upstream's docstring says "average pooling" but the implementation passes
    reduce="amax" (the `reduce='mean'` line is commented out). Following the code,
    not the docstring — mean vs max here is numerically silent in shape and would
    quietly change every downsampled feature.

    Output coordinates come out sorted by the same packed code upstream sorts on, so
    rows stay batch-contiguous as SparseTensor requires.
    """
    import numpy as _np

    c = _np.asarray(x.coords, dtype=_np.int64).copy()
    c[:, 1:] //= factor

    maxs = c[:, 1:].max(axis=0) + 1
    # OFFSET = reversed cumprod, matching upstream's packing
    off = _np.array(
        [maxs[0] * maxs[1] * maxs[2], maxs[1] * maxs[2], maxs[2], 1], dtype=_np.int64
    )
    code = (c * off).sum(axis=1)

    uniq, inv = _np.unique(code, return_inverse=True)
    feats = _np.asarray(x.feats)
    out = _np.full((uniq.shape[0], feats.shape[1]), -_np.inf, dtype=feats.dtype)
    _np.maximum.at(out, inv, feats)

    new_coords = _np.stack(
        [
            uniq // off[0],
            (uniq // off[1]) % maxs[0],
            (uniq // off[2]) % maxs[1],
            uniq % maxs[2],
        ],
        axis=-1,
    ).astype(_np.int32)
    return SparseTensor(
        mx.array(out),
        mx.array(new_coords),
        scale=tuple(s * factor for s in x._scale),
    )