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

_cpp = r'''
#include <torch/extension.h>
#include <vector>
void conv_silu_z_launcher(torch::Tensor xz, torch::Tensor weight, torch::Tensor bias, torch::Tensor xb, torch::Tensor z);
void scan_launcher(torch::Tensor xb, torch::Tensor z, torch::Tensor delta_pre, torch::Tensor Bt, torch::Tensor Ct, torch::Tensor A_log, torch::Tensor D, torch::Tensor y);

std::vector<torch::Tensor> conv_silu_z(torch::Tensor xz, torch::Tensor weight, torch::Tensor bias) {
    auto B = xz.size(0), L = xz.size(1), I = xz.size(2) / 2;
    auto xb = torch::empty({B, L, I}, xz.options());
    auto z = torch::empty({B, L, I}, xz.options());
    conv_silu_z_launcher(xz, weight, bias, xb, z);
    return {xb, z};
}

torch::Tensor selective_scan_gated(torch::Tensor xb, torch::Tensor z, torch::Tensor delta_pre, torch::Tensor Bt, torch::Tensor Ct, torch::Tensor A_log, torch::Tensor D) {
    auto y = torch::empty_like(xb);
    scan_launcher(xb, z, delta_pre, Bt, Ct, A_log, D, y);
    return y;
}
'''

_cuda = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__device__ __forceinline__ float silu_f(float x) { return x / (1.0f + expf(-x)); }
__device__ __forceinline__ float softplus_f(float x) { return (x > 20.0f) ? x : log1pf(expf(x)); }

__global__ void conv_silu_z_kernel(const float* __restrict__ xz, const float* __restrict__ w,
                                   const float* __restrict__ bias, float* __restrict__ xb,
                                   float* __restrict__ z, int B, int L, int I, int K) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)B * L * I;
    if (idx >= total) return;
    int i = idx % I;
    int t = (idx / I) % L;
    int b = idx / ((long long)L * I);
    int twoI = I << 1;
    float acc = bias[i];
    #pragma unroll
    for (int k = 0; k < 4; ++k) {
        int tt = t + k - (K - 1);
        if (tt >= 0 && tt < L && k < K)
            acc += xz[((long long)b * L + tt) * twoI + i] * w[(long long)i * K + k];
    }
    xb[idx] = silu_f(acc);
    z[idx] = silu_f(xz[((long long)b * L + t) * twoI + I + i]);
}

__global__ void scan_shared_kernel(const float* __restrict__ xb,
                                   const float* __restrict__ z,
                                   const float* __restrict__ delta_pre,
                                   const float* __restrict__ Bt,
                                   const float* __restrict__ Ct,
                                   const float* __restrict__ A_log,
                                   const float* __restrict__ Dv,
                                   float* __restrict__ y,
                                   int B, int L, int I, int N) {
    __shared__ float Bs[16];
    __shared__ float Cs[16];
    int b = blockIdx.y;
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    bool valid = (i < I);
    float h[16];
    float A[16];
    #pragma unroll
    for (int n=0; n<16; ++n) {
        h[n] = 0.0f;
        A[n] = (valid && n < N) ? -expf(A_log[(long long)i * N + n]) : 0.0f;
    }
    float Dval = valid ? Dv[i] : 0.0f;
    for (int t=0; t<L; ++t) {
        if (threadIdx.x < 16) {
            Bs[threadIdx.x] = Bt[((long long)b * L + t) * N + threadIdx.x];
            Cs[threadIdx.x] = Ct[((long long)b * L + t) * N + threadIdx.x];
        }
        __syncthreads();
        if (valid) {
            long long xi = ((long long)b * L + t) * I + i;
            float xval = xb[xi];
            float delta = softplus_f(delta_pre[xi]);
            float xd = xval * delta;
            float sum = 0.0f;
            #pragma unroll
            for (int n=0; n<16; ++n) {
                float hv = expf(delta * A[n]) * h[n] + xd * Bs[n];
                h[n] = hv;
                sum += hv * Cs[n];
            }
            y[xi] = (sum + xval * Dval) * z[xi];
        }
        __syncthreads();
    }
}

void conv_silu_z_launcher(torch::Tensor xz, torch::Tensor weight, torch::Tensor bias, torch::Tensor xb, torch::Tensor z) {
    int B=(int)xz.size(0), L=(int)xz.size(1), I=(int)(xz.size(2)/2), K=(int)weight.size(2);
    int threads=256; long long total=(long long)B*L*I; int blocks=(int)((total+threads-1)/threads);
    conv_silu_z_kernel<<<blocks,threads>>>(xz.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), xb.data_ptr<float>(), z.data_ptr<float>(), B,L,I,K);
}
void scan_launcher(torch::Tensor xb, torch::Tensor z, torch::Tensor delta_pre, torch::Tensor Bt, torch::Tensor Ct, torch::Tensor A_log, torch::Tensor D, torch::Tensor y) {
    int B=(int)xb.size(0), L=(int)xb.size(1), I=(int)xb.size(2), N=(int)Bt.size(2);
    dim3 block(64); dim3 grid((I+63)/64, B);
    scan_shared_kernel<<<grid,block>>>(xb.data_ptr<float>(), z.data_ptr<float>(), delta_pre.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);
}
'''

_ext = load_inline(name='mamba_block_ext_final', cpp_sources=_cpp, cuda_sources=_cuda,
                   functions=['conv_silu_z','selective_scan_gated'], 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
        xz=self.in_proj(self.norm(x))
        x_branch,z=_ext.conv_silu_z(xz,self.conv1d.weight,self.conv1d.bias)
        delta_pre=self.proj_delta(x_branch)
        B_t=self.proj_B(x_branch)
        C_t=self.proj_C(x_branch)
        y=_ext.selective_scan_gated(x_branch,z,delta_pre,B_t,C_t,self.A_log,self.D)
        return self.out_proj(y)+residual