import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Full Mamba block: RMSNorm -> linear expansion into two branches -> causal depthwise Conv1d ->
    SiLU -> selective scan (input-dependent SSM) -> elementwise gate -> linear projection -> residual.
    This is the complete Mamba architecture block as described in Gu & Dao (2023), used in Mamba,
    Caduceus (bidirectional DNA), and various protein/genomic language models.

    Args:
        d_model (int): Model dimension.
        d_state (int): SSM state dimension (N).
        d_conv (int): Causal convolution kernel size.
        expand (int): Expansion factor for inner dimension.
    """
    def __init__(self, d_model: int, d_state: int, d_conv: int, expand: int):
        super(Model, self).__init__()
        self.d_model = d_model
        self.d_inner = d_model * expand
        self.d_state = d_state
        self.d_conv = d_conv

        self.norm = nn.RMSNorm(d_model)
        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)
        self.conv1d = nn.Conv1d(self.d_inner, self.d_inner, d_conv,
                                padding=d_conv - 1, groups=self.d_inner, bias=True)
        self.activation = nn.SiLU()

        self.A_log = nn.Parameter(torch.randn(self.d_inner, d_state))
        self.proj_B = nn.Linear(self.d_inner, d_state, bias=False)
        self.proj_C = nn.Linear(self.d_inner, d_state, bias=False)
        self.proj_delta = nn.Linear(self.d_inner, self.d_inner, bias=True)
        self.D = nn.Parameter(torch.ones(self.d_inner))

        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)

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

        xz = self.in_proj(x)
        x_branch, z = xz.chunk(2, dim=-1)

        x_branch = x_branch.transpose(1, 2)
        x_branch = self.conv1d(x_branch)[..., :L]
        x_branch = x_branch.transpose(1, 2)
        x_branch = self.activation(x_branch)

        N = self.d_state
        delta = torch.nn.functional.softplus(self.proj_delta(x_branch))
        B_t = self.proj_B(x_branch)
        C_t = self.proj_C(x_branch)

        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_branch.unsqueeze(-1) * dB

        h = torch.zeros(B, self.d_inner, 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_branch * self.D.unsqueeze(0).unsqueeze(0)

        z = self.activation(z)
        y = y * z
        y = self.out_proj(y)
        y = y + residual
        return y

batch_size = 16
d_model = 256
d_state = 16
d_conv = 4
expand = 2
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, d_conv, expand]

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

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

__global__ void selective_scan_kernel(const float* dA, const float* x_db, const float* C_t, float* 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[32];
    for(int n=0; n<N; n++) h[n]=0.0f;
    for(int t=0; t<L; t++) {
        float yt = 0.0f;
        for(int n=0; n<N; n++) {
            int base = ((b * L + t) * D + d) * N + n;
            float da = dA[base];
            float hnew = da * h[n] + x_db[base];
            h[n] = hnew;
            yt += hnew * C_t[(b * L + t) * N + n];
        }
        y[(b * L + t) * D + d] = yt;
    }
}

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

cpp_source = r'''
torch::Tensor selective_scan_cuda(torch::Tensor dA, torch::Tensor x_db, torch::Tensor Ct);
'''

class ModelNew(nn.Module):
    def __init__(self, d_model: int, d_state: int, d_conv: int, expand: int):
        super(ModelNew, self).__init__()
        self.d_model = d_model
        self.d_inner = d_model * expand
        self.d_state = d_state
        self.d_conv = d_conv

        self.norm = nn.RMSNorm(d_model)
        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)
        self.conv1d = nn.Conv1d(self.d_inner, self.d_inner, d_conv, padding=d_conv - 1, groups=self.d_inner, bias=True)
        self.activation = nn.SiLU()

        self.A_log = nn.Parameter(torch.randn(self.d_inner, d_state))
        self.proj_B = nn.Linear(self.d_inner, d_state, bias=False)
        self.proj_C = nn.Linear(self.d_inner, d_state, bias=False)
        self.proj_delta = nn.Linear(self.d_inner, self.d_inner, bias=True)
        self.D = nn.Parameter(torch.ones(self.d_inner))

        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)

        self.scan_module = load_inline(
            name="mamba_scan",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["selective_scan_cuda"],
            verbose=False,
            with_cuda=True
        )

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

        xz = self.in_proj(x)
        x_branch, z = xz.chunk(2, dim=-1)

        x_branch = x_branch.transpose(1, 2)
        x_branch = self.conv1d(x_branch)[..., :L]
        x_branch = x_branch.transpose(1, 2)
        x_branch = self.activation(x_branch)

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

        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_branch.unsqueeze(-1) * dB

        y = self.scan_module.selective_scan_cuda(dA, x_db, C_t)
        y = y + x_branch * self.D.unsqueeze(0).unsqueeze(0)

        z = self.activation(z)
        y = y * z
        y = self.out_proj(y)
        y = y + residual
        return y