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

cpp_sources = r'''
#include <torch/extension.h>
torch::Tensor fft_conv_forward(torch::Tensor x, torch::Tensor k_f);
'''

cuda_sources = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cufft.h>
#include <c10/util/complex.h>

static cufftHandle plan_r2c = 0;
static cufftHandle plan_c2r = 0;
static int cached_N = 0;
static int cached_batch = 0;

static void ensure_plans(int N, int batch){
    if(plan_r2c != 0 && cached_N == N && cached_batch == batch) return;
    if(plan_r2c != 0){
        cufftDestroy(plan_r2c);
        cufftDestroy(plan_c2r);
        plan_r2c = 0;
        plan_c2r = 0;
    }
    int n[1] = {N};
    int inembed[1] = {N};
    int onembed[1] = {N/2 + 1};
    cufftResult r = cufftPlanMany(&plan_r2c, 1, n, inembed, 1, N, onembed, 1, N/2 + 1, CUFFT_R2C, batch);
    TORCH_CHECK(r == CUFFT_SUCCESS, "cufftPlanMany R2C failed");
    r = cufftPlanMany(&plan_c2r, 1, n, onembed, 1, N/2 + 1, inembed, 1, N, CUFFT_C2R, batch);
    TORCH_CHECK(r == CUFFT_SUCCESS, "cufftPlanMany C2R failed");
    cached_N = N;
    cached_batch = batch;
}

__global__ void pad_kernel(const float* __restrict__ x, float* __restrict__ xp, int L, int N){
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    int t = idx % N;
    long b = idx / N;
    xp[idx] = (t < L) ? x[b * (long)L + t] : 0.0f;
}

__global__ void mul_kernel(float2* __restrict__ xf, const float2* __restrict__ kf, int C, int F, long total){
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if(idx >= total) return;
    int f = idx % F;
    int c = (idx / F) % C;
    float2 a = xf[idx];
    float2 b = kf[(long)c * F + f];
    float2 r;
    r.x = a.x * b.x - a.y * b.y;
    r.y = a.x * b.y + a.y * b.x;
    xf[idx] = r;
}

__global__ void scale_extract_kernel(const float* __restrict__ full, float* __restrict__ out, int L, int N, long total){
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if(idx >= total) return;
    int t = idx % L;
    long b = idx / L;
    out[idx] = full[b * (long)N + t] * (1.0f / (float)N);
}

torch::Tensor fft_conv_forward(torch::Tensor x, torch::Tensor k_f){
    int B = x.size(0);
    int C = x.size(1);
    int L = x.size(2);
    int N = 2 * L;
    int F = N / 2 + 1;
    int batch = B * C;
    ensure_plans(N, batch);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cufftSetStream(plan_r2c, stream);
    cufftSetStream(plan_c2r, stream);

    auto x_pad = torch::empty({batch, N}, x.options());
    auto x_freq = torch::empty({batch, F}, x.options().dtype(torch::kComplexFloat));
    auto full = torch::empty({batch, N}, x.options());
    int th = 256;
    long total_pad = (long)batch * N;
    pad_kernel<<<(total_pad + th - 1) / th, th, 0, stream>>>(x.data_ptr<float>(), x_pad.data_ptr<float>(), L, N);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    cufftResult r = cufftExecR2C(plan_r2c, x_pad.data_ptr<float>(), reinterpret_cast<cufftComplex*>(x_freq.data_ptr<c10::complex<float>>()));
    TORCH_CHECK(r == CUFFT_SUCCESS, "cufftExecR2C failed");

    long total_f = (long)batch * F;
    mul_kernel<<<(total_f + th - 1) / th, th, 0, stream>>>(reinterpret_cast<float2*>(x_freq.data_ptr<c10::complex<float>>()), reinterpret_cast<const float2*>(k_f.data_ptr<c10::complex<float>>()), C, F, total_f);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    r = cufftExecC2R(plan_c2r, reinterpret_cast<cufftComplex*>(x_freq.data_ptr<c10::complex<float>>()), full.data_ptr<float>());
    TORCH_CHECK(r == CUFFT_SUCCESS, "cufftExecC2R failed");

    auto out = torch::empty({B, C, L}, x.options());
    long total_out = (long)batch * L;
    scale_extract_kernel<<<(total_out + th - 1) / th, th, 0, stream>>>(full.data_ptr<float>(), out.data_ptr<float>(), L, N, total_out);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

_ext = load_inline(
    name='cufft_conv_ext_final',
    cpp_sources=cpp_sources,
    cuda_sources=cuda_sources,
    functions=['fft_conv_forward'],
    extra_cuda_cflags=['-O3'],
    extra_ldflags=['-lcufft'],
    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))
        self._k_f = None
        self._k_version = -1
        self._k_device = None
        self._k_dtype = None

    def _get_kf(self):
        if (self._k_f is None or
            self._k_version != self.filter._version or
            self._k_device != self.filter.device or
            self._k_dtype != self.filter.dtype):
            self._k_f = torch.fft.rfft(self.filter, n=self.fft_size, dim=-1).contiguous()
            self._k_version = self.filter._version
            self._k_device = self.filter.device
            self._k_dtype = self.filter.dtype
        return self._k_f

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.fft_conv_forward(x.contiguous(), self._get_kf())