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

class _LinearParams(nn.Module):
    def __init__(self, in_features, out_features, bias=True):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()
    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            nn.init.uniform_(self.bias, -bound, bound)

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor selective_forward_cuda(torch::Tensor x, torch::Tensor A_log, torch::Tensor wB, torch::Tensor wC, torch::Tensor wDelta, torch::Tensor bDelta, torch::Tensor Dp);
"""

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

__inline__ __device__ float warp_sum(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) v += __shfl_down_sync(0xffffffff, v, offset);
    return v;
}

__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__ Dp,
                                      float* __restrict__ y,
                                      int L, int D, int N) {
    int d = blockIdx.x;
    int b = blockIdx.y;
    int lane = threadIdx.x;
    float h = 0.0f;
    float a = 0.0f;
    if (lane < N) a = -__expf(A_log[d * N + lane]);
    float Dval = Dp[d];
    int base_x = b * L * D + d;
    int base_bn = b * L * N;

    for (int t = 0; t < L; ++t) {
        int idx_x = base_x + t * D;
        float xv = 0.0f;
        float del = 0.0f;
        if (lane == 0) {
            xv = x[idx_x];
            del = delta[idx_x];
        }
        xv = __shfl_sync(0xffffffff, xv, 0);
        del = __shfl_sync(0xffffffff, del, 0);

        float contrib = 0.0f;
        if (lane < N) {
            int idx_n = base_bn + t * N + lane;
            h = __expf(del * a) * h + (xv * del) * Bt[idx_n];
            contrib = h * Ct[idx_n];
        }
        float sum = warp_sum(contrib);
        if (lane == 0) y[idx_x] = sum + xv * Dval;
    }
}

torch::Tensor selective_forward_cuda(torch::Tensor x,
                                     torch::Tensor A_log,
                                     torch::Tensor wB,
                                     torch::Tensor wC,
                                     torch::Tensor wDelta,
                                     torch::Tensor bDelta,
                                     torch::Tensor Dp) {
    auto delta_lin = at::linear(x, wDelta, bDelta);
    auto delta = at::softplus(delta_lin, 1.0, 20.0);
    auto Bt = at::linear(x, wB, c10::nullopt);
    auto Ct = at::linear(x, wC, c10::nullopt);
    int L = x.size(1);
    int D = x.size(2);
    int N = A_log.size(1);
    auto y = torch::empty_like(x);
    selective_scan_kernel<<<dim3(D, x.size(0)), dim3(32)>>>(
        x.data_ptr<float>(), delta.data_ptr<float>(), Bt.data_ptr<float>(), Ct.data_ptr<float>(),
        A_log.data_ptr<float>(), Dp.data_ptr<float>(), y.data_ptr<float>(), L, D, N);
    return y;
}
"""

_ext = load_inline(
    name="selective_scan_ext_final",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["selective_forward_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)

class ModelNew(nn.Module):
    def __init__(self, d_model: int, d_state: int):
        super().__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 = _LinearParams(d_model, d_state, bias=False)
        self.proj_C = _LinearParams(d_model, d_state, bias=False)
        self.proj_delta = _LinearParams(d_model, d_model, bias=True)
        self.D = nn.Parameter(torch.ones(d_model))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.selective_forward_cuda(
            x, self.A_log, self.proj_B.weight, self.proj_C.weight,
            self.proj_delta.weight, self.proj_delta.bias, self.D
        )