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>

__device__ void complex_exp(float re, float im, float& out_re, float& out_im) {
    float exp_re = expf(re);
    out_re = exp_re * cosf(im);
    out_im = exp_re * sinf(im);
}

__device__ void complex_log(float re, float im, float& out_re, float& out_im) {
    out_re = 0.5f * logf(re * re + im * im);
    out_im = atan2f(im, re);
}

__device__ void complex_pow(float z_re, float z_im, float p_re, float p_im, float& out_re, float& out_im) {
    float log_re, log_im;
    complex_log(z_re, z_im, log_re, log_im);

    float prod_re = p_re * log_re - p_im * log_im;
    float prod_im = p_re * log_im + p_im * log_re;

    complex_exp(prod_re, prod_im, out_re, out_im);
}

__global__ void complex_ops_kernel(
    const float* __restrict__ z_f,
    float* __restrict__ output,
    float p_re,
    float p_im,
    int batch_size) {

    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < batch_size) {
        float z_re = z_f[idx * 2];
        float z_im = z_f[idx * 2 + 1];

        float y_re, y_im;
        complex_exp(z_re, z_im, y_re, y_im);

        float w_re, w_im;
        complex_log(y_re, y_im, w_re, w_im);

        float out_re, out_im;
        complex_pow(w_re, w_im, p_re, p_im, out_re, out_im);

        output[idx * 2] = out_re;
        output[idx * 2 + 1] = out_im;
    }
}

torch::Tensor complex_ops_cuda(
    torch::Tensor z_f,
    float p_re,
    float p_im) {

    int batch_size = z_f.size(0);
    auto output = torch::empty_like(z_f);

    int threads = 256;
    int blocks = (batch_size + threads - 1) / threads;

    complex_ops_kernel<<<blocks, threads>>>(
        z_f.data_ptr<float>(),
        output.data_ptr<float>(),
        p_re,
        p_im,
        batch_size
    );

    return output;
}
"""

cpp_source = """
torch::Tensor complex_ops_cuda(
    torch::Tensor z_f,
    float p_re,
    float p_im);
"""

cuda_module = load_inline(
    name="complex_ops_module",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["complex_ops_cuda"],
    verbose=True
)


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

    def forward(self, z_f):
        return cuda_module.complex_ops_cuda(z_f, self.p_re, self.p_im)