"""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, mode: str = "max"
) -> "SparseTensor":
    """Downsample by `factor`, reducing colliding voxels with `mode` ("max" or "mean").

    The reduction differs BETWEEN MODELS and is numerically silent if you get it wrong —
    identical shapes, quietly different features:

      LATO.2   hardcodes reduce="amax" (its docstring says "average pooling"; the
               `reduce='mean'` line is commented out — follow the code, not the prose).
      Pixal3D  parameterises it as `mode: Literal['mean','max'] = 'mean'`, i.e. it
               DEFAULTS TO MEAN.

    So callers must be explicit. LATO.2 passes mode="max".

    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)
    if mode == "max":
        out = _np.full((uniq.shape[0], feats.shape[1]), -_np.inf, dtype=feats.dtype)
        _np.maximum.at(out, inv, feats)
    elif mode == "mean":
        out = _np.zeros((uniq.shape[0], feats.shape[1]), dtype=feats.dtype)
        _np.add.at(out, inv, feats)
        out /= _np.bincount(inv, minlength=uniq.shape[0])[:, None].astype(feats.dtype)
    else:
        raise ValueError(f"mode must be 'max' or 'mean', got {mode!r}")

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


def upsample(
    x: "SparseTensor", subdivision: "SparseTensor", factor: int = 2
) -> "SparseTensor":
    """Masked ×`factor` upsample — expand each voxel only into its FLAGGED children.

    Distinct from `subdivide`, which expands every voxel into all 8 children
    unconditionally. Here `subdivision.feats` is a [N, factor**3] 0/1 mask saying which
    children survive, so a voxel can produce anywhere from 0 to 8 outputs. Pixal3D's
    decoder uses this to grow the occupied set selectively rather than blowing up ×8 at
    every level.

    Child index decoding matches upstream: `subidx // factor**i % factor` for axis i,
    i.e. axis 0 is the fastest-varying — note this is the REVERSE of `subdivide`'s
    C-order, and mixing them up misplaces every child.

    Scale divides by `factor` (finer), where downsample multiplies.
    """
    import numpy as _np

    sub = _np.asarray(subdivision.feats)
    if sub.ndim != 2 or sub.shape[0] != x.coords.shape[0]:
        raise ValueError(
            f"subdivision must be [N, factor**3] matching N={x.coords.shape[0]}, "
            f"got {sub.shape}"
        )
    mask = sub > 0
    n_leaf = mask.sum(axis=1)
    parent = _np.repeat(_np.arange(sub.shape[0]), n_leaf)
    subidx = _np.nonzero(mask)[1]

    base = _np.asarray(x.coords, dtype=_np.int64).copy()
    base[:, 1:] *= factor
    new_coords = base[parent]
    for i in range(3):
        new_coords[:, i + 1] += subidx // (factor**i) % factor

    new_feats = mx.take(x.feats, mx.array(parent), axis=0)
    return SparseTensor(
        new_feats,
        mx.array(new_coords.astype(_np.int32)),
        scale=tuple(s / factor for s in x._scale),
    )


def _subidx_of(coords_np, factor: int):
    """Slot index of each voxel within its factor^3 cell.

    subidx = sum_i (coord[i] % factor) * factor**i  — axis 0 is the FASTEST varying,
    which is the reverse of C order. `upsample` and `channel2spatial` decode it back the
    same way (`subidx // factor**i % factor`); using C order in one place and this in the
    other misplaces every child while keeping all shapes valid.
    """
    import numpy as _np

    r = coords_np[:, 1:] % factor
    return sum(r[:, i] * (factor**i) for i in range(3)).astype(_np.int64)


def spatial2channel(x: "SparseTensor", factor: int = 2) -> "SparseTensor":
    """Sparse space-to-depth: [N, C] at res R -> [N', C*factor^3] at res R/factor.

    Voxels sharing a coarse cell are packed into disjoint channel slots; slots with no
    voxel are ZERO. Registers the inverse mapping in the spatial cache so a later
    `channel2spatial` can undo it without being handed a subdivision mask.
    """
    import numpy as _np

    c = _np.asarray(x.coords, dtype=_np.int64)
    feats = _np.asarray(x.feats)
    n, ch = feats.shape
    f3 = factor**3

    coarse = c.copy()
    coarse[:, 1:] //= factor
    maxs = coarse[:, 1:].max(axis=0) + 1
    off = _np.array(
        [maxs[0] * maxs[1] * maxs[2], maxs[1] * maxs[2], maxs[2], 1], dtype=_np.int64
    )
    code = (coarse * off).sum(axis=1)
    uniq, idx = _np.unique(code, return_inverse=True)
    subidx = _subidx_of(c, factor)

    out = _np.zeros((uniq.shape[0] * f3, ch), dtype=feats.dtype)
    out[idx * f3 + subidx] = 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)

    res = SparseTensor(
        mx.array(out.reshape(uniq.shape[0], ch * f3)),
        mx.array(new_coords),
        scale=tuple(s * factor for s in x._scale),
    )
    res.register_spatial_cache(f"channel2spatial_{factor}", (c.astype(_np.int32), idx, subidx))
    return res


def channel2spatial(
    x: "SparseTensor", subdivision: "SparseTensor" = None, factor: int = 2
) -> "SparseTensor":
    """Inverse of `spatial2channel`: [N, C*factor^3] -> [N', C] at finer resolution.

    Uses the cached mapping when this follows a `spatial2channel`; otherwise a
    `subdivision` mask ([N, factor^3], nonzero = keep) selects which children exist.
    Unlike `upsample`, features are SPLIT across children rather than replicated: child
    j takes channel slot j.
    """
    import numpy as _np

    cache = x.cache_get(f"channel2spatial_{factor}")
    f3 = factor**3
    feats = _np.asarray(x.feats)
    ch = feats.shape[1] // f3

    if cache is not None:
        new_coords, idx, subidx = cache
        new_coords = _np.asarray(new_coords, dtype=_np.int32)
    else:
        if subdivision is None:
            raise ValueError(
                "channel2spatial needs a subdivision mask, or to follow a spatial2channel"
            )
        sub = _np.asarray(subdivision.feats) > 0
        n_leaf = sub.sum(axis=1)
        idx = _np.repeat(_np.arange(sub.shape[0]), n_leaf)
        subidx = _np.nonzero(sub)[1]
        base = _np.asarray(x.coords, dtype=_np.int64).copy()
        base[:, 1:] *= factor
        new_coords = base[idx]
        for i in range(3):
            new_coords[:, i + 1] += subidx // (factor**i) % factor
        new_coords = new_coords.astype(_np.int32)

    flat = feats.reshape(feats.shape[0] * f3, ch)
    new_feats = flat[idx * f3 + subidx]
    return SparseTensor(
        mx.array(new_feats),
        mx.array(new_coords),
        scale=tuple(s / factor for s in x._scale),
    )