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

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void complex_conj_mul_div_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    const float* __restrict__ c,
    float* __restrict__ output,
    int N, float eps
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        int offset = idx * 2;

        // Input A = Xa + iYa
        float Xa = a[offset];
        float Ya = a[offset + 1];

        // Input B = Xb + iYb
        float Xb = b[offset];
        float Yb = b[offset + 1];

        // Input C = Xc + iYc
        float Xc = c[offset];
        float Yc = c[offset + 1];

        // Step 1 & 2: P = conj(A) * B = Xp + iYp
        // Xp = Xa*Xb + Ya*Yb
        // Yp = Xa*Yb - Ya*Xb
        float Xp = fmaf(Xa, Xb, Ya * Yb);
        float Yp = fmaf(Xa, Yb, -Ya * Xb);

        // Step 3: Division Out = P / C
        // Denominator D = |C|^2 = Xc^2 + Yc^2
        float D = fmaf(Xc, Xc, Yc * Yc);
        float inv_D = 1.0f / (D + eps);

        // Out_Re = (Xp*Xc + Yp*Yc) / D
        // Out_Im = (Yp*Xc - Xp*Yc) / D
        float Out_Re = fmaf(Xp, Xc, Yp * Yc) * inv_D;
        float Out_Im = fmaf(Yp, Xc, -Xp * Yc) * inv_D;

        output[offset] = Out_Re;
        output[offset + 1] = Out_Im;
    }
}

torch::Tensor complex_conj_mul_div_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor c) {
    auto output = torch::empty_like(a);
    int N = a.size(0);

    const int block_size = 256;
    int num_blocks = (N + block_size - 1) / block_size;

    complex_conj_mul_div_kernel<<<num_blocks, block_size>>>(
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        c.data_ptr<float>(),
        output.data_ptr<float>(),
        N, 1e-12f
    );

    return output;
}
"""

cpp_source = """
torch::Tensor complex_conj_mul_div_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor c);
"""

module = load_inline(
    name="complex_conj_mul_div",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["complex_conj_mul_div_cuda"],
    verbose=True
)


class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.module = module

    def forward(self, a, b, c):
        return self.module.complex_conj_mul_div_cuda(a, b, c)