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

mamba_cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
#include <cfloat>

// RMSNorm kernel - one block per (batch, seq) position
__global__ void __launch_bounds__(256) rmsnorm_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    float* __restrict__ output,
    int B, int L, int D,
    float eps
) {
    int b = blockIdx.x;
    int l = blockIdx.y;
    if (b >= B || l >= L) return;
    
    const float* in_ptr = input + (b * L + l) * D;
    float* out_ptr = output + (b * L + l) * D;
    
    // Compute sum of squares using parallel reduction
    float local_sum = 0.0f;
    for (int d = threadIdx.x; d < D; d += blockDim.x) {
        float val = in_ptr[d];
        local_sum += val * val;
    }
    
    // Warp-level reduction
    for (int offset = 16; offset > 0; offset >>= 1) {
        local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);
    }
    
    // Shared memory for cross-warp reduction
    __shared__ float s_sum[8];
    int warp_id = threadIdx.x / 32;
    int lane_id = threadIdx.x % 32;
    
    if (lane_id == 0) {
        s_sum[warp_id] = local_sum;
    }
    __syncthreads();
    
    // First warp reduces across warps
    float total_sum = 0.0f;
    int num_warps = (blockDim.x + 31) / 32;
    if (warp_id == 0) {
        if (lane_id < num_warps) {
            total_sum = s_sum[lane_id];
        }
        for (int offset = 16; offset > 0; offset >>= 1) {
            total_sum += __shfl_down_sync(0xffffffff, total_sum, offset);
        }
        if (lane_id == 0) {
            s_sum[0] = total_sum;
        }
    }
    __syncthreads();
    
    float rms = rsqrtf(s_sum[0] / D + eps);
    
    // Apply normalization
    for (int d = threadIdx.x; d < D; d += blockDim.x) {
        out_ptr[d] = in_ptr[d] * rms * weight[d];
    }
}

// Selective scan kernel - processes one (batch, channel) pair per thread block
// Each thread handles one state dimension
__global__ void selective_scan_kernel(
    const float* __restrict__ x_branch,     // (B, L, d_inner)
    const float* __restrict__ delta,         // (B, L, d_inner)
    const float* __restrict__ B_t,           // (B, L, d_state)
    const float* __restrict__ C_t,           // (B, L, d_state)
    const float* __restrict__ A,             // (d_inner, d_state) - already negated exp
    const float* __restrict__ D_param,       // (d_inner,)
    float* __restrict__ output,              // (B, L, d_inner)
    int B, int L, int d_inner, int d_state
) {
    int b = blockIdx.x;
    int c = blockIdx.y;  // channel in d_inner
    
    if (b >= B || c >= d_inner) return;
    
    int tid = threadIdx.x;
    int state_idx = tid;
    
    // Initialize state
    float h = 0.0f;
    
    // Pointers for this batch and channel
    const float* x_ptr = x_branch + b * L * d_inner + c;
    const float* delta_ptr = delta + b * L * d_inner + c;
    
    // Output pointer
    float* out_ptr = output + b * L * d_inner + c;
    
    // Shared memory for reduction - pad to avoid bank conflicts
    extern __shared__ float shared_mem[];
    float* y_shared = shared_mem;
    
    for (int t = 0; t < L; t++) {
        float x_val = x_ptr[t * d_inner];
        float delta_val = delta_ptr[t * d_inner];
        
        float y_t = 0.0f;
        
        if (state_idx < d_state) {
            const float* B_ptr = B_t + (b * L + t) * d_state;
            const float* C_ptr = C_t + (b * L + t) * d_state;
            const float* A_ptr = A + c * d_state;
            
            float a = A_ptr[state_idx];
            float b_val = B_ptr[state_idx];
            float c_val = C_ptr[state_idx];
            
            float dA = expf(delta_val * a);
            float dB = delta_val * b_val;
            
            h = dA * h + x_val * dB;
            
            y_t = h * c_val;
        }
        
        y_shared[state_idx] = y_t;
        __syncthreads();
        
        // Reduce y_t across state dimensions - use parallel reduction
        if (state_idx == 0) {
            float sum_y = 0.0f;
            // Unrolled loop for small d_state
            #pragma unroll 16
            for (int s = 0; s < d_state; s++) {
                sum_y += y_shared[s];
            }
            // Add skip connection: y = sum(h*C) + x*D
            sum_y += x_val * D_param[c];
            out_ptr[t * d_inner] = sum_y;
        }
        __syncthreads();
    }
}

// Fused kernel: silu(z) * y -> out_proj -> residual add
// One block per (batch, seq) position, threads process d_inner dimension
__global__ void __launch_bounds__(512) fused_output_kernel(
    const float* __restrict__ y,              // (B, L, d_inner)
    const float* __restrict__ z,              // (B, L, d_inner)
    const float* __restrict__ out_proj_w,     // (D, d_inner)
    const float* __restrict__ residual,       // (B, L, D)
    float* __restrict__ output,               // (B, L, D)
    int B, int L, int d_inner, int D
) {
    int b = blockIdx.x;
    int l = blockIdx.y;
    int d = threadIdx.x;
    
    if (b >= B || l >= L || d >= D) return;
    
    // Shared memory for reduction across d_inner
    __shared__ float s_gated[1024];  // d_inner elements
    
    float sum = 0.0f;
    for (int c = threadIdx.x; c < d_inner; c += blockDim.x) {
        int idx = b * L * d_inner + l * d_inner + c;
        float z_val = z[idx];
        float y_val = y[idx];
        
        // SiLU activation: x * sigmoid(x)
        float silu_z = z_val * (1.0f / (1.0f + expf(-z_val)));
        float gated = silu_z * y_val;
        
        s_gated[c] = gated;
    }
    __syncthreads();
    
    // Compute projection for output dimension d
    sum = 0.0f;
    for (int c = 0; c < d_inner; c++) {
        sum += s_gated[c] * out_proj_w[d * d_inner + c];
    }
    
    // Add residual and write output
    int out_idx = b * L * D + l * D + d;
    output[out_idx] = sum + residual[out_idx];
}

// Entry points
torch::Tensor rmsnorm_entry(torch::Tensor input, torch::Tensor weight) {
    int B = input.size(0);
    int L = input.size(1);
    int D = input.size(2);
    auto output = torch::empty_like(input);
    
    float eps = 1e-5f;
    dim3 grid(B, L);
    int threads = 256;
    
    rmsnorm_kernel<<<grid, threads>>>(
        input.data_ptr<float>(),
        weight.data_ptr<float>(),
        output.data_ptr<float>(),
        B, L, D, eps
    );
    return output;
}

torch::Tensor selective_scan_entry(
    torch::Tensor x_branch,
    torch::Tensor delta,
    torch::Tensor B_t,
    torch::Tensor C_t,
    torch::Tensor A,
    torch::Tensor D_param
) {
    int B = x_branch.size(0);
    int L = x_branch.size(1);
    int d_inner = x_branch.size(2);
    int d_state = B_t.size(2);
    auto output = torch::empty_like(x_branch);
    
    dim3 grid(B, d_inner);
    int threads = max(d_state, 32);
    int shared_mem = threads * sizeof(float);
    
    selective_scan_kernel<<<grid, threads, shared_mem>>>(
        x_branch.data_ptr<float>(),
        delta.data_ptr<float>(),
        B_t.data_ptr<float>(),
        C_t.data_ptr<float>(),
        A.data_ptr<float>(),
        D_param.data_ptr<float>(),
        output.data_ptr<float>(),
        B, L, d_inner, d_state
    );
    return output;
}

torch::Tensor fused_output_entry(
    torch::Tensor y,
    torch::Tensor z,
    torch::Tensor out_proj_w,
    torch::Tensor residual
) {
    int B = y.size(0);
    int L = y.size(1);
    int d_inner = y.size(2);
    int D = residual.size(2);
    auto output = torch::empty_like(residual);
    
    dim3 grid(B, L);
    int threads = min(D, 512);
    
    fused_output_kernel<<<grid, threads>>>(
        y.data_ptr<float>(),
        z.data_ptr<float>(),
        out_proj_w.data_ptr<float>(),
        residual.data_ptr<float>(),
        output.data_ptr<float>(),
        B, L, d_inner, D
    );
    return output;
}
"""

mamba_cpp_source = """
torch::Tensor rmsnorm_entry(torch::Tensor input, torch::Tensor weight);
torch::Tensor selective_scan_entry(
    torch::Tensor x_branch,
    torch::Tensor delta,
    torch::Tensor B_t,
    torch::Tensor C_t,
    torch::Tensor A,
    torch::Tensor D_param
);
torch::Tensor fused_output_entry(
    torch::Tensor y,
    torch::Tensor z,
    torch::Tensor out_proj_w,
    torch::Tensor residual
);
"""

# Load the extension
mamba_extension = load_inline(
    name='mamba_extension_v4',
    cpp_sources=mamba_cpp_source,
    cuda_sources=mamba_cuda_source,
    functions=['rmsnorm_entry', 'selective_scan_entry', 'fused_output_entry'],
    extra_cuda_cflags=['-O3', '-use_fast_math', '-lineinfo'],
    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:
        B, L, D = x.shape
        residual = x
        
        # Custom RMSNorm
        x = mamba_extension.rmsnorm_entry(x.contiguous(), self.norm.weight.contiguous())

        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)
        
        # Use custom selective scan kernel
        y = mamba_extension.selective_scan_entry(
            x_branch.contiguous(),
            delta.contiguous(),
            B_t.contiguous(),
            C_t.contiguous(),
            A.contiguous(),
            self.D.contiguous()
        )

        # Use fused output kernel
        output = mamba_extension.fused_output_entry(
            y.contiguous(),
            z.contiguous(),
            self.out_proj.weight.contiguous(),
            residual.contiguous()
        )
        return output