import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Pre-norm 1D dilated convolution block: LayerNorm, GELU activation, dilated Conv1d, and residual addition.
    Common in genomic sequence models (Enformer, Basenji) for capturing long-range dependencies in DNA/RNA sequences.
    """
    def __init__(self, channels, seq_len, kernel_size, dilation):
        super(Model, self).__init__()
        self.norm = nn.LayerNorm([channels, seq_len])
        self.activation = nn.GELU()
        self.conv = nn.Conv1d(channels, channels, kernel_size, padding=dilation, dilation=dilation)

    def forward(self, x):
        residual = x
        x = self.norm(x)
        x = self.activation(x)
        x = self.conv(x)
        x = x + residual
        return x

batch_size = 32
channels = 256
seq_len = 4096
kernel_size = 3
dilation = 4

def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), channels, seq_len)]

def get_init_inputs():
    return [channels, seq_len, kernel_size, dilation]

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

# Define CUDA source
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

// Constants for GELU approximation
constexpr float GELU_COEF = 0.7978845608028654f;  // sqrt(2/M_PI)
constexpr float GELU_CUBIC = 0.035677408136300125f;  // 0.044715 * sqrt(2/M_PI) approximation adjustment

// Fused LayerNorm + GELU kernel
template <int BLOCK_SIZE>
__global__ void layernorm_gelu_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int batch_size,
    int channels,
    int seq_len,
    const float* __restrict__ weight,
    const float* __restrict__ bias
) {
    const int batch_idx = blockIdx.x;
    const int tid = threadIdx.x;
    const int stride = blockDim.x;
    const int size = channels * seq_len;
    
    if (batch_idx >= batch_size) return;
    
    // Shared memory for reduction
    __shared__ float s_mean[BLOCK_SIZE];
    __shared__ float s_var[BLOCK_SIZE];
    
    const float* batch_input = input + batch_idx * size;
    float* batch_output = output + batch_idx * size;
    
    // First pass: compute mean
    float local_sum = 0.0f;
    for (int i = tid; i < size; i += stride) {
        local_sum += batch_input[i];
    }
    s_mean[tid] = local_sum;
    __syncthreads();
    
    // Reduction for mean
    for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            s_mean[tid] += s_mean[tid + s];
        }
        __syncthreads();
    }
    
    float mean = s_mean[0] / size;
    
    // Second pass: compute variance
    float local_var_sum = 0.0f;
    for (int i = tid; i < size; i += stride) {
        float diff = batch_input[i] - mean;
        local_var_sum += diff * diff;
    }
    s_var[tid] = local_var_sum;
    __syncthreads();
    
    // Reduction for variance
    for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            s_var[tid] += s_var[tid + s];
        }
        __syncthreads();
    }
    
    float variance = s_var[0] / size;
    float inv_std = rsqrtf(variance + 1e-5f);
    
    // Third pass: normalize and apply GELU
    for (int i = tid; i < size; i += stride) {
        int c = i / seq_len;
        int s = i % seq_len;
        int weight_idx = c * seq_len + s;
        
        float normalized = (batch_input[i] - mean) * inv_std;
        normalized = normalized * weight[weight_idx] + bias[weight_idx];
        
        // GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
        float x = normalized;
        float x3 = x * x * x;
        float inner = GELU_COEF * (x + 0.044715f * x3);
        float tanh_out = tanhf(inner);
        batch_output[i] = 0.5f * x * (1.0f + tanh_out);
    }
}

// Dilated Conv1d + Residual Add kernel (optimized with shared memory)
__global__ void dilated_conv1d_residual_kernel(
    const float* __restrict__ input,
    const float* __restrict__ residual,
    float* __restrict__ output,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    int batch_size,
    int channels,
    int seq_len,
    int kernel_size,
    int dilation,
    int padding
) {
    const int batch_idx = blockIdx.z;
    const int out_channel = blockIdx.y;
    const int seq_pos = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (batch_idx >= batch_size || out_channel >= channels || seq_pos >= seq_len) return;
    
    float sum = 0.0f;
    
    // Compute convolution for this output position
    for (int in_channel = 0; in_channel < channels; in_channel++) {
        for (int k = 0; k < kernel_size; k++) {
            int input_pos = seq_pos + padding - k * dilation;
            if (input_pos >= 0 && input_pos < seq_len) {
                int weight_idx = out_channel * channels * kernel_size + in_channel * kernel_size + k;
                int input_idx = batch_idx * channels * seq_len + in_channel * seq_len + input_pos;
                sum += input[input_idx] * weight[weight_idx];
            }
        }
    }
    
    // Add bias
    sum += bias[out_channel];
    
    // Add residual
    int residual_idx = batch_idx * channels * seq_len + out_channel * seq_len + seq_pos;
    sum += residual[residual_idx];
    
    // Write output
    int output_idx = batch_idx * channels * seq_len + out_channel * seq_len + seq_pos;
    output[output_idx] = sum;
}

// Optimized dilated Conv1d kernel using shared memory for input channels
__global__ void dilated_conv1d_residual_optimized_kernel(
    const float* __restrict__ input,
    const float* __restrict__ residual,
    float* __restrict__ output,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    int batch_size,
    int channels,
    int seq_len,
    int kernel_size,
    int dilation,
    int padding
) {
    // Each block processes one (batch, out_channel) pair
    const int batch_idx = blockIdx.y;
    const int out_channel = blockIdx.x;
    const int seq_pos = threadIdx.x;
    
    if (batch_idx >= batch_size || out_channel >= channels) return;
    
    // Shared memory for weights of this output channel
    extern __shared__ float shared_mem[];
    float* channel_weights = shared_mem;  // channels * kernel_size floats
    
    // Cooperatively load weights for this output channel
    const int weight_offset = out_channel * channels * kernel_size;
    for (int i = threadIdx.x; i < channels * kernel_size; i += blockDim.x) {
        channel_weights[i] = weight[weight_offset + i];
    }
    __syncthreads();
    
    float sum = 0.0f;
    
    // Process each input channel
    for (int in_channel = 0; in_channel < channels; in_channel++) {
        // Process each kernel position
        for (int k = 0; k < kernel_size; k++) {
            int input_pos = seq_pos + padding - k * dilation;
            if (input_pos >= 0 && input_pos < seq_len) {
                int input_idx = batch_idx * channels * seq_len + in_channel * seq_len + input_pos;
                int weight_idx = in_channel * kernel_size + k;
                sum += input[input_idx] * channel_weights[weight_idx];
            }
        }
    }
    
    // Add bias
    sum += bias[out_channel];
    
    // Add residual
    int residual_idx = batch_idx * channels * seq_len + out_channel * seq_len + seq_pos;
    sum += residual[residual_idx];
    
    // Write output
    int output_idx = batch_idx * channels * seq_len + out_channel * seq_len + seq_pos;
    output[output_idx] = sum;
}

// Host functions
void layernorm_gelu_cuda(
    const torch::Tensor& input,
    torch::Tensor& output,
    const torch::Tensor& weight,
    const torch::Tensor& bias
) {
    int batch_size = input.size(0);
    int channels = input.size(1);
    int seq_len = input.size(2);
    int size = channels * seq_len;
    
    const int BLOCK_SIZE = 512;
    int blocks = batch_size;
    
    layernorm_gelu_kernel<BLOCK_SIZE><<<blocks, BLOCK_SIZE>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        batch_size,
        channels,
        seq_len,
        weight.data_ptr<float>(),
        bias.data_ptr<float>()
    );
}

void dilated_conv1d_residual_cuda(
    const torch::Tensor& input,
    const torch::Tensor& residual,
    torch::Tensor& output,
    const torch::Tensor& weight,
    const torch::Tensor& bias,
    int dilation
) {
    int batch_size = input.size(0);
    int channels = input.size(1);
    int seq_len = input.size(2);
    int kernel_size = weight.size(2);
    int padding = dilation;
    
    const int BLOCK_SIZE = 256;
    dim3 grid((seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE, channels, batch_size);
    
    dilated_conv1d_residual_kernel<<<grid, BLOCK_SIZE>>>(
        input.data_ptr<float>(),
        residual.data_ptr<float>(),
        output.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        batch_size,
        channels,
        seq_len,
        kernel_size,
        dilation,
        padding
    );
}
"""

cpp_source = """
void layernorm_gelu_cuda(
    const torch::Tensor& input,
    torch::Tensor& output,
    const torch::Tensor& weight,
    const torch::Tensor& bias
);

void dilated_conv1d_residual_cuda(
    const torch::Tensor& input,
    const torch::Tensor& residual,
    torch::Tensor& output,
    const torch::Tensor& weight,
    const torch::Tensor& bias,
    int dilation
);
"""

# Load the CUDA extension
module = load_inline(
    name='dilated_conv_extension',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['layernorm_gelu_cuda', 'dilated_conv1d_residual_cuda'],
    extra_cuda_cflags=['-O3', '-use_fast_math'],
    verbose=False
)


class ModelNew(nn.Module):
    """
    Pre-norm 1D dilated convolution block: LayerNorm, GELU activation, dilated Conv1d, and residual addition.
    Optimized with custom CUDA kernels.
    """
    def __init__(self, channels, seq_len, kernel_size, dilation):
        super(ModelNew, self).__init__()
        
        # LayerNorm parameters
        self.channels = channels
        self.seq_len = seq_len
        
        # LayerNorm weights (gamma and beta)
        self.norm_weight = nn.Parameter(torch.ones(channels, seq_len))
        self.norm_bias = nn.Parameter(torch.zeros(channels, seq_len))
        
        # Conv1d weights
        self.conv_weight = nn.Parameter(torch.empty(channels, channels, kernel_size))
        self.conv_bias = nn.Parameter(torch.zeros(channels))
        self.dilation = dilation
        self.kernel_size = kernel_size
        
        # Initialize weights like Conv1d would
        nn.init.kaiming_uniform_(self.conv_weight, a=math.sqrt(5))
        
        # Register the loaded module
        self.cuda_module = module
    
    def forward(self, x):
        residual = x
        
        # Allocate output tensor for LayerNorm + GELU
        x_norm_gelu = torch.empty_like(x)
        
        # Fused LayerNorm + GELU
        self.cuda_module.layernorm_gelu_cuda(
            x, x_norm_gelu, self.norm_weight, self.norm_bias
        )
        
        # Allocate output tensor for Conv1d + Residual
        output = torch.empty_like(x)
        
        # Fused Dilated Conv1d + Residual Add
        self.cuda_module.dilated_conv1d_residual_cuda(
            x_norm_gelu, residual, output, self.conv_weight, self.conv_bias, self.dilation
        )
        
        return output


import math  # Needed for initialization