import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Dilated convolution tower with exponentially increasing dilation rates and residual connections.
    Each block applies LayerNorm -> GELU -> DilatedConv1d, then adds back the original input as a residual.
    This architecture is a core building block of genomic foundation models (Enformer, Basenji, Caduceus)
    for capturing multi-scale sequence dependencies in DNA/RNA.
    """
    def __init__(self, channels, seq_len, kernel_size, max_dilation_power):
        super(Model, self).__init__()
        self.norms = nn.ModuleList()
        self.convs = nn.ModuleList()
        self.activation = nn.GELU()
        for i in range(max_dilation_power + 1):
            dilation = 2 ** i
            self.norms.append(nn.LayerNorm([channels, seq_len]))
            self.convs.append(nn.Conv1d(channels, channels, kernel_size, padding=dilation, dilation=dilation))

    def forward(self, x):
        residual = x
        for norm, conv in zip(self.norms, self.convs):
            x = norm(x)
            x = self.activation(x)
            x = conv(x)
            x = x + residual
        return x

batch_size = 32
channels = 256
seq_len = 4096
kernel_size = 3
max_dilation_power = 8

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, max_dilation_power]

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

# Define CUDA kernels
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <cmath>

// Fused LayerNorm + GELU kernel
// Normalizes over the entire (channels, seq_len) for each batch sample
template <typename scalar_t>
__global__ void layernorm_gelu_kernel(
    const scalar_t* __restrict__ input,
    scalar_t* __restrict__ output,
    const int channels,
    const int seq_len) {
    
    const int batch_idx = blockIdx.x;
    const int total_elements = channels * seq_len;
    const int tid = threadIdx.x;
    const int block_size = blockDim.x;
    
    // Shared memory for reduction
    extern __shared__ char smem[];
    scalar_t* s_sum = reinterpret_cast<scalar_t*>(smem);
    scalar_t* s_sum_sq = s_sum + block_size;
    
    // Compute mean
    scalar_t local_sum = scalar_t(0.0);
    for (int i = tid; i < total_elements; i += block_size) {
        int idx = batch_idx * total_elements + i;
        local_sum += input[idx];
    }
    s_sum[tid] = local_sum;
    __syncthreads();
    
    for (int stride = block_size >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            s_sum[tid] += s_sum[tid + stride];
        }
        __syncthreads();
    }
    
    const scalar_t mean = s_sum[0] / total_elements;
    
    // Compute variance
    scalar_t local_sum_sq = scalar_t(0.0);
    for (int i = tid; i < total_elements; i += block_size) {
        int idx = batch_idx * total_elements + i;
        scalar_t diff = input[idx] - mean;
        local_sum_sq += diff * diff;
    }
    s_sum_sq[tid] = local_sum_sq;
    __syncthreads();
    
    for (int stride = block_size >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            s_sum_sq[tid] += s_sum_sq[tid + stride];
        }
        __syncthreads();
    }
    
    const scalar_t var = s_sum_sq[0] / total_elements;
    const scalar_t inv_std = rsqrt(var + scalar_t(1e-5));
    
    // Normalize and apply GELU
    for (int i = tid; i < total_elements; i += block_size) {
        int idx = batch_idx * total_elements + i;
        scalar_t val = (input[idx] - mean) * inv_std;
        // GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
        const scalar_t sqrt_2_over_pi = scalar_t(0.7978845608028654);
        const scalar_t coeff = scalar_t(0.044715);
        scalar_t gelu = scalar_t(0.5) * val * (scalar_t(1.0) + tanh(sqrt_2_over_pi * (val + coeff * val * val * val)));
        output[idx] = gelu;
    }
}

// Optimized dilated Conv1d + Residual Add kernel
// Each thread handles one output element (batch, out_channel, seq_pos)
template <typename scalar_t>
__global__ void dilated_conv1d_residual_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ weight,
    const scalar_t* __restrict__ residual,
    scalar_t* __restrict__ output,
    const int channels,
    const int seq_len,
    const int kernel_size,
    const int dilation,
    const int padding) {
    
    const int batch_idx = blockIdx.x;
    const int out_c = blockIdx.y;
    const int out_s = blockIdx.z * blockDim.x + threadIdx.x;
    
    if (out_s >= seq_len) return;
    
    const int total_elements = channels * seq_len;
    scalar_t sum = scalar_t(0.0);
    
    // Conv1d: each output channel connects to all input channels
    // Weight shape: [out_channels, in_channels, kernel_size]
    for (int in_c = 0; in_c < channels; in_c++) {
        for (int k = 0; k < kernel_size; k++) {
            const int in_s = out_s - padding + k * dilation;
            if (in_s >= 0 && in_s < seq_len) {
                const int in_idx = batch_idx * total_elements + in_c * seq_len + in_s;
                const int w_idx = out_c * channels * kernel_size + in_c * kernel_size + k;
                sum += input[in_idx] * weight[w_idx];
            }
        }
    }
    
    const int out_idx = batch_idx * total_elements + out_c * seq_len + out_s;
    output[out_idx] = sum + residual[out_idx];
}

// Optimized fused kernel: LayerNorm + GELU + Conv1d + Residual
// Processes one block at a time
template <typename scalar_t>
__global__ void fused_block_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ weight,
    scalar_t* __restrict__ output,
    const int channels,
    const int seq_len,
    const int kernel_size,
    const int dilation,
    const int padding) {
    
    const int batch_idx = blockIdx.x;
    const int out_c = blockIdx.y;
    const int out_s = blockIdx.z * blockDim.x + threadIdx.x;
    
    if (out_s >= seq_len) return;
    
    const int total_elements = channels * seq_len;
    
    // Load input values for this thread's output position and needed neighbors
    // Also load values for computing layer norm statistics
    // For now, we'll use a two-stage approach within the kernel
    
    // Stage 1: Each thread contributes to normalization stats
    // This is done by the first thread block ( blockIdx.z == 0 )
    __shared__ scalar_t s_mean[1];
    __shared__ scalar_t s_inv_std[1];
    
    // For simplicity, we'll use cooperative groups or separate launches
    // Here we'll just do the straightforward approach
}

// Host functions
void layernorm_gelu_cuda(
    torch::Tensor input,
    torch::Tensor output) {
    
    const int batch_size = input.size(0);
    const int channels = input.size(1);
    const int seq_len = input.size(2);
    const int total_elements = channels * seq_len;
    
    const int threads = 256;
    const int blocks = batch_size;
    const int shared_mem = 2 * threads * sizeof(float);
    
    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "layernorm_gelu_cuda", ([&] {
        layernorm_gelu_kernel<scalar_t><<<blocks, threads, shared_mem>>>(
            input.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            channels,
            seq_len);
    }));
}

void dilated_conv1d_residual_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor residual,
    torch::Tensor output,
    int dilation) {
    
    const int batch_size = input.size(0);
    const int channels = input.size(1);
    const int seq_len = input.size(2);
    const int kernel_size = weight.size(2);
    const int padding = dilation;
    
    const int threads = 256;
    const dim3 blocks(batch_size, channels, (seq_len + threads - 1) / threads);
    
    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "dilated_conv1d_residual_cuda", ([&] {
        dilated_conv1d_residual_kernel<scalar_t><<<blocks, threads>>>(
            input.data_ptr<scalar_t>(),
            weight.data_ptr<scalar_t>(),
            residual.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            channels,
            seq_len,
            kernel_size,
            dilation,
            padding);
    }));
}

// Multi-block fused operation - applies all blocks in sequence
void forward_cuda(
    torch::Tensor x,
    std::vector<torch::Tensor> weights,
    int num_blocks) {
    
    const int batch_size = x.size(0);
    const int channels = x.size(1);
    const int seq_len = x.size(2);
    
    // Allocate temporary buffer
    auto temp = torch::empty_like(x);
    torch::Tensor* current = &x;
    torch::Tensor* next = &temp;
    
    const int threads = 256;
    const int shared_mem = 2 * threads * sizeof(float);
    
    for (int block_idx = 0; block_idx < num_blocks; block_idx++) {
        int dilation = 1 << block_idx;
        int kernel_size = weights[block_idx].size(2);
        int padding = dilation;
        
        // Save residual
        auto residual = *current;
        
        // LayerNorm + GELU
        AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "layernorm_gelu_cuda", ([&] {
            layernorm_gelu_kernel<scalar_t><<<batch_size, threads, shared_mem>>>(
                current->data_ptr<scalar_t>(),
                next->data_ptr<scalar_t>(),
                channels,
                seq_len);
        }));
        
        // Dilated Conv1d + Residual
        const dim3 blocks(batch_size, channels, (seq_len + threads - 1) / threads);
        AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "dilated_conv1d_residual_cuda", ([&] {
            dilated_conv1d_residual_kernel<scalar_t><<<blocks, threads>>>(
                next->data_ptr<scalar_t>(),
                weights[block_idx].data_ptr<scalar_t>(),
                residual.data_ptr<scalar_t>(),
                current->data_ptr<scalar_t>(),
                channels,
                seq_len,
                kernel_size,
                dilation,
                padding);
        }));
    }
}

torch::Tensor forward_multi_block_cuda(
    torch::Tensor x,
    torch::Tensor weights_flat,
    int num_blocks,
    int kernel_size) {
    
    const int batch_size = x.size(0);
    const int channels = x.size(1);
    const int seq_len = x.size(2);
    
    // Allocate temporary buffer
    auto temp = torch::empty_like(x);
    
    const int threads = 256;
    const int shared_mem = 2 * threads * sizeof(float);
    const dim3 conv_blocks(batch_size, channels, (seq_len + threads - 1) / threads);
    
    torch::Tensor current = x;
    
    for (int block_idx = 0; block_idx < num_blocks; block_idx++) {
        int dilation = 1 << block_idx;
        int padding = dilation;
        
        // Save residual (need to clone since we'll modify the tensor)
        auto residual = current.clone();
        
        // LayerNorm + GELU
        AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "layernorm_gelu_cuda", ([&] {
            layernorm_gelu_kernel<scalar_t><<<batch_size, threads, shared_mem>>>(
                current.data_ptr<scalar_t>(),
                temp.data_ptr<scalar_t>(),
                channels,
                seq_len);
        }));
        
        // Dilated Conv1d + Residual
        // Weight offset for this block
        int weight_offset = block_idx * channels * channels * kernel_size;
        auto weight_slice = weights_flat.slice(0, weight_offset, weight_offset + channels * channels * kernel_size);
        
        AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "dilated_conv1d_residual_cuda", ([&] {
            dilated_conv1d_residual_kernel<scalar_t><<<conv_blocks, threads>>>(
                temp.data_ptr<scalar_t>(),
                weight_slice.data_ptr<scalar_t>(),
                residual.data_ptr<scalar_t>(),
                current.data_ptr<scalar_t>(),
                channels,
                seq_len,
                kernel_size,
                dilation,
                padding);
        }));
    }
    
    return current;
}
"""

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


class ModelNew(nn.Module):
    """
    Optimized implementation of dilated convolution tower with residual connections.
    Uses custom CUDA kernels for LayerNorm+GELU fusion and dilated Conv1d+Residual fusion.
    """
    def __init__(self, channels, seq_len, kernel_size, max_dilation_power):
        super(ModelNew, self).__init__()
        self.channels = channels
        self.seq_len = seq_len
        self.kernel_size = kernel_size
        self.max_dilation_power = max_dilation_power
        self.num_blocks = max_dilation_power + 1
        
        # Initialize weights as a single flat tensor for efficient memory access
        # Each block has shape [channels, channels, kernel_size]
        self.weights = nn.Parameter(torch.randn(self.num_blocks, channels, channels, kernel_size) * 0.02)
        
        # Create layer norm parameters (weight=1, bias=0 as in default LayerNorm)
        # We'll use default LayerNorm behavior (no learnable params in this implementation)

    def forward(self, x):
        # Use custom CUDA kernels for the forward pass
        batch_size = x.size(0)
        
        # Clone input to avoid modifying the original tensor
        current = x.clone()
        temp = torch.empty_like(current)
        residual = torch.empty_like(current)
        
        threads = 256
        shared_mem = 2 * threads * 4  # 2 arrays of float
        
        for block_idx in range(self.num_blocks):
            dilation = 1 << block_idx
            padding = dilation
            
            # Save residual
            residual.copy_(current)
            
            # LayerNorm + GELU
            dilated_conv_ext.layernorm_gelu_cuda(current, temp)
            
            # Dilated Conv1d + Residual
            weight = self.weights[block_idx]
            dilated_conv_ext.dilated_conv1d_residual_cuda(temp, weight, residual, current, dilation)
        
        return current