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


class ModelNew(nn.Module):
    def __init__(self, num_bins=10, min_val=0.0, max_val=1.0):
        super().__init__()
        self.num_bins = num_bins
        self.min_val = min_val
        self.max_val = max_val
        self.step = (max_val - min_val) / num_bins
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        torch::Tensor histogramloss_cuda(torch::Tensor pos, torch::Tensor neg, int num_bins, float min_val, float max_val);
        """

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

        __global__ void histogram_kernel(
            const float* __restrict__ data,
            float* __restrict__ hist,
            const int n_elements,
            const int num_bins,
            const float min_val,
            const float step)
        {
            const int tid = blockIdx.x * blockDim.x + threadIdx.x;
            const int stride = blockDim.x * gridDim.x;

            for (int i = tid; i < n_elements; i += stride) {
                float val = data[i];
                for (int b = 0; b < num_bins; ++b) {
                    float center = min_val + (b + 0.5f) * step;
                    float diff = fabsf(val - center);
                    float weight = fmaxf(0.0f, 1.0f - diff / step);
                    atomicAdd(&hist[b], weight);
                }
            }
        }

        torch::Tensor histogramloss_cuda(torch::Tensor pos, torch::Tensor neg, int num_bins, float min_val, float max_val) {
            auto pos_c = pos.contiguous();
            auto neg_c = neg.contiguous();

            auto pos_hist = torch::zeros({num_bins}, pos.options());
            auto neg_hist = torch::zeros({num_bins}, neg.options());

            float step = (max_val - min_val) / num_bins;

            const int threads = 256;

            int pos_blocks = min((int)((pos_c.numel() + threads - 1) / threads), 65535);
            histogram_kernel<<<pos_blocks, threads>>>(
                pos_c.data_ptr<float>(),
                pos_hist.data_ptr<float>(),
                pos_c.numel(),
                num_bins,
                min_val,
                step
            );

            int neg_blocks = min((int)((neg_c.numel() + threads - 1) / threads), 65535);
            histogram_kernel<<<neg_blocks, threads>>>(
                neg_c.data_ptr<float>(),
                neg_hist.data_ptr<float>(),
                neg_c.numel(),
                num_bins,
                min_val,
                step
            );

            auto pos_cdf = torch::cumsum(pos_hist, 0);
            pos_cdf = pos_cdf / (pos_cdf[num_bins - 1] + 1e-8);

            auto neg_pdf = neg_hist / (neg_hist.sum() + 1e-8);

            return (neg_pdf * pos_cdf).sum();
        }
        """

        self.op = load_inline(
            name="histogramloss_op",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["histogramloss_cuda"],
            extra_cuda_cflags=["-O3"],
            verbose=False
        )

    def forward(self, pos, neg):
        return self.op.histogramloss_cuda(pos, neg, self.num_bins, self.min_val, self.max_val)