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

B, COORD = 32, 4


class ModelNew(nn.Module):
    def __init__(self, block_size=256):
        super().__init__()
        self.block_size = block_size
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor iou_fused_cuda(torch::Tensor pred, torch::Tensor target);
        """

        cuda_source = f"""
        #include <cuda_runtime.h>
        #include <algorithm> // For std::min, std::max
        #include <cmath>     // For fmaxf, fminf

        #define BLOCK_SIZE {self.block_size}
        #define EPSILON 1e-6f

        // 融合核函数：计算 IoU Loss (1 - IoU)
        // 每个线程处理一个批次样本 (一个边界框对)
        __global__ void iou_loss_fused_kernel(
            const float* __restrict__ P, // B x 4 (x1, y1, x2, y2)
            const float* __restrict__ T, // B x 4 (x1, y1, x2, y2)
            float* __restrict__ D,       // B x 1 (Loss per item)
            int B,
            int COORD
        ) {{
            int b = blockIdx.x * blockDim.x + threadIdx.x;
            if (b < B) {{
                
                // 预测框坐标 (P[b, 0]...P[b, 3])
                const float px1 = P[b * COORD + 0];
                const float py1 = P[b * COORD + 1];
                const float px2 = P[b * COORD + 2];
                const float py2 = P[b * COORD + 3];
                
                // 目标框坐标 (T[b, 0]...T[b, 3])
                const float tx1 = T[b * COORD + 0];
                const float ty1 = T[b * COORD + 1];
                const float tx2 = T[b * COORD + 2];
                const float ty2 = T[b * COORD + 3];

                // 1. 交集区域坐标
                const float ix1 = fmaxf(px1, tx1);
                const float iy1 = fmaxf(py1, ty1);
                const float ix2 = fminf(px2, tx2);
                const float iy2 = fminf(py2, ty2);

                // 2. 交集区域边长和面积
                const float iw = fmaxf(ix2 - ix1, 0.0f);
                const float ih = fmaxf(iy2 - iy1, 0.0f);
                const float intersection = iw * ih;

                // 3. 预测框和目标框面积
                const float area_p = (px2 - px1) * (py2 - py1);
                const float area_t = (tx2 - tx1) * (ty2 - ty1);

                // 4. 并集区域面积: A_union = A_p + A_t - A_intersection
                const float union_area = area_p + area_t - intersection;

                // 5. IoU 和 Loss
                const float iou = intersection / (union_area + EPSILON);
                const float loss = 1.0f - iou;
                
                // 存储每个样本的损失
                D[b] = loss;
            }}
        }}

        torch::Tensor iou_fused_cuda(torch::Tensor pred, torch::Tensor target) {{

            TORCH_CHECK(pred.is_cuda() && target.is_cuda(), "Inputs must be CUDA tensors");
            TORCH_CHECK(pred.dim() == 2 && target.dim() == 2, "Inputs must be 2D (B, COORD)");

            pred = pred.contiguous();
            target = target.contiguous();

            int B = pred.size(0);
            int COORD = pred.size(1);
            
            // 输出 D 的形状是 B x 1 (每批次样本的损失)
            auto D = torch::empty({{B, 1}}, pred.options());
            
            int blocks = std::min((B + BLOCK_SIZE - 1) / BLOCK_SIZE, 1024);
            
            iou_loss_fused_kernel<<<blocks, BLOCK_SIZE>>>(
                pred.data_ptr<float>(),
                target.data_ptr<float>(),
                D.data_ptr<float>(),
                B,
                COORD
            );

            return D;
        }}
        """

        self.op = load_inline(
            name='iou_loss_opt',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['iou_fused_cuda'],
            extra_cuda_cflags=['-O3', '--use_fast_math'],
            verbose=False
        )

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        if not pred.is_cuda: pred = pred.cuda()
        if not target.is_cuda: target = target.cuda()
        
        # 调用融合内核，返回每批次样本的损失 (B x 1)
        per_item_loss = self.op.iou_fused_cuda(pred, target)
        
        # 在 C++ 侧对 B x 1 结果求平均，实现最终的标量损失
        return per_item_loss.mean()