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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor causal_conv_silu_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias);
torch::Tensor selective_scan_forward(torch::Tensor x, torch::Tensor delta, torch::Tensor Bt, torch::Tensor Ct, torch::Tensor A_log, torch::Tensor D);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void causal_conv_silu_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    long sB, long sL, long sI,
    int B, int L, int I, int K)
{
    long total = (long)B * L * I;
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long stride = (long)blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int i = (int)(idx % I);
        long tmp = idx / I;
        int t = (int)(tmp % L);
        int b = (int)(tmp / L);
        float acc = bias[i];
        // PyTorch Conv1d with padding K-1, returning the first L positions:
        // input index = t + k - (K-1).
        #pragma unroll
        for (int k = 0; k < 8; ++k) {
            if (k < K) {
                int ti = t + k - (K - 1);
                if (ti >= 0) {
                    acc += x[(long)b * sB + (long)ti * sL + (long)i * sI] * w[(long)i * K + k];
                }
            }
        }
        float sig = 1.0f / (1.0f + expf(-acc));
        out[idx] = acc * sig;
    }
}

__global__ void selective_scan_kernel(
    const float* __restrict__ x,
    const float* __restrict__ delta,
    const float* __restrict__ Bt,
    const float* __restrict__ Ct,
    const float* __restrict__ A_log,
    const float* __restrict__ D,
    float* __restrict__ y,
    int B, int L, int I, int N)
{
    int lane = threadIdx.x & 31;
    int warp_in_block = threadIdx.x >> 5;
    int warps_per_block = blockDim.x >> 5;
    int wi = blockIdx.x * warps_per_block + warp_in_block;
    int total_warps = B * I;
    if (wi >= total_warps) return;

    int b = wi / I;
    int i = wi - b * I;
    int n = lane;

    float h = 0.0f;
    float aval = 0.0f;
    if (n < N) {
        aval = -expf(A_log[(long)i * N + n]);
    }
    float dskip = D[i];

    for (int t = 0; t < L; ++t) {
        long xi = ((long)b * L + t) * I + i;
        float xv = x[xi];
        float del = delta[xi];
        float prod = 0.0f;
        if (n < N) {
            long bcn = ((long)b * L + t) * N + n;
            float da = expf(del * aval);
            h = da * h + (xv * del) * Bt[bcn];
            prod = h * Ct[bcn];
        }
        // warp reduction over the N=16 active lanes (remaining lanes hold zero)
        prod += __shfl_down_sync(0xffffffff, prod, 16);
        prod += __shfl_down_sync(0xffffffff, prod, 8);
        prod += __shfl_down_sync(0xffffffff, prod, 4);
        prod += __shfl_down_sync(0xffffffff, prod, 2);
        prod += __shfl_down_sync(0xffffffff, prod, 1);
        if (lane == 0) {
            y[xi] = prod + xv * dskip;
        }
    }
}

torch::Tensor causal_conv_silu_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda() && bias.is_cuda(), "tensors must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(w.scalar_type() == at::kFloat && bias.scalar_type() == at::kFloat, "only float32 is supported");
    int B = (int)x.size(0);
    int L = (int)x.size(1);
    int I = (int)x.size(2);
    int K = (int)w.size(2);
    auto out = torch::empty({B, L, I}, x.options());
    long total = (long)B * L * I;
    int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    causal_conv_silu_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(),
        x.stride(0), x.stride(1), x.stride(2), B, L, I, K);
    return out;
}

torch::Tensor selective_scan_forward(torch::Tensor x, torch::Tensor delta, torch::Tensor Bt, torch::Tensor Ct, torch::Tensor A_log, torch::Tensor D) {
    TORCH_CHECK(x.is_cuda() && delta.is_cuda() && Bt.is_cuda() && Ct.is_cuda() && A_log.is_cuda() && D.is_cuda(), "tensors must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    int B = (int)x.size(0);
    int L = (int)x.size(1);
    int I = (int)x.size(2);
    int N = (int)Bt.size(2);
    auto y = torch::empty_like(x);
    int threads = 128; // 4 warps; each warp scans one (batch, channel)
    int warps_per_block = threads / 32;
    int blocks = (B * I + warps_per_block - 1) / warps_per_block;
    selective_scan_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), delta.data_ptr<float>(), Bt.data_ptr<float>(), Ct.data_ptr<float>(),
        A_log.data_ptr<float>(), D.data_ptr<float>(), y.data_ptr<float>(), B, L, I, N);
    return y;
}
"""

_mamba_ext = load_inline(
    name="mamba_block_custom_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["causal_conv_silu_forward", "selective_scan_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

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)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x_norm = self.norm(x)
        xz = F.linear(x_norm, self.in_proj.weight, None)
        I = self.d_inner
        x_branch = xz[..., :I]
        z = xz[..., I:]

        # Fused causal depthwise convolution over the x branch and SiLU activation.
        x_branch = _mamba_ext.causal_conv_silu_forward(x_branch, self.conv1d.weight, self.conv1d.bias)

        delta = F.softplus(F.linear(x_branch, self.proj_delta.weight, self.proj_delta.bias))
        B_t = F.linear(x_branch, self.proj_B.weight, None)
        C_t = F.linear(x_branch, self.proj_C.weight, None)

        y = _mamba_ext.selective_scan_forward(x_branch, delta, B_t, C_t, self.A_log, self.D)
        z = F.silu(z)
        y = y * z
        y = F.linear(y, self.out_proj.weight, None)
        return y + residual