# examples/mamba.py
# Copyright (c) 2023, Tri Dao, Albert Gu.
# SPDX-License-Identifier: Apache-2.0
"""Transparent, differentiable Mamba CPU reference with persistent decode buffers.

Adapted from state-spaces/mamba v2.2.4, commit
95d8aba8a8c75aedcaa6143713b11e745e7cd0d9, mamba_simple.py and
selective_scan_interface.py. Licensed under Apache-2.0:
https://www.apache.org/licenses/LICENSE-2.0
This educational adaptation uses no optional CUDA packages. It implements the
real diagonal recurrence, not the fused hardware-aware kernel or a language model.
The outside prenorm/residual wrapper is separate from MambaBlock (the mixer).
Run: python examples/mamba.py --smoke --backward
Check: python examples/mamba.py --self-check
"""

import argparse
import json
import math

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


def _finite_shape(tensor, shape, name):
    """Reject accidental broadcasting and nonfinite inputs."""
    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_scan(x, delta, a, b, c, initial_state):
    if x.ndim != 3 or min(x.shape) <= 0 or a.ndim != 2 or min(a.shape) <= 0:
        raise ValueError(
            "Expected nonempty x [batch,length,channels], A [channels,state]"
        )
    batch, length, channels = x.shape
    size = a.shape[-1]
    _finite_shape(x, (batch, length, channels), "x")
    _finite_shape(delta, x.shape, "delta")
    _finite_shape(a, (channels, size), "A")
    _finite_shape(b, (batch, length, size), "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, channels, size), "initial_state")


# region: recurrence
def selective_scan(x, delta, a, b, c, d=None, initial_state=None):
    """Return outputs and all states of a diagonal selective recurrence.

    Args:
        x: Content [batch, length, channels].
        delta: Already positive step sizes, same shape as x.
        a: Negative continuous dynamics [channels, state].
        b: Input-dependent write coefficients [batch, length, state].
        c: Input-dependent read coefficients, same shape as b.
        d: Optional direct skip [channels], separate from the outside residual.
        initial_state: Optional state [batch, channels, state].

    Returns:
        Outputs [batch,length,channels] and states [batch,length,channels,state].
        Input injection is delta*B*x, not the exact zero-order-hold integral.
    """
    _validate_scan(x, delta, a, b, c, initial_state)
    if d is not None:
        _finite_shape(d, (x.shape[-1],), "D")
    state = (
        x.new_zeros(x.shape[0], x.shape[2], a.shape[1])
        if initial_state is None
        else initial_state
    )
    outputs, states = [], []
    for index in range(x.shape[1]):
        step = delta[:, index, :, None]
        retention = torch.exp(step * a)
        write = step * b[:, index, None, :] * x[:, index, :, None]
        state = retention * state + write
        readout = (state * c[:, index, None, :]).sum(-1)
        outputs.append(readout if d is None else readout + d * x[:, index])
        states.append(state)
    return torch.stack(outputs, 1), torch.stack(states, 1)


# endregion: recurrence


# region: scan
def affine_prefix(x, delta, a, b, initial_state=None):
    """Hillis-Steele scan; each round reads a snapshot, not its own writes."""
    _validate_scan(x, delta, a, b, b, initial_state)
    retention = torch.exp(delta[..., None] * a)
    write = delta[..., None] * b[:, :, None, :] * x[..., None]
    pairs = [(retention, write)]
    stride = 1
    while stride < x.shape[1]:
        previous_a, previous_b = pairs[-1]
        # later(earlier(h)) = later.a*earlier.a*h + later.a*earlier.b+later.b
        combined_a = previous_a[:, stride:] * previous_a[:, :-stride]
        combined_b = (
            previous_a[:, stride:] * previous_b[:, :-stride]
            + previous_b[:, stride:]
        )
        pairs.append(
            (
                torch.cat((previous_a[:, :stride], combined_a), 1),
                torch.cat((previous_b[:, :stride], combined_b), 1),
            )
        )
        stride *= 2
    final_a, final_b = pairs[-1]
    initial = (
        x.new_zeros(x.shape[0], x.shape[2], a.shape[1])
        if initial_state is None
        else initial_state
    )
    return final_a * initial[:, None] + final_b, pairs


# endregion: scan


# region: convolution
def causal_convolution(x, convolution):
    """Convolve [batch,length,channels] with left context only, then SiLU."""
    width = convolution.kernel_size[0]
    padded = F.pad(x.transpose(1, 2), (width - 1, 0))
    return F.silu(
        F.conv1d(
            padded, convolution.weight, convolution.bias, groups=x.shape[-1]
        )
    ).transpose(1, 2)


# endregion: convolution


def rms_norm(x, weight=None, eps=1e-5):
    """RMS normalization; callers own the outside prenorm/residual wrapper."""
    normalized = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + eps)
    return normalized if weight is None else normalized * weight


class MambaBlock(nn.Module):
    """Mamba mixer, with neither the outside prenorm nor outside residual."""

    def __init__(self, model_width=4, state_size=3, expand=2, conv_width=3):
        super().__init__()
        if any(
            not isinstance(value, int) or value <= 0
            for value in (model_width, state_size, expand, conv_width)
        ):
            raise ValueError("Block dimensions must be positive integers")
        self.model_width = model_width
        self.channels = model_width * expand
        self.state_size = state_size
        self.conv_width = conv_width
        self.rank = math.ceil(model_width / 16)
        self.in_proj = nn.Linear(model_width, 2 * self.channels, bias=False)
        self.conv = nn.Conv1d(
            self.channels, self.channels, conv_width, groups=self.channels
        )
        self.x_proj = nn.Linear(
            self.channels, self.rank + 2 * state_size, bias=False
        )
        self.dt_proj = nn.Linear(self.rank, self.channels)
        self.A_log = nn.Parameter(
            torch.arange(1, state_size + 1)
            .float()
            .log()
            .repeat(self.channels, 1)
        )
        self.D = nn.Parameter(torch.ones(self.channels))
        self.out_proj = nn.Linear(self.channels, model_width, bias=False)
        nn.init.constant_(self.dt_proj.bias, -2.0)

    # region: selection
    def coefficients(self, x):
        """Produce token-dependent delta/B/C after the local convolution."""
        low_rank, b, c = self.x_proj(x).split(
            (self.rank, self.state_size, self.state_size), -1
        )
        delta = F.softplus(self.dt_proj(low_rank))
        a = -self.A_log.exp()  # Learned, input-independent continuous dynamics.
        return delta, a, b, c

    # endregion: selection

    # region: block
    def forward(self, u, return_trace=False):
        """Trace the mixer; an enclosing block supplies prenorm and residual."""
        projected = self.in_proj(u)
        x, z = projected.chunk(2, -1)
        convolved = causal_convolution(x, self.conv)
        delta, a, b, c = self.coefficients(convolved)
        scanned, states = selective_scan(convolved, delta, a, b, c, self.D)
        gated = scanned * F.silu(z)
        output = self.out_proj(gated)
        if return_trace:
            return output, {
                "projected": projected,
                "x": x,
                "z": z,
                "convolved": convolved,
                "delta": delta,
                "A": a,
                "B": b,
                "C": c,
                "states": states,
                "scanned": scanned,
                "gated": gated,
                "output": output,
            }
        return output

    # endregion: block

    # region: cache
    def initial_cache(self, batch):
        """Store K projected samples, including the latest, as the source does."""
        reference = self.in_proj.weight
        return (
            reference.new_zeros(batch, self.channels, self.conv_width),
            reference.new_zeros(batch, self.channels, self.state_size),
        )

    def step(self, token, cache):
        """Process one [batch,model_width] token with functional cache updates."""
        conv_state, ssm_state = cache
        x, z = self.in_proj(token).chunk(2, -1)
        conv_state = torch.cat((conv_state[..., 1:], x[..., None]), -1)
        x = F.silu(
            (conv_state * self.conv.weight[:, 0]).sum(-1) + self.conv.bias
        )
        delta, a, b, c = self.coefficients(x)
        ssm_state = (
            torch.exp(delta[..., None] * a) * ssm_state
            + delta[..., None] * b[:, None, :] * x[..., None]
        )
        readout = (ssm_state * c[:, None, :]).sum(-1) + self.D * x
        output = self.out_proj(readout * F.silu(z))
        return output, (conv_state, ssm_state)

    # endregion: cache


def self_check():
    """Check analytic values, scan order, gradients, causality and both caches."""
    torch.manual_seed(17)
    for dtype, tolerance in ((torch.float64, 1e-9), (torch.float32, 1e-5)):
        for batch, channels, size in ((1, 2, 3), (2, 4, 2)):
            x = torch.randn(batch, 6, channels, dtype=dtype, requires_grad=True)
            delta = torch.rand_like(x) + 0.1
            a = -torch.rand(channels, size, dtype=dtype) - 0.1
            b = torch.randn(batch, 6, size, dtype=dtype)
            c = torch.randn_like(b)
            initial = torch.randn(batch, channels, size, dtype=dtype)
            y, states = selective_scan(x, delta, a, b, c, initial_state=initial)
            prefix, _ = affine_prefix(x, delta, a, b, initial)
            torch.testing.assert_close(
                prefix, states, atol=tolerance, rtol=tolerance
            )
            y.square().mean().backward()
            assert x.grad is not None and torch.isfinite(x.grad).all()
        model = MambaBlock().to(dtype)
        u = torch.randn(2, 6, 4, dtype=dtype, requires_grad=True)
        full, trace = model(u, return_trace=True)
        cache = model.initial_cache(2)
        streamed = []
        for token in u.unbind(1):
            output, cache = model.step(token, cache)
            streamed.append(output)
        torch.testing.assert_close(
            torch.stack(streamed, 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["x"][:, -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(
            parameter.grad is not None and torch.isfinite(parameter.grad).all()
            for parameter in model.parameters()
        )
    print(
        "Mamba self-check passed: scan, causality, full/stream caches, finite gradients"
    )


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 = MambaBlock()
    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 random-weight mechanics",
                "shape": list(output.shape),
                "finite": True,
                "backward": args.backward,
                "device": "cpu",
            }
        )
    )


if __name__ == "__main__":
    main()