import torch
from torch.utils.cpp_extension import load_inline

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

// Hexpo Kernel: Piecewise function defined by a, b, c, d
__global__ void hexpo_kernel(const float* x, float* y, int size, float a, float b, float c, float d) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        float input = x[idx];
        float output;
        
        if (input >= 0.0f) {
            // x >= 0: -a * (exp(-x/b) - 1)
            output = -a * (expf(-input / b) - 1.0f);
        } else {
            // x < 0: c * (exp(x/d) - 1)
            output = c * (expf(input / d) - 1.0f);
        }
        
        y[idx] = output;
    }
}

torch::Tensor hexpo_cuda(torch::Tensor x, float a, float b, float c, float d) {
    TORCH_CHECK(x.is_cuda() && x.dtype() == torch::kFloat32, "Input must be a float32 CUDA tensor.");
    
    auto size = x.numel();
    auto y = torch::empty_like(x);
    
    const int block_size = 256;
    int num_blocks = (size + block_size - 1) / block_size;
    
    hexpo_kernel<<<num_blocks, block_size>>>(
        x.data_ptr<float>(), 
        y.data_ptr<float>(), 
        size, 
        a, b, c, d
    );
    
    return y;
}
"""

hexpo_cpp_source = """
torch::Tensor hexpo_cuda(torch::Tensor x, float a, float b, float c, float d);
"""

hexpo = load_inline(
    name="hexpo",
    cpp_sources=hexpo_cpp_source,
    cuda_sources=hexpo_source,
    functions=["hexpo_cuda"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    """
    Model using the custom Hexpo CUDA kernel.
    """
    def __init__(self, a: float = 1.0, b: float = 1.0, c: float = 1.0, d: float = 1.0):
        super(ModelNew, self).__init__()
        self.hexpo = hexpo
        self.a = a
        self.b = b
        self.c = c
        self.d = d

    def forward(self, x):
        return self.hexpo.hexpo_cuda(x, self.a, self.b, self.c, self.d)