import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    """
    Full Hyena operator (order 2): projects input into value (v), query-like (q), and key-like (k)
    branches, applies short depthwise convolutions to q and k, generates a long convolution filter
    via a small MLP, convolves v with the filter via FFT, then applies two rounds of elementwise
    gating. This is the core sequence mixing operator in StripedHyena, HyenaDNA, and Evo
    for subquadratic long-range dependency modeling.

    Args:
        d_model (int): Model dimension.
        seq_len (int): Maximum sequence length.
        order (int): Hyena recurrence order (number of gating rounds, typically 2).
        short_filter_size (int): Kernel size for short depthwise convolutions.
    """
    def __init__(self, d_model: int, seq_len: int, order: int, short_filter_size: int):
        super(Model, self).__init__()
        self.d_model = d_model
        self.seq_len = seq_len
        self.order = order
        self.fft_size = 2 * seq_len

        self.in_proj = nn.Linear(d_model, d_model * (order + 1), bias=True)

        self.short_convs = nn.ModuleList([
            nn.Conv1d(d_model, d_model, short_filter_size,
                      padding=short_filter_size // 2, groups=d_model)
            for _ in range(order)
        ])

        self.filter_mlp = nn.Sequential(
            nn.Linear(1, 64),
            nn.SiLU(),
            nn.Linear(64, d_model),
        )
        t = torch.linspace(0, 1, seq_len).unsqueeze(-1)
        self.register_buffer('filter_positions', t)

        self.out_proj = nn.Linear(d_model, d_model, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, L, D = x.shape

        projected = self.in_proj(x)
        chunks = projected.chunk(self.order + 1, dim=-1)
        v = chunks[0].transpose(1, 2)
        gates = [c.transpose(1, 2) for c in chunks[1:]]

        for i in range(self.order):
            gates[i] = self.short_convs[i](gates[i])

        h = self.filter_mlp(self.filter_positions).transpose(0, 1)

        v_padded = F.pad(v, (0, self.seq_len))
        h_padded = F.pad(h, (0, self.seq_len))
        v_f = torch.fft.rfft(v_padded, dim=-1)
        h_f = torch.fft.rfft(h_padded, dim=-1)
        y_f = v_f * h_f.unsqueeze(0)
        y = torch.fft.irfft(y_f, n=self.fft_size, dim=-1)[..., :L]

        for gate in gates:
            y = y * gate

        y = y.transpose(1, 2)
        y = self.out_proj(y)
        return y

batch_size = 16
d_model = 256
seq_len = 8192
order = 2
short_filter_size = 3

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

def get_init_inputs():
    return [d_model, seq_len, order, short_filter_size]

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

cpp_src = r'''
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> prepare_v_gate(torch::Tensor projected, torch::Tensor w0, torch::Tensor b0, torch::Tensor w1, torch::Tensor b1, int64_t seq_len);
torch::Tensor prepare_h(torch::Tensor fw1, torch::Tensor fb1, torch::Tensor fw2, torch::Tensor fb2, int64_t seq_len);
void mul_freq_inplace(torch::Tensor vf, torch::Tensor hf);
torch::Tensor apply_gate_transpose(torch::Tensor y, torch::Tensor gate, int64_t L);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <cmath>
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

__global__ void prep_v_gate_kernel(const float* __restrict__ proj,
                                   const float* __restrict__ cw0, const float* __restrict__ cb0,
                                   const float* __restrict__ cw1, const float* __restrict__ cb1,
                                   float* __restrict__ vpad, float* __restrict__ gate,
                                   int B, int L, int D, int K, int N) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)B * L * D;
    if (idx >= total) return;
    int d = idx % D;
    long long q = idx / D;
    int t = q % L;
    int b = q / L;
    int threeD = 3 * D;
    const float* base = proj + ((long long)b * L + t) * threeD;
    vpad[((long long)b * D + d) * N + t] = base[d];
    int pad = K >> 1;
    float g0 = cb0[d];
    float g1 = cb1[d];
    #pragma unroll
    for (int r = 0; r < 8; ++r) {
        if (r < K) {
            int tt = t + r - pad;
            if ((unsigned)tt < (unsigned)L) {
                const float* p = proj + ((long long)b * L + tt) * threeD;
                g0 += cw0[d * K + r] * p[D + d];
                g1 += cw1[d * K + r] * p[2 * D + d];
            }
        }
    }
    gate[((long long)b * L + t) * D + d] = g0 * g1;
}

__global__ void zero_second_half_kernel(float* __restrict__ vpad, int B, int D, int L, int N) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)B * D * L;
    if (idx >= total) return;
    int t = idx % L;
    long long bd = idx / L;
    vpad[bd * N + L + t] = 0.0f;
}

__global__ void prep_h_kernel(const float* __restrict__ fw1, const float* __restrict__ fb1,
                              const float* __restrict__ fw2, const float* __restrict__ fb2,
                              float* __restrict__ hpad, int D, int L, int N) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)D * N;
    if (idx >= total) return;
    int t = idx % N;
    int d = idx / N;
    if (t >= L) { hpad[(long long)d * N + t] = 0.0f; return; }
    float x = (L > 1) ? ((float)t / (float)(L - 1)) : 0.0f;
    float acc = fb2[d];
    #pragma unroll
    for (int j = 0; j < 64; ++j) {
        float z = x * fw1[j] + fb1[j];
        float a = z / (1.0f + expf(-z));
        acc += a * fw2[d * 64 + j];
    }
    hpad[(long long)d * N + t] = acc;
}

__global__ void cmul_broadcast_inplace_kernel(float2* __restrict__ vf, const float2* __restrict__ hf, int B, int D, int F) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)B * D * F;
    if (idx >= total) return;
    int f = idx % F;
    int d = (idx / F) % D;
    float2 a = vf[idx];
    float2 b = hf[(long long)d * F + f];
    vf[idx] = make_float2(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
}

__global__ void apply_gate_transpose_kernel(const float* __restrict__ y, const float* __restrict__ gate,
                                            float* __restrict__ out, int B, int L, int D, int N) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)B * L * D;
    if (idx >= total) return;
    int d = idx % D;
    long long q = idx / D;
    int t = q % L;
    int b = q / L;
    out[idx] = y[((long long)b * D + d) * N + t] * gate[((long long)b * L + t) * D + d];
}

std::vector<torch::Tensor> prepare_v_gate(torch::Tensor projected, torch::Tensor w0, torch::Tensor b0, torch::Tensor w1, torch::Tensor b1, int64_t seq_len) {
    CHECK_INPUT(projected); CHECK_INPUT(w0); CHECK_INPUT(b0); CHECK_INPUT(w1); CHECK_INPUT(b1);
    int B = projected.size(0), L = projected.size(1), D = projected.size(2) / 3, K = w0.size(2), N = (int)(2 * seq_len);
    auto vpad = torch::empty({B, D, N}, projected.options());
    auto gate = torch::empty({B, L, D}, projected.options());
    const int threads = 256;
    long long total = (long long)B * L * D;
    prep_v_gate_kernel<<<(total + threads - 1) / threads, threads>>>(projected.data_ptr<float>(), w0.data_ptr<float>(), b0.data_ptr<float>(), w1.data_ptr<float>(), b1.data_ptr<float>(), vpad.data_ptr<float>(), gate.data_ptr<float>(), B, L, D, K, N);
    long long ztotal = (long long)B * D * (N - L);
    zero_second_half_kernel<<<(ztotal + threads - 1) / threads, threads>>>(vpad.data_ptr<float>(), B, D, L, N);
    return {vpad, gate};
}

torch::Tensor prepare_h(torch::Tensor fw1, torch::Tensor fb1, torch::Tensor fw2, torch::Tensor fb2, int64_t seq_len) {
    CHECK_INPUT(fw1); CHECK_INPUT(fb1); CHECK_INPUT(fw2); CHECK_INPUT(fb2);
    int D = fw2.size(0), L = (int)seq_len, N = 2 * L;
    auto hpad = torch::empty({D, N}, fw2.options());
    const int threads = 256;
    long long total = (long long)D * N;
    prep_h_kernel<<<(total + threads - 1) / threads, threads>>>(fw1.data_ptr<float>(), fb1.data_ptr<float>(), fw2.data_ptr<float>(), fb2.data_ptr<float>(), hpad.data_ptr<float>(), D, L, N);
    return hpad;
}

void mul_freq_inplace(torch::Tensor vf, torch::Tensor hf) {
    CHECK_CUDA(vf); CHECK_CUDA(hf); CHECK_CONTIGUOUS(vf); CHECK_CONTIGUOUS(hf);
    int B = vf.size(0), D = vf.size(1), F = vf.size(2);
    const int threads = 256;
    long long total = (long long)B * D * F;
    cmul_broadcast_inplace_kernel<<<(total + threads - 1) / threads, threads>>>(reinterpret_cast<float2*>(vf.data_ptr<c10::complex<float>>()), reinterpret_cast<const float2*>(hf.data_ptr<c10::complex<float>>()), B, D, F);
}

torch::Tensor apply_gate_transpose(torch::Tensor y, torch::Tensor gate, int64_t Larg) {
    CHECK_INPUT(y); CHECK_INPUT(gate);
    int B = y.size(0), D = y.size(1), N = y.size(2), L = (int)Larg;
    auto out = torch::empty({B, L, D}, y.options());
    const int threads = 256;
    long long total = (long long)B * L * D;
    apply_gate_transpose_kernel<<<(total + threads - 1) / threads, threads>>>(y.data_ptr<float>(), gate.data_ptr<float>(), out.data_ptr<float>(), B, L, D, N);
    return out;
}
'''

_ext = load_inline(name='hyena_order2_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['prepare_v_gate', 'prepare_h', 'mul_freq_inplace', 'apply_gate_transpose'],
                   extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, d_model: int, seq_len: int, order: int, short_filter_size: int):
        super(ModelNew, self).__init__()
        self.d_model = d_model
        self.seq_len = seq_len
        self.order = order
        self.fft_size = 2 * seq_len
        self.in_proj = nn.Linear(d_model, d_model * (order + 1), bias=True)
        self.short_convs = nn.ModuleList([
            nn.Conv1d(d_model, d_model, short_filter_size,
                      padding=short_filter_size // 2, groups=d_model)
            for _ in range(order)
        ])
        self.filter_mlp = nn.Sequential(nn.Linear(1, 64), nn.SiLU(), nn.Linear(64, d_model))
        self.register_buffer('filter_positions', torch.linspace(0, 1, seq_len).unsqueeze(-1))
        self.out_proj = nn.Linear(d_model, d_model, bias=True)
        self._hf_cache = None
        self._hf_cache_key = None

    def _get_hf(self):
        key = (self.filter_mlp[0].weight._version, self.filter_mlp[0].bias._version,
               self.filter_mlp[2].weight._version, self.filter_mlp[2].bias._version,
               self.filter_mlp[0].weight.data_ptr(), self.filter_mlp[0].bias.data_ptr(),
               self.filter_mlp[2].weight.data_ptr(), self.filter_mlp[2].bias.data_ptr(),
               self.filter_mlp[2].weight.device.index)
        if self._hf_cache is None or self._hf_cache_key != key:
            hpad = _ext.prepare_h(self.filter_mlp[0].weight.contiguous(), self.filter_mlp[0].bias.contiguous(),
                                  self.filter_mlp[2].weight.contiguous(), self.filter_mlp[2].bias.contiguous(),
                                  self.seq_len)
            self._hf_cache = torch.fft.rfft(hpad, dim=-1)
            self._hf_cache_key = key
        return self._hf_cache

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        projected = F.linear(x, self.in_proj.weight, self.in_proj.bias)
        vpad, gate = _ext.prepare_v_gate(projected.contiguous(),
                                         self.short_convs[0].weight.contiguous(), self.short_convs[0].bias.contiguous(),
                                         self.short_convs[1].weight.contiguous(), self.short_convs[1].bias.contiguous(),
                                         self.seq_len)
        vf = torch.fft.rfft(vpad, dim=-1)
        _ext.mul_freq_inplace(vf, self._get_hf())
        y = torch.fft.irfft(vf, n=self.fft_size, dim=-1)
        yt = _ext.apply_gate_transpose(y, gate, x.shape[1])
        return F.linear(yt, self.out_proj.weight, self.out_proj.bias)