# examples/vmamba.py
"""CPU teaching reference for final VMamba's four scans and no-z VSS block.

This compact model uses RANDOM weights; it is not the pretrained classifier.
Run: python examples/vmamba.py --smoke --backward
Genuine released predictions use vmamba_export.py and the pinned official model.
Reference: MzeroMiko/VMamba at 2ed52ead062a51a64521ed3871d52914bf532876.
"""

import argparse
import json
import math

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


# region: cross-scan
def cross_scan(x):
    """Map [batch, channels, H, W] to four [batch, 4, channels, H*W] routes."""
    if x.ndim != 4 or min(x.shape) <= 0:
        raise ValueError("Expected a nonempty NCHW feature map")
    row = x.flatten(2)
    column = x.transpose(2, 3).flatten(2)
    forward = torch.stack((row, column), dim=1)
    return torch.cat((forward, forward.flip(-1)), dim=1)


# endregion: cross-scan


def _validate_scan(u, delta, a, b, c):
    """Reject ambiguous grouping instead of silently broadcasting channels."""
    if u.ndim != 3 or b.ndim != 4 or min(u.shape) <= 0 or min(b.shape) <= 0:
        raise ValueError(
            "Expected nonempty u [batch, channels, length] and B [batch, groups, state, length]"
        )
    batch, channels, length = u.shape
    if (
        delta.shape != u.shape
        or c.shape != b.shape
        or b.shape[0] != batch
        or b.shape[-1] != length
        or channels % b.shape[1]
        or a.shape != (channels, b.shape[2])
    ):
        raise ValueError("Incompatible scan shapes or channel grouping")


# region: recurrence
def selective_scan(
    u, delta, a, b, c, d=None, delta_bias=None, delta_softplus=True
):
    """Scan grouped B/C, using the official delta*B*u input discretization.

    u/delta: [batch, channels, length]; a: [channels, state].
    b/c: [batch, groups, state, length], shared within each channel group.
    d/delta_bias: [channels]. States start at zero for each channel/route.
    Computation is float32, matching the upstream torch reference.
    """
    _validate_scan(u, delta, a, b, c)
    if delta_bias is not None:
        delta = delta + delta_bias[None, :, None]
    if delta_softplus:
        delta = F.softplus(delta)
    u, delta, a, b, c = (v.float() for v in (u, delta, a, b, c))
    batch, channels, length = u.shape
    channels_per_group = channels // b.shape[1]
    b = b.repeat_interleave(channels_per_group, dim=1)
    c = c.repeat_interleave(channels_per_group, dim=1)
    state = u.new_zeros(batch, channels, a.shape[-1])
    outputs = []
    for t in range(length):
        step = delta[:, :, t, None]
        retention = torch.exp(step * a)
        injection = step * b[:, :, :, t] * u[:, :, t, None]
        state = retention * state + injection
        outputs.append((c[:, :, :, t] * state).sum(-1))
    y = torch.stack(outputs, dim=-1)
    return y if d is None else y + d[None, :, None] * u


# endregion: recurrence


# region: cross-merge
def cross_merge(routes, height, width):
    """Undo reversals and column transpose, then SUM aligned outputs."""
    if (
        routes.ndim != 4
        or routes.shape[1] != 4
        or height <= 0
        or width <= 0
        or routes.shape[-1] != height * width
    ):
        raise ValueError("Expected four routes with length height * width")
    aligned = routes[:, :2] + routes[:, 2:].flip(-1)
    batch, _, channels, _ = aligned.shape
    row = aligned[:, 0].reshape(batch, channels, height, width)
    column = aligned[:, 1].reshape(batch, channels, width, height)
    return row + column.transpose(2, 3)


# endregion: cross-merge


class ChannelNorm(nn.LayerNorm):
    """Layer normalization over channels at each spatial position."""

    def forward(self, x):
        return super().forward(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)


class SS2D(nn.Module):
    """Final no-z SS2D topology, with a readable sequential CPU scan."""

    def __init__(self, channels, state_size=1):
        super().__init__()
        self.channels = channels
        self.state_size = state_size
        self.rank = math.ceil(channels / 16)
        self.in_proj = nn.Conv2d(channels, channels, 1, bias=False)
        self.local = nn.Conv2d(
            channels, channels, 3, padding=1, groups=channels, bias=False
        )
        self.x_proj = nn.Parameter(
            torch.randn(4, self.rank + 2 * state_size, channels)
            / math.sqrt(channels)
        )
        self.dt_proj = nn.Parameter(
            torch.randn(4, channels, self.rank) / math.sqrt(self.rank)
        )
        self.dt_bias = nn.Parameter(torch.full((4 * channels,), -3.0))
        self.a_log = nn.Parameter(
            torch.arange(1, state_size + 1)
            .float()
            .log()
            .repeat(4 * channels, 1)
        )
        self.d = nn.Parameter(torch.ones(4 * channels))
        self.norm = ChannelNorm(channels)
        self.out_proj = nn.Conv2d(channels, channels, 1, bias=False)

    # region: selectivity
    def selective_parameters(self, routes):
        """Each direction projects input into its delta, B and C sequences."""
        projected = torch.einsum("bkdl,kpd->bkpl", routes, self.x_proj)
        low_rank_delta, b, c = projected.split(
            (self.rank, self.state_size, self.state_size), dim=2
        )
        raw_delta = torch.einsum("bkrl,kdr->bkdl", low_rank_delta, self.dt_proj)
        a = (
            -self.a_log.exp()
        )  # Learned stable negative rates, not input-derived.
        return raw_delta.flatten(1, 2), a, b, c

    # endregion: selectivity

    # region: ss2d
    def forward(self, x):
        """Mix locally, scan globally, align routes, normalize and project."""
        x = F.silu(self.local(self.in_proj(x)))
        batch, channels, height, width = x.shape
        routes = cross_scan(x)
        delta, a, b, c = self.selective_parameters(routes)
        outputs = selective_scan(
            routes.flatten(1, 2), delta, a, b, c, self.d, self.dt_bias
        )
        outputs = outputs.reshape(batch, 4, channels, height * width)
        merged = cross_merge(outputs, height, width)
        return self.out_proj(self.norm(merged))

    # endregion: ss2d


# region: vss
class VSSBlock(nn.Module):
    """Two pre-normalized residuals; the final block has no separate z gate.

    Drop-path is omitted in this teaching model (identity during evaluation).
    The scan's D*u skip is inside SS2D and separate from these two residuals.
    """

    def __init__(self, channels):
        super().__init__()
        self.norm1 = ChannelNorm(channels)
        self.ss2d = SS2D(channels)
        self.norm2 = ChannelNorm(channels)
        self.mlp = nn.Sequential(
            nn.Conv2d(channels, 4 * channels, 1),
            nn.GELU(),
            nn.Conv2d(4 * channels, channels, 1),
        )

    def forward(self, x):
        x = x + self.ss2d(self.norm1(x))
        return x + self.mlp(self.norm2(x))


# endregion: vss


# region: backbone
class VMamba(nn.Module):
    """Four-stage teaching backbone; defaults match Tiny stage dimensions.

    Initialization and state-dictionary names are teaching choices. Use the
    official exporter for pretrained results, never load released weights here.
    """

    def __init__(
        self, depths=(2, 2, 8, 2), dims=(96, 192, 384, 768), num_classes=1000
    ):
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(3, dims[0] // 2, 3, stride=2, padding=1),
            ChannelNorm(dims[0] // 2),
            nn.GELU(),
            nn.Conv2d(dims[0] // 2, dims[0], 3, stride=2, padding=1),
            ChannelNorm(dims[0]),
        )
        self.stages = nn.ModuleList(
            [
                nn.Sequential(*(VSSBlock(dim) for _ in range(depth)))
                for dim, depth in zip(dims, depths)
            ]
        )
        self.downsamples = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Conv2d(left, right, 3, stride=2, padding=1),
                    ChannelNorm(right),
                )
                for left, right in zip(dims[:-1], dims[1:])
            ]
        )
        self.norm = ChannelNorm(dims[-1])
        self.head = nn.Linear(dims[-1], num_classes)

    def forward(self, image):
        x = self.stem(image)
        for index, stage in enumerate(self.stages):
            x = stage(x)
            if index < len(self.downsamples):
                x = self.downsamples[index](x)
        return self.head(self.norm(x).mean(dim=(2, 3)))


# endregion: backbone


def main():
    """Run a small, reproducible CPU forward/backward without downloads."""
    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 = VMamba(depths=(1, 1, 1, 1), dims=(8, 16, 32, 64))
    logits = model(torch.randn(2, 3, 32, 48))
    if args.backward:
        logits.square().mean().backward()
        assert all(
            p.grad is not None and torch.isfinite(p.grad).all()
            for p in model.parameters()
        )
    print(
        json.dumps(
            {
                "weights": "random teaching weights; not pretrained",
                "device": "cpu",
                "inputShape": [2, 3, 32, 48],
                "logitsShape": list(logits.shape),
                "finite": bool(torch.isfinite(logits).all()),
                "backward": args.backward,
            }
        )
    )


if __name__ == "__main__":
    main()