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

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

// Complex multiplication kernel
__global__ void complex_multiply_kernel(const cufftComplex* __restrict__ x_f, 
                                        const cufftComplex* __restrict__ k_f, 
                                        cufftComplex* __restrict__ out_f, 
                                        int B, int C, int rfft_size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * C * rfft_size;
    
    if (idx < total) {
        int f = idx % rfft_size;
        int remainder = idx / rfft_size;
        int c = remainder % C;
        
        int k_idx = c * rfft_size + f;
        
        cufftComplex x_val = x_f[idx];
        cufftComplex k_val = k_f[k_idx];
        
        out_f[idx].x = x_val.x * k_val.x - x_val.y * k_val.y;
        out_f[idx].y = x_val.x * k_val.y + x_val.y * k_val.x;
    }
}

// Pad input to 2*L (properly laid out for batched FFT)
__global__ void pad_input_kernel(const float* __restrict__ input, float* __restrict__ padded, 
                                  int B, int C, int L) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total_elements = B * C * L;
    
    if (idx < total_elements) {
        int l = idx % L;
        int bc = idx / L;
        int c = bc % C;
        int b = bc / C;
        
        int padded_idx = b * C * 2 * L + c * 2 * L + l;
        padded[padded_idx] = input[idx];
        padded[padded_idx + L] = 0.0f;
    }
}

// Pad filter to 2*L
__global__ void pad_filter_kernel(const float* __restrict__ filter, float* __restrict__ padded, 
                                   int C, int L) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total_elements = C * L;
    
    if (idx < total_elements) {
        int l = idx % L;
        int c = idx / L;
        
        int padded_idx = c * 2 * L + l;
        padded[padded_idx] = filter[idx];
        padded[padded_idx + L] = 0.0f;
    }
}

// Extract result (first L elements from IRFFT output) and normalize
__global__ void extract_and_normalize_kernel(const float* __restrict__ irfft_out, 
                                              float* __restrict__ result, 
                                              int B, int C, int L, float scale) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total_elements = B * C * L;
    
    if (idx < total_elements) {
        int l = idx % L;
        int bc = idx / L;
        int c = bc % C;
        int b = bc / C;
        
        int irfft_idx = b * C * 2 * L + c * 2 * L + l;
        result[idx] = irfft_out[irfft_idx] * scale;
    }
}

torch::Tensor fft_conv_forward(torch::Tensor input, torch::Tensor filter) {
    int B = input.size(0);
    int C = input.size(1);
    int L = input.size(2);
    int fft_size = 2 * L;
    int rfft_size = L + 1;
    
    auto output = torch::empty({B, C, L}, input.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    float* d_x_padded;
    float* d_k_padded;
    cufftComplex* d_x_f;
    cufftComplex* d_k_f;
    cufftComplex* d_out_f;
    float* d_irfft_out;
    
    size_t input_padded_size = B * C * fft_size * sizeof(float);
    size_t filter_padded_size = C * fft_size * sizeof(float);
    size_t input_complex_size = B * C * rfft_size * sizeof(cufftComplex);
    size_t filter_complex_size = C * rfft_size * sizeof(cufftComplex);
    
    cudaMalloc(&d_x_padded, input_padded_size);
    cudaMalloc(&d_k_padded, filter_padded_size);
    cudaMalloc(&d_x_f, input_complex_size);
    cudaMalloc(&d_k_f, filter_complex_size);
    cudaMalloc(&d_out_f, input_complex_size);
    cudaMalloc(&d_irfft_out, input_padded_size);
    
    cufftHandle plan_rfft_x, plan_rfft_k, plan_irfft;
    
    int n = fft_size;
    int istride = 1, ostride = 1;
    int idist = fft_size, odist = rfft_size;
    
    cufftPlanMany(&plan_rfft_x, 1, &n, &n, istride, idist, &n, ostride, odist, CUFFT_R2C, B * C);
    cufftPlanMany(&plan_rfft_k, 1, &n, &n, istride, idist, &n, ostride, odist, CUFFT_R2C, C);
    cufftPlanMany(&plan_irfft, 1, &n, &n, ostride, odist, &n, istride, idist, CUFFT_C2R, B * C);
    
    cufftSetStream(plan_rfft_x, stream);
    cufftSetStream(plan_rfft_k, stream);
    cufftSetStream(plan_irfft, stream);
    
    int threads = 256;
    int total_input = B * C * L;
    int blocks = (total_input + threads - 1) / threads;
    pad_input_kernel<<<blocks, threads, 0, stream>>>(input.data_ptr<float>(), d_x_padded, B, C, L);
    
    int total_filter = C * L;
    blocks = (total_filter + threads - 1) / threads;
    pad_filter_kernel<<<blocks, threads, 0, stream>>>(filter.data_ptr<float>(), d_k_padded, C, L);
    
    cufftExecR2C(plan_rfft_x, (cufftReal*)d_x_padded, d_x_f);
    cufftExecR2C(plan_rfft_k, (cufftReal*)d_k_padded, d_k_f);
    
    int total_complex = B * C * rfft_size;
    blocks = (total_complex + threads - 1) / threads;
    complex_multiply_kernel<<<blocks, threads, 0, stream>>>(d_x_f, d_k_f, d_out_f, B, C, rfft_size);
    
    cufftExecC2R(plan_irfft, d_out_f, (cufftReal*)d_irfft_out);
    
    float scale = 1.0f / fft_size;
    blocks = (total_input + threads - 1) / threads;
    extract_and_normalize_kernel<<<blocks, threads, 0, stream>>>(d_irfft_out, output.data_ptr<float>(), B, C, L, scale);
    
    cufftDestroy(plan_rfft_x);
    cufftDestroy(plan_rfft_k);
    cufftDestroy(plan_irfft);
    
    cudaFree(d_x_padded);
    cudaFree(d_k_padded);
    cudaFree(d_x_f);
    cudaFree(d_k_f);
    cudaFree(d_out_f);
    cudaFree(d_irfft_out);
    
    return output;
}
'''

cpp_source = "torch::Tensor fft_conv_forward(torch::Tensor input, torch::Tensor filter);"

fft_module = load_inline(
    name='fft_conv',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['fft_conv_forward'],
    extra_cuda_cflags=['-O3'],
    verbose=False
)

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:
        return fft_module.fft_conv_forward(x, self.filter)