import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs 1D convolution via FFT for long learnable filters. Computes circular convolution by
    zero-padding to avoid wraparound, performing FFT on both input and filter, multiplying in
    frequency domain, and inverse FFT. This is the core computational primitive in Hyena, H3,
    S4, and other SSM-based architectures for subquadratic sequence mixing.

    Args:
        channels (int): Number of independent channels (each has its own filter).
        seq_len (int): Length of input sequences.
    """
    def __init__(self, channels: int, seq_len: int):
        super(Model, self).__init__()
        self.channels = channels
        self.seq_len = seq_len
        self.fft_size = 2 * seq_len
        self.filter = nn.Parameter(torch.randn(channels, seq_len))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, C, L = x.shape
        x_padded = torch.nn.functional.pad(x, (0, self.seq_len))
        k_padded = torch.nn.functional.pad(self.filter, (0, self.seq_len))
        x_f = torch.fft.rfft(x_padded, dim=-1)
        k_f = torch.fft.rfft(k_padded, dim=-1)
        out_f = x_f * k_f.unsqueeze(0)
        out = torch.fft.irfft(out_f, n=self.fft_size, dim=-1)
        return out[..., :self.seq_len]

batch_size = 16
channels = 256
seq_len = 8192

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

def get_init_inputs():
    return [channels, seq_len]

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

# Define CUDA kernel and C++ wrapper
cuda_source = '''
#include <cuda_runtime.h>
#include <cufft.h>
#include <torch/extension.h>

// CUDA kernel to multiply complex numbers element-wise
__global__ void complex_mul_kernel(cufftComplex* out, const cufftComplex* x, const cufftComplex* k, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        cufftComplex x_val = x[idx];
        cufftComplex k_val = k[idx];
        cufftComplex result;
        result.x = x_val.x * k_val.x - x_val.y * k_val.y;
        result.y = x_val.x * k_val.y + x_val.y * k_val.x;
        out[idx] = result;
    }
}

// C++ wrapper function to perform complex multiplication
torch::Tensor complex_mul(torch::Tensor x, torch::Tensor k) {
    auto out = torch::empty_like(x);
    int size = x.numel();
    cufftComplex* x_data = reinterpret_cast<cufftComplex*>(x.data_ptr<cfloat>());
    cufftComplex* k_data = reinterpret_cast<cufftComplex*>(k.data_ptr<cfloat>());
    cufftComplex* out_data = reinterpret_cast<cufftComplex*>(out.data_ptr<cfloat>());
    
    int blockSize = 256;
    int numBlocks = (size + blockSize - 1) / blockSize;
    complex_mul_kernel<<<numBlocks, blockSize>>>(out_data, x_data, k_data, size);
    return out;
}
'''

# Compile CUDA kernel
complex_mul_extension = load_inline(
    name='complex_mul_extension',
    cuda_sources=cuda_source,
    functions=['complex_mul'],
    extra_cuda_cflags=['--use_fast_math'],
    extra_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self, channels: int, seq_len: int):
        super(ModelNew, self).__init__()
        self.channels = channels
        self.seq_len = seq_len
        self.fft_size = 2 * seq_len
        self.filter = nn.Parameter(torch.randn(channels, seq_len))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, C, L = x.shape
        x_padded = torch.nn.functional.pad(x, (0, self.seq_len))
        k_padded = torch.nn.functional.pad(self.filter, (0, self.seq_len)).unsqueeze(0).repeat(B, 1, 1)
        x_f = torch.fft.rfft(x_padded, dim=-1)
        k_f = torch.fft.rfft(k_padded, dim=-1)
        out_f = complex_mul_extension.complex_mul(x_f, k_f)
        out = torch.fft.irfft(out_f, n=self.fft_size, dim=-1)
        return out[..., :self.seq_len]