import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Selective scan (Mamba's S6 layer): projects input to derive input-dependent SSM parameters
    (delta, B, C), discretizes a learned state matrix A, runs the linear recurrence, and
    applies a skip connection via D. The fusion opportunity is keeping all intermediates
    (discretized parameters, hidden states) in SRAM rather than materializing them to global memory.

    Args:
        d_model (int): Input feature dimension.
        d_state (int): Hidden state dimension per channel (N in Mamba, typically 16).
    """
    def __init__(self, d_model: int, d_state: int):
        super(Model, self).__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.A_log = nn.Parameter(torch.randn(d_model, d_state))
        self.proj_B = nn.Linear(d_model, d_state, bias=False)
        self.proj_C = nn.Linear(d_model, d_state, bias=False)
        self.proj_delta = nn.Linear(d_model, d_model, bias=True)
        self.D = nn.Parameter(torch.ones(d_model))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, L, D = x.shape
        N = self.d_state

        delta = torch.nn.functional.softplus(self.proj_delta(x))
        B_t = self.proj_B(x)
        C_t = self.proj_C(x)

        A = -torch.exp(self.A_log)
        dA = torch.exp(delta.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0))
        dB = delta.unsqueeze(-1) * B_t.unsqueeze(2)
        x_db = x.unsqueeze(-1) * dB

        h = torch.zeros(B, D, N, device=x.device, dtype=x.dtype)
        outputs = []
        for t in range(L):
            h = dA[:, t] * h + x_db[:, t]
            y_t = (h * C_t[:, t].unsqueeze(1)).sum(dim=-1)
            outputs.append(y_t)

        y = torch.stack(outputs, dim=1)
        y = y + x * self.D.unsqueeze(0).unsqueeze(0)
        return y

batch_size = 16
d_model = 256
d_state = 16
seq_len = 2048

def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), popcorn_pri.jitter_int(seq_len), d_model)]

def get_init_inputs():
    return [d_model, d_state]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void selective_scan_kernel(
    const float* __restrict__ x,
    const float* __restrict__ delta,
    const float* __restrict__ B_t,
    const float* __restrict__ C_t,
    const float* __restrict__ dA,
    const float* __restrict__ x_db,
    float* __restrict__ y,
    int B, int L, int D, int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * D;
    if (idx >= total) return;
    int b = idx / D;
    int d = idx % D;

    float h[16];  // assuming N <=16
    for (int n=0; n<N; n++) h[n] = 0.0f;

    for (int t=0; t<L; t++) {
        int base_dn = (b * L + t) * D * N + d * N;
        for (int n=0; n<N; n++) {
            float da = dA[base_dn + n];
            float xdb = x_db[base_dn + n];
            h[n] = da * h[n] + xdb;
        }
        float yt = 0.0f;
        int base_cn = (b * L + t) * N;
        for (int n=0; n<N; n++) {
            yt += h[n] * C_t[base_cn + n];
        }
        y[(b * L + t) * D + d] = yt;
    }
}

torch::Tensor selective_scan_cuda(
    torch::Tensor x,
    torch::Tensor delta,
    torch::Tensor B_t,
    torch::Tensor C_t,
    torch::Tensor dA,
    torch::Tensor x_db
) {
    int B = x.size(0);
    int L = x.size(1);
    int D = x.size(2);
    int N = B_t.size(2);
    auto y = torch::empty({B, L, D}, x.options());
    int total = B * D;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    selective_scan_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(), delta.data_ptr<float>(), B_t.data_ptr<float>(),
        C_t.data_ptr<float>(), dA.data_ptr<float>(), x_db.data_ptr<float>(),
        y.data_ptr<float>(), B, L, D, N
    );
    return y;
}
'''

cpp_source = '''
torch::Tensor selective_scan_cuda(
    torch::Tensor x,
    torch::Tensor delta,
    torch::Tensor B_t,
    torch::Tensor C_t,
    torch::Tensor dA,
    torch::Tensor x_db
);
'''

module = load_inline(
    name='selective_scan',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['selective_scan_cuda'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self, d_model: int, d_state: int):
        super(ModelNew, self).__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.A_log = nn.Parameter(torch.randn(d_model, d_state))
        self.proj_B = nn.Linear(d_model, d_state, bias=False)
        self.proj_C = nn.Linear(d_model, d_state, bias=False)
        self.proj_delta = nn.Linear(d_model, d_model, bias=True)
        self.D = nn.Parameter(torch.ones(d_model))
        self.scan_module = module

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, L, D = x.shape
        N = self.d_state
        delta = torch.nn.functional.softplus(self.proj_delta(x))
        B_t = self.proj_B(x)
        C_t = self.proj_C(x)
        A = -torch.exp(self.A_log)
        dA = torch.exp(delta.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0))
        dB = delta.unsqueeze(-1) * B_t.unsqueeze(2)
        x_db = x.unsqueeze(-1) * dB
        y = self.scan_module.selective_scan_cuda(x, delta, B_t, C_t, dA, x_db)
        y = y + x * self.D.unsqueeze(0).unsqueeze(0)
        return y