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

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

torch::Tensor ln_gelu_blc(torch::Tensor x_blc, torch::Tensor gamma_cl, torch::Tensor beta_cl);
'''

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

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

__global__ void ln_gelu_blc_kernel(const float* __restrict__ x,
                                  const float* __restrict__ gamma,
                                  const float* __restrict__ beta,
                                  float* __restrict__ out,
                                  int B, int L, int C) {
    extern __shared__ float smem[];
    float* ssum = smem;
    float* ssq = smem + blockDim.x;
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int N = L * C;
    const float* xb = x + ((long long)b) * N;

    float sum = 0.0f;
    float sq = 0.0f;
    for (int idx = tid; idx < N; idx += blockDim.x) {
        float v = xb[idx];
        sum += v;
        sq += v * v;
    }
    ssum[tid] = sum;
    ssq[tid] = sq;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            ssum[tid] += ssum[tid + stride];
            ssq[tid] += ssq[tid + stride];
        }
        __syncthreads();
    }

    float mean = ssum[0] / (float)N;
    float var = ssq[0] / (float)N - mean * mean;
    var = fmaxf(var, 0.0f);
    float inv = rsqrtf(var + 1.0e-5f);

    const float kAlpha = 0.7978845608028654f; // sqrt(2/pi)
    const float kBeta = 0.044715f;
    for (int idx = tid; idx < N; idx += blockDim.x) {
        int l = idx / C;
        int c = idx - l * C;
        // gamma/beta are stored as [C,L] in the PyTorch LayerNorm module.
        float y = (xb[idx] - mean) * inv;
        y = y * gamma[(long long)c * L + l] + beta[(long long)c * L + l];
        float u = kAlpha * (y + kBeta * y * y * y);
        float gelu = 0.5f * y * (1.0f + tanhf(u));
        out[((long long)b) * N + idx] = gelu;
    }
}

torch::Tensor ln_gelu_blc(torch::Tensor x_blc, torch::Tensor gamma_cl, torch::Tensor beta_cl) {
    CHECK_CUDA(x_blc);
    CHECK_CUDA(gamma_cl);
    CHECK_CUDA(beta_cl);
    TORCH_CHECK(x_blc.scalar_type() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(gamma_cl.scalar_type() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(beta_cl.scalar_type() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(x_blc.is_contiguous(), "x_blc must be contiguous");
    TORCH_CHECK(gamma_cl.is_contiguous(), "gamma must be contiguous");
    TORCH_CHECK(beta_cl.is_contiguous(), "beta must be contiguous");
    int B = (int)x_blc.size(0);
    int L = (int)x_blc.size(1);
    int C = (int)x_blc.size(2);
    auto out = torch::empty_like(x_blc);
    int threads = 256;
    size_t shmem = threads * 2 * sizeof(float);
    ln_gelu_blc_kernel<<<B, threads, shmem, at::cuda::getCurrentCUDAStream()>>>(
        x_blc.data_ptr<float>(), gamma_cl.data_ptr<float>(), beta_cl.data_ptr<float>(),
        out.data_ptr<float>(), B, L, C);
    return out;
}
'''

_ext = load_inline(
    name='dilated_ln_gelu_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['ln_gelu_blc'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    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.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):
        # Work internally in [B,L,C] so each dilated-convolution tap is a large GEMM.
        # This is mathematically the same as Conv1d with kernel size 3 and padding=dilation.
        B = x.shape[0]
        C = self.channels
        L = self.seq_len
        xb = x.permute(0, 2, 1).contiguous()
        residual = xb
        for i in range(self.max_dilation_power + 1):
            norm = self.norms[i]
            conv = self.convs[i]
            z = _ext.ln_gelu_blc(xb, norm.weight, norm.bias)
            w = conv.weight
            # center tap plus bias
            out = torch.matmul(z.reshape(B * L, C), w[:, :, 1].t()).reshape(B, L, C)
            out = out + conv.bias.view(1, 1, C)
            d = 1 << i
            # left tap: output position l receives input l-d
            if d < L:
                left = torch.matmul(z[:, :L-d, :].reshape(B * (L-d), C), w[:, :, 0].t()).reshape(B, L-d, C)
                out[:, d:, :].add_(left)
                right = torch.matmul(z[:, d:, :].reshape(B * (L-d), C), w[:, :, 2].t()).reshape(B, L-d, C)
                out[:, :L-d, :].add_(right)
            xb = out + residual
        return xb.permute(0, 2, 1).contiguous()