import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange

class Model(nn.Module):
    def __init__(self, batch_size, seq_length, n_heads, d_head, d_state, block_len=64):
        """
        Mamba Structured State Space model implementation for benchmarking.
        
        :param batch_size: Size of the batch
        :param seq_length: Length of the input sequence
        :param n_heads: Number of attention heads
        :param d_head: Dimension of each head
        :param d_state: Dimension of the state space
        :param block_len: Length of each block for chunked computation
        """
        super(Model, self).__init__()
        
        assert seq_length % block_len == 0, "Sequence length must be divisible by block length"
        
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.n_heads = n_heads
        self.d_head = d_head
        self.d_state = d_state
        self.block_len = block_len
        
        # Initialize parameters
        self.A = nn.Parameter(torch.randn(batch_size, seq_length, n_heads))
        self.B = nn.Parameter(torch.randn(batch_size, seq_length, n_heads, d_state))
        self.C = nn.Parameter(torch.randn(batch_size, seq_length, n_heads, d_state))
        
    def segsum(self, x):
        """Naive segment sum calculation."""
        T = x.size(-1)
        x_cumsum = torch.cumsum(x, dim=-1)
        x_segsum = x_cumsum[..., :, None] - x_cumsum[..., None, :]
        mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=0)
        x_segsum = x_segsum.masked_fill(~mask, -torch.inf)
        return x_segsum
    
    def forward(self, X, initial_states=None):
        """
        Forward pass implementing the SSD operation.
        
        :param X: Input tensor of shape (batch, length, n_heads, d_head)
        :param initial_states: Optional initial states
        :return: Output tensor Y and final state
        """
        # Rearrange into blocks/chunks
        X_blocks, A_blocks, B_blocks, C_blocks = [
            rearrange(x, "b (c l) ... -> b c l ...", l=self.block_len)
            for x in (X, self.A, self.B, self.C)
        ]
        
        A_blocks = rearrange(A_blocks, "b c l h -> b h c l")
        A_cumsum = torch.cumsum(A_blocks, dim=-1)
        
        # 1. Compute diagonal block outputs
        L = torch.exp(self.segsum(A_blocks))
        Y_diag = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", 
                             C_blocks, B_blocks, L, X_blocks)
        
        # 2. Compute intra-chunk states
        decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum))
        states = torch.einsum("bclhn,bhcl,bclhp->bchpn", 
                            B_blocks, decay_states, X_blocks)
        
        # 3. Compute inter-chunk recurrence
        if initial_states is None:
            initial_states = torch.zeros_like(states[:, :1])
        states = torch.cat([initial_states, states], dim=1)
        
        decay_chunk = torch.exp(self.segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
        new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
        states = new_states[:, :-1]
        
        # 4. Compute state-to-output conversion
        state_decay_out = torch.exp(A_cumsum)
        Y_off = torch.einsum('bclhn,bchpn,bhcl->bclhp', 
                           C_blocks, states, state_decay_out)
        
        # Combine diagonal and off-diagonal terms
        Y = rearrange(Y_diag + Y_off, "b c l h p -> b (c l) h p")
        
        
        return Y

# Test parameters
batch_size = 2048
seq_length = 128
n_heads = 8
d_head = 64
d_state = 16
block_len = 64

def get_inputs():
    return [torch.rand(batch_size, seq_length, n_heads, d_head)]

def get_init_inputs():
    return [batch_size, seq_length, n_heads, d_head, d_state, block_len]

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


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

torch::Tensor mamba_forward_cuda(torch::Tensor X, torch::Tensor A, torch::Tensor Bp, torch::Tensor Cp);
torch::Tensor mamba_forward_init_cuda(torch::Tensor X, torch::Tensor A, torch::Tensor Bp, torch::Tensor Cp, torch::Tensor init);

torch::Tensor mamba_forward(torch::Tensor X, torch::Tensor A, torch::Tensor Bp, torch::Tensor Cp) {
    return mamba_forward_cuda(X, A, Bp, Cp);
}

torch::Tensor mamba_forward_init(torch::Tensor X, torch::Tensor A, torch::Tensor Bp, torch::Tensor Cp, torch::Tensor init) {
    return mamba_forward_init_cuda(X, A, Bp, Cp, init);
}
'''

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

#ifndef CHECK_CUDA
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#endif
#ifndef CHECK_CONTIGUOUS
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#endif

// One CUDA block computes one (batch, head) plane.  Threads 0..P-1 compute
// the P channels independently, while A/B/C for the whole sequence are cached
// once in shared memory and broadcast to all P channels.
template <bool HAS_INIT>
__global__ void ssd_recurrence_kernel(
    const float* __restrict__ X,
    const float* __restrict__ A,
    const float* __restrict__ Bp,
    const float* __restrict__ Cp,
    const float* __restrict__ Init,
    float* __restrict__ Y,
    int Bsz, int T, int H, int P, int N)
{
    extern __shared__ float smem[];
    float* __restrict__ E  = smem;          // [T] exp(A)
    float* __restrict__ Bs = E + T;         // [T, N]
    float* __restrict__ Cs = Bs + T * N;    // [T, N]

    const int bh = blockIdx.x;
    const int b = bh / H;
    const int h = bh - b * H;
    const int tid = threadIdx.x;

    const int TN = T * N;
    const int total = T + 2 * TN;

    for (int i = tid; i < total; i += blockDim.x) {
        if (i < T) {
            int t = i;
            E[t] = expf(A[(b * T + t) * H + h]);
        } else if (i < T + TN) {
            int j = i - T;
            int t = j / N;
            int n = j - t * N;
            Bs[j] = Bp[((b * T + t) * H + h) * N + n];
        } else {
            int j = i - T - TN;
            int t = j / N;
            int n = j - t * N;
            Cs[j] = Cp[((b * T + t) * H + h) * N + n];
        }
    }
    __syncthreads();

    // The benchmark uses P=64, N=16.  The generic arguments are still used for
    // indexing, but the state is intentionally kept in registers with fixed max.
    if (tid < P) {
        const int p = tid;
        float s0, s1, s2, s3, s4, s5, s6, s7;
        float s8, s9, s10, s11, s12, s13, s14, s15;
        if (HAS_INIT) {
            const long long base = (((long long)b * H + h) * P + p) * N;
            s0  = (N >  0) ? Init[base +  0] : 0.0f;
            s1  = (N >  1) ? Init[base +  1] : 0.0f;
            s2  = (N >  2) ? Init[base +  2] : 0.0f;
            s3  = (N >  3) ? Init[base +  3] : 0.0f;
            s4  = (N >  4) ? Init[base +  4] : 0.0f;
            s5  = (N >  5) ? Init[base +  5] : 0.0f;
            s6  = (N >  6) ? Init[base +  6] : 0.0f;
            s7  = (N >  7) ? Init[base +  7] : 0.0f;
            s8  = (N >  8) ? Init[base +  8] : 0.0f;
            s9  = (N >  9) ? Init[base +  9] : 0.0f;
            s10 = (N > 10) ? Init[base + 10] : 0.0f;
            s11 = (N > 11) ? Init[base + 11] : 0.0f;
            s12 = (N > 12) ? Init[base + 12] : 0.0f;
            s13 = (N > 13) ? Init[base + 13] : 0.0f;
            s14 = (N > 14) ? Init[base + 14] : 0.0f;
            s15 = (N > 15) ? Init[base + 15] : 0.0f;
        } else {
            s0=s1=s2=s3=s4=s5=s6=s7=s8=s9=s10=s11=s12=s13=s14=s15=0.0f;
        }

        #pragma unroll 1
        for (int t = 0; t < T; ++t) {
            const float x = X[((long long)(b * T + t) * H + h) * P + p];
            const float e = E[t];
            const int off = t * N;

            if (N >  0) s0  = fmaf(e, s0,  Bs[off +  0] * x);
            if (N >  1) s1  = fmaf(e, s1,  Bs[off +  1] * x);
            if (N >  2) s2  = fmaf(e, s2,  Bs[off +  2] * x);
            if (N >  3) s3  = fmaf(e, s3,  Bs[off +  3] * x);
            if (N >  4) s4  = fmaf(e, s4,  Bs[off +  4] * x);
            if (N >  5) s5  = fmaf(e, s5,  Bs[off +  5] * x);
            if (N >  6) s6  = fmaf(e, s6,  Bs[off +  6] * x);
            if (N >  7) s7  = fmaf(e, s7,  Bs[off +  7] * x);
            if (N >  8) s8  = fmaf(e, s8,  Bs[off +  8] * x);
            if (N >  9) s9  = fmaf(e, s9,  Bs[off +  9] * x);
            if (N > 10) s10 = fmaf(e, s10, Bs[off + 10] * x);
            if (N > 11) s11 = fmaf(e, s11, Bs[off + 11] * x);
            if (N > 12) s12 = fmaf(e, s12, Bs[off + 12] * x);
            if (N > 13) s13 = fmaf(e, s13, Bs[off + 13] * x);
            if (N > 14) s14 = fmaf(e, s14, Bs[off + 14] * x);
            if (N > 15) s15 = fmaf(e, s15, Bs[off + 15] * x);

            float y = 0.0f;
            if (N >  0) y = fmaf(Cs[off +  0], s0,  y);
            if (N >  1) y = fmaf(Cs[off +  1], s1,  y);
            if (N >  2) y = fmaf(Cs[off +  2], s2,  y);
            if (N >  3) y = fmaf(Cs[off +  3], s3,  y);
            if (N >  4) y = fmaf(Cs[off +  4], s4,  y);
            if (N >  5) y = fmaf(Cs[off +  5], s5,  y);
            if (N >  6) y = fmaf(Cs[off +  6], s6,  y);
            if (N >  7) y = fmaf(Cs[off +  7], s7,  y);
            if (N >  8) y = fmaf(Cs[off +  8], s8,  y);
            if (N >  9) y = fmaf(Cs[off +  9], s9,  y);
            if (N > 10) y = fmaf(Cs[off + 10], s10, y);
            if (N > 11) y = fmaf(Cs[off + 11], s11, y);
            if (N > 12) y = fmaf(Cs[off + 12], s12, y);
            if (N > 13) y = fmaf(Cs[off + 13], s13, y);
            if (N > 14) y = fmaf(Cs[off + 14], s14, y);
            if (N > 15) y = fmaf(Cs[off + 15], s15, y);

            Y[((long long)(b * T + t) * H + h) * P + p] = y;
        }
    }
}

static torch::Tensor launch_common(torch::Tensor X, torch::Tensor A, torch::Tensor Bp, torch::Tensor Cp, torch::Tensor init, bool has_init) {
    CHECK_CUDA(X); CHECK_CUDA(A); CHECK_CUDA(Bp); CHECK_CUDA(Cp);
    CHECK_CONTIGUOUS(X); CHECK_CONTIGUOUS(A); CHECK_CONTIGUOUS(Bp); CHECK_CONTIGUOUS(Cp);
    TORCH_CHECK(X.scalar_type() == at::kFloat, "X must be float32");
    TORCH_CHECK(A.scalar_type() == at::kFloat && Bp.scalar_type() == at::kFloat && Cp.scalar_type() == at::kFloat, "parameters must be float32");
    const int Bsz = (int)X.size(0);
    const int T   = (int)X.size(1);
    const int H   = (int)X.size(2);
    const int P   = (int)X.size(3);
    const int N   = (int)Bp.size(3);
    TORCH_CHECK(N <= 16, "this optimized kernel supports d_state <= 16");
    auto Y = torch::empty_like(X);
    const int threads = 128;
    const dim3 blocks(Bsz * H);
    const size_t shmem = (size_t)(T + 2 * T * N) * sizeof(float);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (has_init) {
        CHECK_CUDA(init); CHECK_CONTIGUOUS(init);
        ssd_recurrence_kernel<true><<<blocks, threads, shmem, stream>>>(
            X.data_ptr<float>(), A.data_ptr<float>(), Bp.data_ptr<float>(), Cp.data_ptr<float>(), init.data_ptr<float>(), Y.data_ptr<float>(),
            Bsz, T, H, P, N);
    } else {
        ssd_recurrence_kernel<false><<<blocks, threads, shmem, stream>>>(
            X.data_ptr<float>(), A.data_ptr<float>(), Bp.data_ptr<float>(), Cp.data_ptr<float>(), nullptr, Y.data_ptr<float>(),
            Bsz, T, H, P, N);
    }
    return Y;
}

torch::Tensor mamba_forward_cuda(torch::Tensor X, torch::Tensor A, torch::Tensor Bp, torch::Tensor Cp) {
    return launch_common(X, A, Bp, Cp, torch::Tensor(), false);
}

torch::Tensor mamba_forward_init_cuda(torch::Tensor X, torch::Tensor A, torch::Tensor Bp, torch::Tensor Cp, torch::Tensor init) {
    return launch_common(X, A, Bp, Cp, init, true);
}
'''

_ext = load_inline(
    name='ssd_recurrence_ext_v1',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['mamba_forward', 'mamba_forward_init'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, batch_size, seq_length, n_heads, d_head, d_state, block_len=64):
        super(ModelNew, self).__init__()
        assert seq_length % block_len == 0, "Sequence length must be divisible by block length"
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.n_heads = n_heads
        self.d_head = d_head
        self.d_state = d_state
        self.block_len = block_len
        self.A = nn.Parameter(torch.randn(batch_size, seq_length, n_heads))
        self.B = nn.Parameter(torch.randn(batch_size, seq_length, n_heads, d_state))
        self.C = nn.Parameter(torch.randn(batch_size, seq_length, n_heads, d_state))

    def forward(self, X, initial_states=None):
        # The reference SSD computation is algebraically the linear recurrence
        # S_t = exp(A_t) S_{t-1} + B_t * X_t,  Y_t = C_t dot S_t.
        if initial_states is None:
            return _ext.mamba_forward(X, self.A, self.B, self.C)
        return _ext.mamba_forward_init(X, self.A, self.B, self.C, initial_states)