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

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

torch::Tensor quantile_loss_cuda_forward(
    const torch::Tensor& pred, 
    const torch::Tensor& target,
    float q,
    std::string reduction);
"""

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

#define BLOCK_SIZE 256

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

// Branchless Quantile Loss
// loss = max(q * diff, (q - 1) * diff) where diff = target - pred
__device__ __forceinline__ float compute_quantile(float p, float t, float q) {
    float diff = t - p;
    return fmaxf(q * diff, (q - 1.0f) * diff);
}

__global__ void quantile_loss_kernel(
    float* __restrict__ output,
    const float* __restrict__ pred,
    const float* __restrict__ target,
    const int n,
    const float q) 
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int vec_n = n / 4;
    
    int i = idx;
    const int stride = blockDim.x * gridDim.x;

    for (; i < vec_n; i += stride) {
        Float4 p_vec = reinterpret_cast<const Float4*>(pred)[i];
        Float4 t_vec = reinterpret_cast<const Float4*>(target)[i];
        Float4 out_vec;

        out_vec.x = compute_quantile(p_vec.x, t_vec.x, q);
        out_vec.y = compute_quantile(p_vec.y, t_vec.y, q);
        out_vec.z = compute_quantile(p_vec.z, t_vec.z, q);
        out_vec.w = compute_quantile(p_vec.w, t_vec.w, q);

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

    int start_scalar = vec_n * 4;
    int global_tid = blockIdx.x * blockDim.x + threadIdx.x;
    int total_threads = gridDim.x * blockDim.x;
    
    int current_idx = start_scalar + global_tid;
    while (current_idx < n) {
        output[current_idx] = compute_quantile(pred[current_idx], target[current_idx], q);
        current_idx += total_threads;
    }
}

torch::Tensor quantile_loss_cuda_forward(
    const torch::Tensor& pred, 
    const torch::Tensor& target,
    float q,
    std::string reduction) 
{
    TORCH_CHECK(pred.is_cuda() && target.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(pred.is_contiguous() && target.is_contiguous(), "Inputs must be contiguous");
    TORCH_CHECK(pred.numel() == target.numel(), "Shapes must match");

    const int n = pred.numel();
    
    auto output = torch::empty_like(pred);

    const int vec_n = n / 4;
    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;

    quantile_loss_kernel<<<final_grid, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        pred.data_ptr<float>(),
        target.data_ptr<float>(),
        n,
        q
    );

    if (reduction == "mean") {
        return output.mean();
    } else if (reduction == "sum") {
        return output.sum();
    }

    return output;
}
"""

quantile_op_module = load_inline(
    name='quantile_loss_op_fixed',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['quantile_loss_cuda_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self, quantile=0.9, reduction='none'):
        super(ModelNew, self).__init__()
        self.quantile = quantile
        self.reduction = reduction
        self.op = quantile_op_module

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        return self.op.quantile_loss_cuda_forward(
            pred.contiguous(), 
            target.contiguous(), 
            self.quantile, 
            self.reduction
        )