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 bpr_loss_cuda_forward(
    const torch::Tensor& pos_scores, 
    const torch::Tensor& neg_scores,
    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;
};

// Softplus(x) = log(1 + exp(x))
__device__ __forceinline__ float stable_softplus(float x) {
    if (x > 0) {
        return x + logf(1.0f + expf(-x));
    } else {
        return logf(1.0f + expf(x));
    }
}

__global__ void bpr_loss_kernel(
    float* __restrict__ output,
    const float* __restrict__ pos,
    const float* __restrict__ neg,
    const int n_elements) 
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int vec_n = n_elements / 4;
    int stride = blockDim.x * gridDim.x;
    
    // 1. Vectorized Loop
    for (int i = idx; i < vec_n; i += stride) {
        Float4 p_vec = reinterpret_cast<const Float4*>(pos)[i];
        Float4 n_vec = reinterpret_cast<const Float4*>(neg)[i];
        Float4 out_vec;

        // BPR Loss: -log_sigmoid(p - n) <=> softplus(n - p)
        float d1 = n_vec.x - p_vec.x;
        float d2 = n_vec.y - p_vec.y;
        float d3 = n_vec.z - p_vec.z;
        float d4 = n_vec.w - p_vec.w;

        out_vec.x = stable_softplus(d1);
        out_vec.y = stable_softplus(d2);
        out_vec.z = stable_softplus(d3);
        out_vec.w = stable_softplus(d4);

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

    // 2. Scalar Tail
    int tail_start = vec_n * 4;
    int global_tid = idx * 4;
    
    int scalar_idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (scalar_idx < (n_elements - tail_start)) {
        int real_idx = tail_start + scalar_idx;
        float p_val = pos[real_idx];
        float n_val = neg[real_idx];
        output[real_idx] = stable_softplus(n_val - p_val);
    }
}

torch::Tensor bpr_loss_cuda_forward(
    const torch::Tensor& pos, 
    const torch::Tensor& neg,
    std::string reduction) 
{
    TORCH_CHECK(pos.is_cuda() && neg.is_cuda(), "Inputs must be CUDA");
    TORCH_CHECK(pos.is_contiguous() && neg.is_contiguous(), "Inputs must be contiguous");
    
    // 形状必须完全一致
    TORCH_CHECK(pos.sizes() == neg.sizes(), "Input shapes must match exactly");

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

    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;

    bpr_loss_kernel<<<final_grid, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        pos.data_ptr<float>(),
        neg.data_ptr<float>(),
        n
    );

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

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self, reduction='none'):
        super(ModelNew, self).__init__()
        self.reduction = reduction
        self.op = load_inline(
            name='bpr_loss_opt',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['bpr_loss_cuda_forward'],
            verbose=False,
            extra_cuda_cflags=['-O3']
        )

    def forward(self, pos: torch.Tensor, neg: torch.Tensor) -> torch.Tensor:
        return self.op.bpr_loss_cuda_forward(pos.contiguous(), neg.contiguous(), self.reduction)