"""
Simplified selective SSM block for image tokens.
O(N) complexity, O(1) memory per token.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class SimplifiedMambaBlock(nn.Module):
    """Minimal selective SSM block without cuda-specific selective_scan."""
    def __init__(self, d_model: int, d_state: int = 16, d_conv: int = 4, expand: int = 2):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.d_inner = int(expand * d_model)
        self.dt_rank = math.ceil(d_model / 16)
        self.d_conv = d_conv
        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)
        self.conv1d = nn.Conv1d(
            self.d_inner,
            self.d_inner,
            kernel_size=d_conv,
            padding=d_conv - 1,
            groups=self.d_inner,
            bias=True,
        )
        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + d_state * 2, bias=False)
        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)
        A = torch.arange(1, d_state + 1, dtype=torch.float32).repeat(self.d_inner, 1)
        self.A_log = nn.Parameter(torch.log(A))
        self.D = nn.Parameter(torch.ones(self.d_inner))
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
        self.norm = nn.LayerNorm(d_model)

    def _selective_scan(self, x, dt, A, B, C, D):
        Bb, L, d_in = x.shape
        d_state = A.shape[1]
        dtA = dt.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0)
        A_bar = torch.exp(dtA)
        dtB = dt.unsqueeze(-1) * B.unsqueeze(2)
        h = torch.zeros(Bb, d_in, d_state, device=x.device, dtype=x.dtype)
        ys = []
        for t in range(L):
            h = A_bar[:, t] * h + dtB[:, t] * x[:, t].unsqueeze(-1)
            y = torch.sum(h * C[:, t].unsqueeze(1), dim=-1)
            ys.append(y)
        y = torch.stack(ys, dim=1)
        y = y + D.unsqueeze(0).unsqueeze(0) * x
        return y

    def forward(self, x: torch.Tensor):
        x_norm = self.norm(x)
        xz = self.in_proj(x_norm)
        x_gate, z_gate = xz.chunk(2, dim=-1)
        x_conv = self.conv1d(x_gate.transpose(1, 2))[:, :, :x_gate.shape[1]].transpose(1, 2)
        x_conv = F.silu(x_conv)
        xbc = self.x_proj(x_conv)
        dt_un, B_un, C_un = torch.split(xbc, [self.dt_rank, self.d_state, self.d_state], dim=-1)
        dt = F.softplus(self.dt_proj(dt_un))
        A = -torch.exp(self.A_log.float())
        B = B_un
        C = C_un
        y = self._selective_scan(x_conv, dt, A, B, C, self.D)
        y = y * F.silu(z_gate)
        out = self.out_proj(y)
        return out + x