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

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

__global__ void fast_ap_loss_kernel(
    const float* __restrict__ x, 
    const long* __restrict__ labels,
    float* __restrict__ ap_out,
    int batch_size, 
    int dim,
    int num_bins,
    float max_dist) 
{
    int i = blockIdx.x;
    if (i >= batch_size) return;

    int tid = threadIdx.x;
    int label_i = labels[i];


    extern __shared__ float s_query[];
    for (int k = tid; k < dim; k += blockDim.x) {
        s_query[k] = x[i * dim + k];
    }


    float* s_pos_hist = &s_query[dim];
    float* s_neg_hist = &s_pos_hist[num_bins];

    for (int k = tid; k < num_bins; k += blockDim.x) {
        s_pos_hist[k] = 0.0f;
        s_neg_hist[k] = 0.0f;
    }

    __syncthreads();

    float width = max_dist / (float)num_bins;
    float half_width = width * 0.5f;

    for (int j = tid; j < batch_size; j += blockDim.x) {
        if (i == j) continue;

        float dist_sq = 0.0f;
        const float* row_j = x + j * dim;


        int k = 0;
        for (; k + 3 < dim; k += 4) {
            float d0 = s_query[k] - row_j[k];
            float d1 = s_query[k+1] - row_j[k+1];
            float d2 = s_query[k+2] - row_j[k+2];
            float d3 = s_query[k+3] - row_j[k+3];
            dist_sq += d0*d0 + d1*d1 + d2*d2 + d3*d3;
        }
        for (; k < dim; ++k) {
            float d = s_query[k] - row_j[k];
            dist_sq += d*d;
        }

        int label_j = labels[j];
        bool is_pos = (label_i == label_j);


        for (int m = 0; m < num_bins; ++m) {
            float center = half_width + m * width;
            float diff = fabsf(dist_sq - center);
            float weight = 1.0f - diff / width;
            if (weight > 0.0f) {

                if (is_pos) {
                    atomicAdd(&s_pos_hist[m], weight);
                } else {
                    atomicAdd(&s_neg_hist[m], weight);
                }
            }
        }
    }

    __syncthreads();

    if (tid == 0) {
        float pos_cdf = 0.0f;
        float neg_cdf = 0.0f;
        float ap = 0.0f;


        float total_pos = 0.0f;
        for (int m = 0; m < num_bins; ++m) {
            total_pos += s_pos_hist[m];
        }


        total_pos += 1e-10f;


        for (int m = 0; m < num_bins; ++m) {
            float p_hist = s_pos_hist[m];
            float n_hist = s_neg_hist[m];

            pos_cdf += p_hist;
            neg_cdf += n_hist;

            float precision = pos_cdf / (pos_cdf + neg_cdf + 1e-10f);
            float delta_recall = p_hist / total_pos;

            ap += precision * delta_recall;
        }

        ap_out[i] = ap;
    }
}

torch::Tensor launch_fast_ap(torch::Tensor x, torch::Tensor labels, int num_bins) {
    auto batch_size = x.size(0);
    auto dim = x.size(1);
    auto ap_out = torch::empty({batch_size}, x.options());

    int shared_mem = (dim + 2 * num_bins) * sizeof(float);

    dim3 blocks(batch_size);
    dim3 threads(256);

    fast_ap_loss_kernel<<<blocks, threads, shared_mem>>>(
        x.data_ptr<float>(),
        labels.data_ptr<long>(),
        ap_out.data_ptr<float>(),
        batch_size,
        dim,
        num_bins,
        4.0f
    );

    return ap_out;
}
"""

cpp_source = """
torch::Tensor launch_fast_ap(torch::Tensor x, torch::Tensor labels, int num_bins);
"""

fast_ap_module = load_inline(
    name='fast_ap_loss_ext',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['launch_fast_ap'],
    verbose=False
)


class ModelNew(nn.Module):
    def __init__(self, num_bins=10):
        super(ModelNew, self).__init__()
        self.num_bins = num_bins
        self.op = fast_ap_module

    def forward(self, x: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        x = F.normalize(x, p=2, dim=1)

        ap_per_sample = self.op.launch_fast_ap(x.contiguous(), labels.contiguous(), self.num_bins)

        return 1.0 - ap_per_sample.mean()