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

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

torch::Tensor pgelu_cuda_forward(
    const torch::Tensor& input, 
    const torch::Tensor& alpha, 
    const torch::Tensor& beta);
"""

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

#define BLOCK_SIZE 256
#define PI_HALF 1.5707963267948966
#define EPS 1e-7

// double2 for 128-bit vectorization
struct __align__(16) Double2 {
    double x, y;
};

// P-GELU for double
__device__ __forceinline__ double compute_pgelu_double(double x, double alpha, double beta) {
    double inner = alpha * x + beta * x * x * x;
    double clamped_inner = fmin(fmax(inner, -PI_HALF + EPS), PI_HALF - EPS);
    return x * (1.0 + tan(clamped_inner));
}

template <typename T>
__global__ void pgelu_kernel(
    T* __restrict__ output,
    const T* __restrict__ input,
    const int n,
    T alpha,
    T beta) 
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int vec_n = n / 2; // double2
    
    int i = idx;
    const int stride = blockDim.x * gridDim.x;

    for (; i < vec_n; i += stride) {
        Double2 in_vec = reinterpret_cast<const Double2*>(input)[i];
        Double2 out_vec;

        out_vec.x = compute_pgelu_double(in_vec.x, alpha, beta);
        out_vec.y = compute_pgelu_double(in_vec.y, alpha, beta);

        reinterpret_cast<Double2*>(output)[i] = out_vec;
    }

    int tail_idx = vec_n * 2 + idx;
    if (idx == 0 && tail_idx < n) {
        output[tail_idx] = compute_pgelu_double(input[tail_idx], alpha, beta);
    }
}

torch::Tensor pgelu_cuda_forward(
    const torch::Tensor& input, 
    const torch::Tensor& alpha_t, 
    const torch::Tensor& beta_t) 
{
    TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
    TORCH_CHECK(input.is_contiguous(), "Input must be contiguous");
    
    const int n = input.numel();
    auto output = torch::empty_like(input);

    const double alpha = alpha_t.item<double>();
    const double beta = beta_t.item<double>();

    const int vec_n = n / 2;
    const int grid_size = (vec_n + BLOCK_SIZE - 1) / BLOCK_SIZE;
    
    int final_grid = (grid_size < 1) ? 1 : grid_size;
    if (final_grid > 65535) final_grid = 65535;
    
    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "pgelu_kernel", ([&]{
        pgelu_kernel<scalar_t><<<final_grid, BLOCK_SIZE>>>(
            output.data_ptr<scalar_t>(),
            input.data_ptr<scalar_t>(),
            n,
            static_cast<scalar_t>(alpha),
            static_cast<scalar_t>(beta)
        );
    }));

    return output;
}
"""

pgelu_op_module = load_inline(
    name='pgelu_tan_op_double',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['pgelu_cuda_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self, alpha_init=1.0, beta_init=1.0):
        super(ModelNew, self).__init__()
        self.alpha = nn.Parameter(torch.tensor(alpha_init, dtype=torch.float64))
        self.beta = nn.Parameter(torch.tensor(beta_init, dtype=torch.float64))
        self.op = pgelu_op_module

    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        return self.op.pgelu_cuda_forward(input_tensor.contiguous(), self.alpha, self.beta)