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

cpp_source = """
#include <torch/extension.h>

torch::Tensor angle_cuda_forward(const torch::Tensor& input);
"""

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

// Vectorized types
struct __align__(16) Float4 {
    float x, y, z, w;
};

struct __align__(8) Float2 {
    float x, y;
};

// Core logic: angle(z) = atan2(imag, real)
__device__ __forceinline__ float compute_angle(float real, float imag) {
    return atan2f(imag, real);
}

__global__ void angle_kernel(
    const float* __restrict__ input, // Treated as raw float array of complex parts
    float* __restrict__ output,
    const int n_complex_elements) 
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    
    // Total float4 vectors to process
    int vec_loops = n_complex_elements / 2;

    const Float4* vec_input = reinterpret_cast<const Float4*>(input);
    Float2* vec_output = reinterpret_cast<Float2*>(output);

    for (int i = idx; i < vec_loops; i += stride) {
        // Load 2 complex numbers: [r1, i1, r2, i2]
        Float4 in_val = vec_input[i];
        
        Float2 out_val;
        // z1 = x + iy -> in.x + i*in.y
        out_val.x = compute_angle(in_val.x, in_val.y);
        // z2 = z + iw -> in.z + i*in.w
        out_val.y = compute_angle(in_val.z, in_val.w);
        
        vec_output[i] = out_val;
    }

    // Handle tail
    int tail_idx = vec_loops * 2;
    if (tail_idx + idx < n_complex_elements) {
        // We need to access the float array at offset = tail_idx * 2 (real) and +1 (imag)
        // because each complex element has 2 floats.
        int complex_idx = tail_idx + idx;
        // Only one thread will typically handle the single tail element if aligned
        // But stride logic applies if grid is small.
        // Simplified tail check for grid-stride:
        if (complex_idx < n_complex_elements) {
             float real = input[complex_idx * 2];
             float imag = input[complex_idx * 2 + 1];
             output[complex_idx] = compute_angle(real, imag);
        }
    }
}

torch::Tensor angle_cuda_forward(const torch::Tensor& input) {
    TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
    TORCH_CHECK(input.is_contiguous(), "Input must be contiguous");
    TORCH_CHECK(input.scalar_type() == torch::kComplexFloat, "Input must be ComplexFloat (complex64)");

    int n_elements = input.numel();
    // Output is real float32
    auto output = torch::empty(input.sizes(), input.options().dtype(torch::kFloat32));

    const int block_size = 256;
    int num_vectors = (n_elements + 1) / 2; 
    int grid_size = (num_vectors + block_size - 1) / block_size;
    if (grid_size > 65535) grid_size = 65535;

    angle_kernel<<<grid_size, block_size>>>(
        reinterpret_cast<float*>(input.data_ptr<c10::complex<float>>()),
        output.data_ptr<float>(),
        n_elements
    );

    return output;
}
"""

angle_op = load_inline(
    name='angle_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['angle_cuda_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3', '--use_fast_math']
)

class AngleNew(nn.Module):
    def __init__(self):
        super(AngleNew, self).__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return angle_op.angle_cuda_forward(x)

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.act(x)