"""Original selective-SSM teaching block based on Mamba's paper equations.

Width 4, expansion 8, diagonal state size 3, causal depthwise convolution width 3.
Uses exact diagonal zero-order-hold B discretization from paper Eq. 4, not a
fused CUDA kernel or the reference code's simplified delta*B input term.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


class Selection(nn.Module):
    def __init__(self):
        super().__init__()
        self.delta_rank = nn.Linear(8, 1, bias=False)
        self.delta = nn.Linear(1, 8)
        self.input_map = nn.Linear(8, 3, bias=False)
        self.readout = nn.Linear(8, 3, bias=False)

    def forward(self, u):
        delta = F.softplus(self.delta(self.delta_rank(u)))
        return delta, self.input_map(u), self.readout(u)


class SelectiveScan(nn.Module):
    def __init__(self):
        super().__init__()
        self.A_log = nn.Parameter(
            torch.arange(1, 4, dtype=torch.float32).log().repeat(8, 1)
        )
        self.direct = nn.Parameter(torch.ones(8))

    def forward(self, u, delta, b, c):
        a = -self.A_log.exp()
        state = torch.zeros(u.shape[0], 8, 3, device=u.device, dtype=u.dtype)
        outputs = []
        for t in range(u.shape[1]):
            dt = delta[:, t, :, None]
            decay = torch.exp(dt * a)
            drive = torch.expm1(dt * a) / a * b[:, t, None, :] * u[:, t, :, None]
            state = decay * state + drive
            outputs.append((state * c[:, t, None, :]).sum(-1) + self.direct * u[:, t])
        return torch.stack(outputs, dim=1)


class TinyMamba(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = nn.LayerNorm(4)
        self.in_projection = nn.Linear(4, 16, bias=False)
        self.local_conv = nn.Conv1d(8, 8, 3, groups=8)
        self.selection = Selection()
        self.scan = SelectiveScan()
        self.project = nn.Linear(8, 4, bias=False)

    def forward(self, x):
        raw, gate = self.in_projection(self.norm(x)).chunk(2, -1)
        u = F.silu(self.local_conv(F.pad(raw.transpose(1, 2), (2, 0))).transpose(1, 2))
        delta, b, c = self.selection(u)
        mixed = self.scan(u, delta, b, c)
        return x + self.project(mixed * F.silu(gate))


def scan_trace(scan, u, delta, b, c):
    state = torch.zeros(u.shape[0], 8, 3, dtype=u.dtype, device=u.device)
    states = []
    decays = []
    drives = []
    outputs = []
    a = -scan.A_log.exp()
    for t in range(u.shape[1]):
        decay = (delta[:, t, :, None] * a).exp()
        drive = (
            torch.expm1(delta[:, t, :, None] * a)
            / a
            * b[:, t, None, :]
            * u[:, t, :, None]
        )
        state = decay * state + drive
        states.append(state)
        decays.append(decay)
        drives.append(drive)
        outputs.append((state * c[:, t, None, :]).sum(-1) + scan.direct * u[:, t])
    return (
        torch.stack(outputs, 1),
        torch.stack(states, 1),
        torch.stack(decays, 1),
        torch.stack(drives, 1),
    )


def intermediates(model, x):
    raw, gate = model.in_projection(model.norm(x)).chunk(2, -1)
    u = F.silu(model.local_conv(F.pad(raw.transpose(1, 2), (2, 0))).transpose(1, 2))
    delta, b, c = model.selection(u)
    return u, delta, b, c


def stream(model, x):
    conv = torch.zeros(x.shape[0], 8, 2, dtype=x.dtype)
    state = torch.zeros(x.shape[0], 8, 3, dtype=x.dtype)
    out = []
    a = -model.scan.A_log.exp()
    for token in x.unbind(1):
        raw, gate = model.in_projection(model.norm(token)).chunk(2, -1)
        window = torch.cat([conv, raw[:, :, None]], -1)
        conv = window[:, :, 1:]
        u = F.silu(
            F.conv1d(
                window, model.local_conv.weight, model.local_conv.bias, groups=8
            ).squeeze(-1)
        )
        delta, b, c = model.selection(u)
        state = (delta[:, :, None] * a).exp() * state + torch.expm1(
            delta[:, :, None] * a
        ) / a * b[:, None, :] * u[:, :, None]
        mixed = (state * c[:, None, :]).sum(-1) + model.scan.direct * u
        out.append(token + model.project(mixed * F.silu(gate)))
    return torch.stack(out, 1)


def setup():
    torch.manual_seed(59)
    return TinyMamba().double(), torch.randn(1, 6, 4, dtype=torch.float64)


def verify():
    model, x = setup()
    changed = x.clone()
    changed[:, 4:] += torch.tensor([1.0, -2.0, 3.0, -1.0])
    torch.testing.assert_close(model(x)[:, :4], model(changed)[:, :4])
    torch.testing.assert_close(stream(model, x), model(x))
    u, dt, b, c = intermediates(model, x)
    out, states, decays, drives = scan_trace(model.scan, u, dt, b, c)
    torch.testing.assert_close(out, model.scan(u, dt, b, c))
    # Independent sum of all past drives with intervening decay products.
    closed = torch.zeros_like(states[:, 0])
    for j in range(6):
        weight = torch.ones_like(closed)
        for k in range(j + 1, 6):
            weight = weight * decays[:, k]
        closed = closed + weight * drives[:, j]
    torch.testing.assert_close(closed, states[:, -1])
    assert torch.all((decays > 0) & (decays < 1))
    assert not torch.allclose(dt[:, 0], dt[:, 1])
    # Suppress input writes to isolate forgetting of an existing state.
    assert torch.all(decays[:, 0] * torch.ones_like(closed) < 1)
    model(x).square().sum().backward()
    assert (
        model.selection.delta.weight.grad.norm() > 0
        and model.selection.input_map.weight.grad.norm() > 0
    )
    return [
        "Future changes leave every earlier output unchanged",
        "Token streaming with a 16-value convolution cache and 24-value SSM state matches the full forward",
        "The recurrent state matches an independent sum of decayed past writes",
        "Negative diagonal A and positive input-dependent delta give decay factors between zero and one; selection maps receive gradients",
    ]


def experiment():
    model, x = setup()
    cases = []
    with torch.no_grad():
        for identity, label, at, change in [
            ("first", "After token 0", 0, False),
            ("middle", "After token 2", 2, False),
            ("last", "After token 5", 5, False),
            ("changed", "Change token 2 · inspect final state", 5, True),
        ]:
            tokens = x.clone()
            if change:
                tokens[:, 2] += torch.tensor([2.0, -3.0, 1.0, 0.0])
            u, dt, b, c = intermediates(model, tokens)
            out, states, decays, drives = scan_trace(model.scan, u, dt, b, c)
            cases.append(
                dict(
                    id=identity,
                    label=label,
                    target="scan",
                    note=f"The recurrence has processed tokens 0 through {at}. This table shows channel 0’s three state coordinates; the complete state holds eight channels. Delta, B and C come from the current convolved token. Exact diagonal ZOH gives Ā = exp(ΔA) and B̄ = (exp(ΔA)−1)B/A. "
                    + (
                        "Token 2 changes while later input tokens stay fixed, exposing the carried influence."
                        if change
                        else "This is a recorded CPU loop, without the fused parallel scan or measured GPU speed."
                    ),
                    matrices=[
                        dict(
                            label="Channel 0 state history · processed tokens × state coordinates",
                            values=states[0, : at + 1, 0].tolist(),
                        ),
                        dict(
                            label="Current full state · channels × coordinates",
                            values=states[0, at].tolist(),
                        ),
                    ],
                    vectors=[
                        dict(
                            label="Current input token", values=tokens[0, at].tolist()
                        ),
                        dict(
                            label="Current delta · eight channels",
                            values=dt[0, at].tolist(),
                        ),
                        dict(
                            label="Channel 0 decay factors",
                            values=decays[0, at, 0].tolist(),
                        ),
                        dict(
                            label="Channel 0 write contribution",
                            values=drives[0, at, 0].tolist(),
                        ),
                    ],
                    metrics=[
                        dict(label="SSM state values per sequence", value=24),
                        dict(label="Convolution history values per sequence", value=16),
                    ],
                )
            )
    return dict(
        kind="matrices",
        title="Read, write and retain a recurrent state.",
        description="Recorded selective-state evolution in a tiny Mamba-style block. Streaming state size is fixed for this architecture; these counts do not include weights or training activations.",
        controlLabel="Scan checkpoint",
        cases=cases,
    )