# examples/mamba2.py
# Copyright (c) 2024, Tri Dao, Albert Gu.
# SPDX-License-Identifier: Apache-2.0
"""Self-contained Mamba-2 CPU reference: recurrence, dense SSD and local chunks.

Adapted from state-spaces/mamba v2.2.4, commit
95d8aba8a8c75aedcaa6143713b11e745e7cd0d9, mamba2_simple.py and
ssd_minimal.py. The latter credits Albert Gu and Tri Dao (2024).
Licensed under Apache-2.0: https://www.apache.org/licenses/LICENSE-2.0
The dense reference is quadratic. Chunk boundaries here pass by a linear
recurrence rather than the minimal listing's dense inter-chunk matrix.
Uses one B/C group shared by heads in the block; the core accepts per-head B/C.
This is an untrained mechanics reference, not a language model or speed benchmark.
Run: python examples/mamba2.py --smoke --backward
Check: python examples/mamba2.py --self-check
"""

import argparse
import json

import torch
from torch import nn
from torch.nn import functional as F


def _finite_shape(tensor, shape, name):
    if tuple(tensor.shape) != tuple(shape) or not torch.isfinite(tensor).all():
        raise ValueError(f"{name} must have finite values and shape {shape}")


def _validate_ssd(x, delta, a, b, c, initial_state):
    if x.ndim != 4 or min(x.shape) <= 0 or b.ndim != 4 or b.shape[-1] <= 0:
        raise ValueError(
            "Expected nonempty x [batch,length,heads,width], B [batch,length,heads,state]"
        )
    batch, length, heads, width = x.shape
    _finite_shape(x, x.shape, "x")
    _finite_shape(delta, (batch, length, heads), "delta")
    _finite_shape(a, (heads,), "A")
    _finite_shape(b, (batch, length, heads, b.shape[-1]), "B")
    _finite_shape(c, b.shape, "C")
    if not (delta > 0).all() or not (a < 0).all():
        raise ValueError("delta must be positive and continuous A negative")
    if initial_state is not None:
        _finite_shape(
            initial_state, (batch, heads, width, b.shape[-1]), "initial_state"
        )


# region: recurrence
def recurrent_ssd(x, delta, a, b, c, initial_state=None):
    """Return y [B,T,H,P], states [B,T,H,P,N] from continuous scalar A [H].

    delta is [B,T,H]; B/C are [B,T,H,N]. Scale the input by delta once.
    The pre-sequence state is [B,H,P,N]; retention at index zero acts on it.
    """
    _validate_ssd(x, delta, a, b, c, initial_state)
    state = (
        x.new_zeros(x.shape[0], x.shape[2], x.shape[3], b.shape[-1])
        if initial_state is None
        else initial_state
    )
    outputs, states = [], []
    for index in range(x.shape[1]):
        retention = torch.exp(delta[:, index] * a)
        write = (
            delta[:, index, :, None, None]
            * x[:, index, :, :, None]
            * b[:, index, :, None, :]
        )
        state = retention[..., None, None] * state + write
        outputs.append(torch.einsum("bhpn,bhn->bhp", state, c[:, index]))
        states.append(state)
    return torch.stack(outputs, 1), torch.stack(states, 1)


# endregion: recurrence


# region: matrix
def segment_sum(log_decay):
    """Stable sums over j+1..i without subtracting long cumulative sums.

    Input shape [...,T]; output [...,T,T]. The inclusive diagonal has log 0;
    the upper triangle has -inf, so its exponential is exactly zero.
    """
    length = log_decay.shape[-1]
    expanded = log_decay[..., :, None].expand(*log_decay.shape, length)
    strict_lower = torch.ones(
        length, length, dtype=torch.bool, device=log_decay.device
    ).tril(-1)
    lower = strict_lower | torch.eye(
        length, dtype=torch.bool, device=log_decay.device
    )
    sums = expanded.masked_fill(~strict_lower, 0).cumsum(-2)
    return sums.masked_fill(~lower, -torch.inf)


def dense_ssd(x, delta, a, b, c, initial_state=None):
    """Materialize a quadratic teaching matrix; return y, final, M, initial_y.

    M[i,j] = delta[j] * (C[i] dot B[j]) * exp(sum(delta[k]*A, k=j+1..i)).
    This is signed causal mixing, not softmax-normalized attention.
    """
    _validate_ssd(x, delta, a, b, c, initial_state)
    log_decay = (delta * a).transpose(1, 2)
    decay = segment_sum(log_decay).exp()
    mixing = torch.einsum("bthn,bshn,bhts,bsh->bhts", c, b, decay, delta)
    state = (
        x.new_zeros(x.shape[0], x.shape[2], x.shape[3], b.shape[-1])
        if initial_state is None
        else initial_state
    )
    initial_outputs = torch.einsum(
        "bhpn,bthn,bht->bthp", state, c, log_decay.cumsum(-1).exp()
    )
    output = torch.einsum("bhts,bshp->bthp", mixing, x) + initial_outputs
    final = torch.einsum(
        "bht,bthn,bthp,bth->bhpn", decay[:, :, -1], b, x, delta
    )
    final = final + log_decay.sum(-1).exp()[..., None, None] * state
    return output, final, mixing, initial_outputs


# endregion: matrix


# region: chunks
def chunked_ssd(x, delta, a, b, c, chunk_size, initial_state=None):
    """Dense local blocks plus linear boundary passing; return y, final, chunks."""
    _validate_ssd(x, delta, a, b, c, initial_state)
    if (
        isinstance(chunk_size, bool)
        or not isinstance(chunk_size, int)
        or chunk_size <= 0
        or x.shape[1] % chunk_size
    ):
        raise ValueError(
            "chunk_size must be a positive integer dividing sequence length"
        )
    incoming = (
        x.new_zeros(x.shape[0], x.shape[2], x.shape[3], b.shape[-1])
        if initial_state is None
        else initial_state
    )
    outputs, chunks = [], []
    for start in range(0, x.shape[1], chunk_size):
        end = start + chunk_size
        local, local_final, _, _ = dense_ssd(
            x[:, start:end],
            delta[:, start:end],
            a,
            b[:, start:end],
            c[:, start:end],
        )
        log_decay = delta[:, start:end] * a
        carried = torch.einsum(
            "bhpn,bthn,bth->bthp",
            incoming,
            c[:, start:end],
            log_decay.cumsum(1).exp(),
        )
        final = log_decay.sum(1).exp()[..., None, None] * incoming + local_final
        output = local + carried
        chunks.append(
            {
                "start": start,
                "end": end,
                "localOutputs": local,
                "localFinal": local_final,
                "initialState": incoming,
                "finalState": final,
                "carriedOutputs": carried,
                "outputs": output,
            }
        )
        outputs.append(output)
        incoming = final
    return torch.cat(outputs, 1), incoming, chunks


# endregion: chunks


def rms_norm(x, weight=None, eps=1e-5):
    normalized = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + eps)
    return normalized if weight is None else normalized * weight


# region: heads
def split_heads(x, head_width):
    """Split expanded width D into H heads of P channels: [B,T,D] -> [B,T,H,P]."""
    if (
        not isinstance(head_width, int)
        or head_width <= 0
        or x.shape[-1] % head_width
    ):
        raise ValueError("head_width must divide expanded width")
    return x.reshape(*x.shape[:-1], x.shape[-1] // head_width, head_width)


# endregion: heads


class Mamba2Block(nn.Module):
    """One-group Mamba-2 mixer; the outside prenorm/residual stays separate."""

    def __init__(
        self, model_width=4, state_size=3, expand=2, conv_width=3, head_width=2
    ):
        super().__init__()
        if any(
            not isinstance(value, int) or value <= 0
            for value in (
                model_width,
                state_size,
                expand,
                conv_width,
                head_width,
            )
        ):
            raise ValueError("Block dimensions must be positive integers")
        self.model_width = model_width
        self.channels = model_width * expand
        if self.channels % head_width:
            raise ValueError("head_width must divide expanded width")
        self.head_width = head_width
        self.heads = self.channels // head_width
        self.state_size = state_size
        self.conv_width = conv_width
        self.conv_channels = self.channels + 2 * state_size
        self.in_proj = nn.Linear(
            model_width,
            self.channels + self.conv_channels + self.heads,
            bias=False,
        )
        self.conv = nn.Conv1d(
            self.conv_channels,
            self.conv_channels,
            conv_width,
            groups=self.conv_channels,
        )
        self.dt_bias = nn.Parameter(torch.full((self.heads,), -2.0))
        self.A_log = nn.Parameter(torch.arange(1, self.heads + 1).float().log())
        self.D = nn.Parameter(torch.ones(self.heads))
        self.norm_weight = nn.Parameter(torch.ones(self.channels))
        self.out_proj = nn.Linear(self.channels, model_width, bias=False)

    def _split_bc(self, xbc):
        x, b, c = xbc.split(
            (self.channels, self.state_size, self.state_size), -1
        )
        # One group: B and C are shared across all heads, not across time.
        b = b[..., None, :].expand(*b.shape[:-1], self.heads, self.state_size)
        c = c[..., None, :].expand_as(b)
        return x, b, c

    # region: block
    def forward(self, u, return_trace=False):
        """Follow pinned source order: parallel projection, conv, SSD, skip, norm."""
        projected = self.in_proj(u)  # Order [z, x, B, C, dt].
        z, xbc, dt = projected.split(
            (self.channels, self.conv_channels, self.heads), -1
        )
        delta = F.softplus(dt + self.dt_bias)
        padded = F.pad(xbc.transpose(1, 2), (self.conv_width - 1, 0))
        convolved = F.silu(
            F.conv1d(
                padded,
                self.conv.weight,
                self.conv.bias,
                groups=self.conv_channels,
            )
        ).transpose(1, 2)
        x, b, c = self._split_bc(convolved)
        headed = split_heads(x, self.head_width)
        mixed, states = recurrent_ssd(headed, delta, -self.A_log.exp(), b, c)
        skipped = mixed + self.D[None, None, :, None] * headed
        # norm_before_gate=False: multiply by SiLU(z) BEFORE RMS normalization.
        gated = skipped.flatten(-2) * F.silu(z)
        normalized = rms_norm(gated, self.norm_weight)
        output = self.out_proj(normalized)
        if return_trace:
            return output, {
                "projected": projected,
                "z": z,
                "xbc": xbc,
                "delta": delta,
                "A": -self.A_log.exp(),
                "convolved": convolved,
                "headed": headed,
                "B": b,
                "C": c,
                "states": states,
                "mixed": mixed,
                "skipped": skipped,
                "gated": gated,
                "normalized": normalized,
                "output": output,
            }
        return output

    # endregion: block

    # region: cache
    def initial_cache(self, batch):
        """Persistent buffers: [B,D+2N,K] and [B,H,P,N], with K stored samples."""
        reference = self.in_proj.weight
        return (
            reference.new_zeros(batch, self.conv_channels, self.conv_width),
            reference.new_zeros(
                batch, self.heads, self.head_width, self.state_size
            ),
        )

    def step(self, token, cache):
        """Stream one token with the same convolution, scalar-A SSD and gate norm."""
        conv_state, state = cache
        z, xbc, dt = self.in_proj(token).split(
            (self.channels, self.conv_channels, self.heads), -1
        )
        delta = F.softplus(dt + self.dt_bias)
        conv_state = torch.cat((conv_state[..., 1:], xbc[..., None]), -1)
        convolved = F.silu(
            (conv_state * self.conv.weight[:, 0]).sum(-1) + self.conv.bias
        )
        x, b, c = self._split_bc(convolved)
        x = split_heads(x, self.head_width)
        state = (
            torch.exp(-delta * self.A_log.exp())[..., None, None] * state
            + delta[..., None, None] * x[..., None] * b[..., None, :]
        )
        mixed = (state * c[..., None, :]).sum(-1) + self.D[None, :, None] * x
        normalized = rms_norm(mixed.flatten(-2) * F.silu(z), self.norm_weight)
        return self.out_proj(normalized), (conv_state, state)

    # endregion: cache


def self_check():
    """Compare outputs, final states, gradients, causality and decode caches."""
    torch.manual_seed(23)
    for dtype, tolerance in ((torch.float64, 1e-9), (torch.float32, 1e-5)):
        for batch, heads, width, size in ((1, 1, 2, 3), (2, 3, 2, 4)):
            x = torch.randn(
                batch, 6, heads, width, dtype=dtype, requires_grad=True
            )
            delta = (
                torch.rand(batch, 6, heads, dtype=dtype) + 0.1
            ).requires_grad_()
            a = (-torch.rand(heads, dtype=dtype) - 0.2).requires_grad_()
            b = torch.randn(
                batch, 6, heads, size, dtype=dtype, requires_grad=True
            )
            c = torch.randn_like(b, requires_grad=True)
            initial = torch.randn(
                batch, heads, width, size, dtype=dtype, requires_grad=True
            )
            args = (x, delta, a, b, c)
            recurrent, states = recurrent_ssd(*args, initial)
            final = states[:, -1]
            loss = recurrent.square().mean() + final.square().mean()
            reference_grads = torch.autograd.grad(
                loss, (*args, initial), retain_graph=True
            )
            dense, dense_final, _, _ = dense_ssd(*args, initial)
            variants = [(dense, dense_final)]
            variants.extend(
                chunked_ssd(*args, q, initial)[:2] for q in (1, 2, 3, 6)
            )
            for y, end_state in variants:
                torch.testing.assert_close(
                    y, recurrent, atol=tolerance, rtol=tolerance
                )
                torch.testing.assert_close(
                    end_state, final, atol=tolerance, rtol=tolerance
                )
                grads = torch.autograd.grad(
                    y.square().mean() + end_state.square().mean(),
                    (*args, initial),
                    retain_graph=True,
                )
                for actual, expected in zip(grads, reference_grads):
                    assert torch.isfinite(actual).all()
                    torch.testing.assert_close(
                        actual, expected, atol=tolerance, rtol=tolerance
                    )
            changed = x.detach().clone()
            changed[:, 4:] += 30
            causal, _ = recurrent_ssd(changed, delta, a, b, c, initial)
            torch.testing.assert_close(
                causal[:, :4], recurrent[:, :4], atol=tolerance, rtol=tolerance
            )
            for invalid in (0, -1, 4, 1.5):
                try:
                    chunked_ssd(*args, invalid, initial)
                    raise AssertionError("Invalid chunk length accepted")
                except ValueError:
                    pass
        model = Mamba2Block().to(dtype)
        u = torch.randn(2, 6, 4, dtype=dtype, requires_grad=True)
        full, trace = model(u, return_trace=True)
        cache, outputs = model.initial_cache(2), []
        for token in u.unbind(1):
            output, cache = model.step(token, cache)
            outputs.append(output)
        torch.testing.assert_close(
            torch.stack(outputs, 1), full, atol=tolerance, rtol=tolerance
        )
        torch.testing.assert_close(
            cache[1], trace["states"][:, -1], atol=tolerance, rtol=tolerance
        )
        torch.testing.assert_close(
            cache[0],
            trace["xbc"][:, -3:].transpose(1, 2),
            atol=tolerance,
            rtol=tolerance,
        )
        changed = u.detach().clone()
        changed[:, 4:] += 20
        torch.testing.assert_close(
            model(changed)[:, :4], full[:, :4], atol=tolerance, rtol=tolerance
        )
        full.square().mean().backward()
        assert torch.isfinite(u.grad).all()
        assert all(
            p.grad is not None and torch.isfinite(p.grad).all()
            for p in model.parameters()
        )
    print(
        "Mamba-2 self-check passed: recurrence/dense/chunks, nonzero initial, finite gradients, full/stream caches"
    )


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--backward", action="store_true")
    parser.add_argument("--self-check", action="store_true")
    args = parser.parse_args()
    if args.self_check:
        self_check()
        return
    torch.manual_seed(7)
    model = Mamba2Block()
    u = torch.randn(2, 6, 4, requires_grad=args.backward)
    output = u + model(
        rms_norm(u)
    )  # Explicit outside prenorm/residual wrapper.
    assert torch.isfinite(output).all()
    if args.backward:
        output.square().mean().backward()
        assert torch.isfinite(u.grad).all()
        assert all(
            p.grad is not None and torch.isfinite(p.grad).all()
            for p in model.parameters()
        )
    print(
        json.dumps(
            {
                "example": "Mamba-2 random-weight mechanics",
                "shape": list(output.shape),
                "finite": True,
                "backward": args.backward,
                "device": "cpu",
            }
        )
    )


if __name__ == "__main__":
    main()