# examples/vision_mamba.py
"""Small, differentiable CPU reference for Vision Mamba's bidirectional v2 block.

This random-weight teaching model is not the pretrained classifier. Run:
    python examples/vision_mamba.py --smoke --backward

Architecture: Zhu et al., Vision Mamba (ICML 2024), hustvl/Vim at
 dd0358ad1e42701f22afbefa0717cc8825cf9f45 (Apache-2.0).
Selective recurrence follows Mamba, copyright (c) 2023 Tri Dao and Albert Gu.
The transparent PyTorch implementation below favors readability over speed.
"""

import argparse
import json
import math

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


# region: patches
def patch_tokens(image, projection):
    """Project RGB windows and flatten H then W, with columns changing fastest."""
    grid = projection(image)  # [batch, width, rows, columns]
    return grid.flatten(2).transpose(1, 2)  # [batch, patches, width]


# endregion: patches


# region: tokens
def assemble_tokens(patches, cls_token, positions):
    """Insert the learned CLS at floor(patch_count / 2), then add positions."""
    midpoint = patches.shape[1] // 2
    cls = cls_token.expand(patches.shape[0], -1, -1)
    tokens = torch.cat(
        (patches[:, :midpoint], cls, patches[:, midpoint:]), dim=1
    )
    if tokens.shape[1:] != positions.shape[1:]:
        raise ValueError("Position embeddings must match this token grid")
    return tokens + positions, midpoint


# endregion: tokens


def causal_convolution(x, weight, bias):
    """Apply per-channel left-padded convolution followed by SiLU (B,D,L)."""
    padded = F.pad(x, (weight.shape[-1] - 1, 0))
    return F.silu(F.conv1d(padded, weight, bias, groups=x.shape[1]))


# region: recurrence
def selective_scan(
    u, delta, a, b, c, d, z=None, delta_bias=None, delta_softplus=True
):
    """Scan a real diagonal state; B/C vary per token and share channels.

    Args:
        u: Content, shape [batch, channels, length].
        delta: Step-size preactivation with the same shape as u.
        a: Negative learned dynamics, shape [channels, state].
        b: Token-dependent input write, shape [batch, state, length].
        c: Token-dependent output read, shape [batch, state, length].
        d: Learned direct skip, shape [channels].
        z: Optional signed SiLU gate, with the same shape as u.
        delta_bias: Optional learned bias, shape [channels].
        delta_softplus: Whether to make step sizes positive with softplus.

    Returns:
        A tensor with u's shape. This uses exp(delta*A) retention and the
        delta*B*u input approximation, not the exact ZOH input integral.
    """
    if delta_bias is not None:
        delta = delta + delta_bias[None, :, None]
    if delta_softplus:
        delta = F.softplus(delta)
    state = u.new_zeros(u.shape[0], u.shape[1], a.shape[1])
    outputs = []
    for index in range(u.shape[-1]):
        step = delta[:, :, index, None]
        retention = torch.exp(step * a[None])
        write = step * b[:, None, :, index] * u[:, :, index, None]
        state = retention * state + write
        read = (state * c[:, None, :, index]).sum(dim=-1)
        outputs.append(read + d[None] * u[:, :, index])
    output = torch.stack(outputs, dim=-1)
    return output if z is None else output * F.silu(z)


# endregion: recurrence


class Direction(nn.Module):
    """Independent local convolution and state-space parameters for one order."""

    def __init__(self, channels, state_size, delta_rank, kernel_size=4):
        super().__init__()
        self.state_size = state_size
        self.delta_rank = delta_rank
        self.conv = nn.Conv1d(channels, channels, kernel_size, groups=channels)
        self.x_proj = nn.Linear(
            channels, delta_rank + 2 * state_size, bias=False
        )
        self.dt_proj = nn.Linear(delta_rank, channels)
        self.a_log = nn.Parameter(
            torch.arange(1.0, state_size + 1).log().repeat(channels, 1)
        )
        self.d = nn.Parameter(torch.ones(channels))

    # region: selectivity
    def forward(self, content):
        """Build token-dependent delta/B/C after this direction's causal conv."""
        u = causal_convolution(content, self.conv.weight, self.conv.bias)
        projected = self.x_proj(u.transpose(1, 2))
        dt, b, c = projected.split(
            [self.delta_rank, self.state_size, self.state_size], dim=-1
        )
        delta = self.dt_proj(dt).transpose(1, 2)
        a = -self.a_log.exp()  # Learned, input-independent negative dynamics.
        return selective_scan(
            u, delta, a, b.transpose(1, 2), c.transpose(1, 2), self.d
        )

    # endregion: selectivity


# region: bidirectional
def align_average_gate(forward, reverse_order, gate):
    """Align the reverse scan, average both directions and apply signed SiLU."""
    aligned = reverse_order.flip(-1)
    return (forward + aligned) * 0.5 * F.silu(gate)


class BidirectionalMixer(nn.Module):
    """Share the content/gate projection; keep the two directional SSMs separate."""

    def __init__(self, width, state_size=16):
        super().__init__()
        channels = 2 * width
        rank = math.ceil(width / 16)
        self.in_proj = nn.Linear(width, 2 * channels, bias=False)
        self.forward_direction = Direction(channels, state_size, rank)
        self.backward_direction = Direction(channels, state_size, rank)
        self.out_proj = nn.Linear(channels, width, bias=False)

    def forward(self, tokens):
        content, gate = self.in_proj(tokens).transpose(1, 2).chunk(2, dim=1)
        forward = self.forward_direction(content)
        backward = self.backward_direction(content.flip(-1))
        mixed = align_average_gate(forward, backward, gate)
        # Official code gates each direction before averaging. The shared
        # aligned gate makes this after-average form algebraically equivalent.
        return self.out_proj(mixed.transpose(1, 2))


# endregion: bidirectional


# region: block
class RMSNorm(nn.Module):
    """Normalize the root mean square across channels, without mean removal."""

    def __init__(self, width, eps=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(width))
        self.eps = eps

    def forward(self, x):
        return (
            x
            * torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + self.eps)
            * self.weight
        )


class VimBlock(nn.Module):
    """Carry residual separately, matching the released add/norm/mixer order."""

    def __init__(self, width, state_size):
        super().__init__()
        self.norm = RMSNorm(width)
        self.mixer = BidirectionalMixer(width, state_size)

    def forward(self, hidden, residual=None):
        residual = hidden if residual is None else residual + hidden
        contribution = self.mixer(self.norm(residual))
        return contribution, residual


# endregion: block


class VisionMamba(nn.Module):
    """Small configurable teaching backbone; no pretrained weights or dropout."""

    def __init__(
        self,
        image_size=(8, 12),
        patch_size=4,
        width=16,
        depth=2,
        state_size=4,
        classes=5,
    ):
        super().__init__()
        rows, columns = (side // patch_size for side in image_size)
        self.projection = nn.Conv2d(3, width, patch_size, stride=patch_size)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, width))
        self.positions = nn.Parameter(torch.zeros(1, rows * columns + 1, width))
        self.blocks = nn.ModuleList(
            VimBlock(width, state_size) for _ in range(depth)
        )
        self.norm = RMSNorm(width)
        self.head = nn.Linear(width, classes)

    # region: readout
    def forward(self, image):
        patches = patch_tokens(image, self.projection)
        hidden, cls_index = assemble_tokens(
            patches, self.cls_token, self.positions
        )
        residual = None
        for block in self.blocks:
            hidden, residual = block(hidden, residual)
        # The last mixer contribution has not yet entered the carried residual.
        features = self.norm(hidden if residual is None else residual + hidden)
        cls_feature = features[:, cls_index]
        return self.head(cls_feature)  # CLS takes precedence over mean pooling.

    # endregion: readout


# region: long-sequence
def sequence_length(image_size, patch_size, stride):
    """Count actual overlapping windows, plus CLS; this does not resize weights."""
    windows_per_side = (image_size - patch_size) // stride + 1
    return windows_per_side**2 + 1


# 224px, kernel 16: stride 16 -> 197 tokens; stride 8 -> 730 tokens.
# Stride 8 needs separately fine-tuned weights and position embeddings.
# endregion: long-sequence


def main():
    """Run a finite random-weight CPU illustration, optionally with gradients."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--backward", action="store_true")
    args = parser.parse_args()
    torch.manual_seed(0)
    torch.set_num_threads(2)
    model = VisionMamba()
    image = torch.randn(2, 3, 8, 12, requires_grad=args.backward)
    logits = model(image)
    if not torch.isfinite(logits).all():
        raise RuntimeError("Nonfinite forward")
    if args.backward:
        logits.square().mean().backward()
        if not all(
            p.grad is not None and torch.isfinite(p.grad).all()
            for p in model.parameters()
        ):
            raise RuntimeError("Nonfinite or absent gradient")
    print(
        json.dumps(
            {
                "kind": "random-weight teaching smoke",
                "device": "cpu",
                "logitsShape": list(logits.shape),
                "finiteBackward": args.backward,
            }
        )
    )


if __name__ == "__main__":
    main()