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_src = r'''
#include <torch/extension.h>

torch::Tensor complex_mul_cuda(torch::Tensor xf, torch::Tensor kf);

torch::Tensor complex_mul(torch::Tensor xf, torch::Tensor kf) {
    return complex_mul_cuda(xf, kf);
}
'''

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

__global__ void complex_mul_kernel(const float2* __restrict__ x,
                                   const float2* __restrict__ k,
                                   float2* __restrict__ y,
                                   int F) {
    int f = blockIdx.x * blockDim.x + threadIdx.x;
    int row = blockIdx.y;
    if (f < F) {
        long idx = (long)row * F + f;
        float2 a = x[idx];
        float2 b = k[((row & 255) * (long)F) + f];
        float2 o;
        o.x = a.x * b.x - a.y * b.y;
        o.y = a.x * b.y + a.y * b.x;
        y[idx] = o;
    }
}

torch::Tensor complex_mul_cuda(torch::Tensor xf, torch::Tensor kf) {
    auto y = torch::empty_like(xf);
    const int F = (int)xf.size(2);
    const int rows = (int)(xf.numel() / F);
    const int threads = 256;
    dim3 blocks((F + threads - 1) / threads, rows);
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
    complex_mul_kernel<<<blocks, threads, 0, stream>>>(
        reinterpret_cast<const float2*>(xf.data_ptr<c10::complex<float>>()),
        reinterpret_cast<const float2*>(kf.data_ptr<c10::complex<float>>()),
        reinterpret_cast<float2*>(y.data_ptr<c10::complex<float>>()), F);
    return y;
}
'''

_ext = load_inline(name='fft_conv_mul_ext_v2',
                   cpp_sources=[_cpp_src],
                   cuda_sources=[_cuda_src],
                   functions=['complex_mul'],
                   extra_cflags=['-O3'],
                   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))
        self._cached_kf = None
        self._cached_version = -1

    def _get_kf(self):
        v = self.filter._version
        if self._cached_kf is None or self._cached_version != v:
            self._cached_kf = torch.fft.rfft(self.filter, n=self.fft_size, dim=-1).contiguous()
            self._cached_version = v
        return self._cached_kf

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        kf = self._get_kf()
        xf = torch.fft.rfft(x, n=self.fft_size, dim=-1)
        yf = _ext.complex_mul(xf, kf)
        y = torch.fft.irfft(yf, n=self.fft_size, dim=-1)
        return y[..., :self.seq_len]