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

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

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

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

// Vectorized type for 128-bit access with doubles
struct __align__(16) Double2 {
    double x, y;
};

// Core computation
// Formula: x * sin(x)
__device__ __forceinline__ double asu_op(double x) {
    return x * sin(x);
}

__global__ void asu_kernel_double(
    const double* __restrict__ input,
    double* __restrict__ output,
    const int n_elements) 
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;

    // 1. Vectorized Loop
    int vec_loops = n_elements / 2;
    const Double2* vec_input = reinterpret_cast<const Double2*>(input);
    Double2* vec_output = reinterpret_cast<Double2*>(output);

    for (int i = idx; i < vec_loops; i += stride) {
        Double2 in_val = vec_input[i];
        Double2 out_val;
        
        out_val.x = asu_op(in_val.x);
        out_val.y = asu_op(in_val.y);
        
        vec_output[i] = out_val;
    }

    // 2. Scalar Loop
    int tail_start = vec_loops * 2;
    for (int i = tail_start + idx; i < n_elements; i += stride) {
        output[i] = asu_op(input[i]);
    }
}

torch::Tensor asu_cuda_forward(const torch::Tensor& input) {
    TORCH_CHECK(input.is_cuda(), "Input tensor must be a CUDA tensor");
    TORCH_CHECK(input.scalar_type() == torch::kDouble, "Input tensor must be float64");
    TORCH_CHECK(input.is_contiguous(), "Input tensor must be contiguous");

    auto output = torch::empty_like(input);
    const int n_elements = input.numel();

    const int block_size = 256;
    // Grid size for Double2 (2 elements per thread)
    int grid_size = (n_elements + block_size * 2 - 1) / (block_size * 2);
    if (grid_size > 65535) grid_size = 65535;

    asu_kernel_double<<<grid_size, block_size>>>(
        input.data_ptr<double>(),
        output.data_ptr<double>(),
        n_elements
    );

    return output;
}
"""

asu_op = load_inline(
    name='asu_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['asu_cuda_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

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

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

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

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