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 giou_loss_cuda_forward(
    const torch::Tensor& b1, 
    const torch::Tensor& b2,
    float eps,
    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;
};

__global__ void giou_loss_kernel(
    float* __restrict__ output,
    const float* __restrict__ b1,
    const float* __restrict__ b2,
    const int num_boxes,
    const float eps) 
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= num_boxes) return;

    // 1. Vectorized Load
    Float4 box1 = reinterpret_cast<const Float4*>(b1)[idx];
    Float4 box2 = reinterpret_cast<const Float4*>(b2)[idx];

    float b1_x1 = box1.x; float b1_y1 = box1.y; float b1_x2 = box1.z; float b1_y2 = box1.w;
    float b2_x1 = box2.x; float b2_y1 = box2.y; float b2_x2 = box2.z; float b2_y2 = box2.w;

    // 2. Area
    float area1 = (b1_x2 - b1_x1) * (b1_y2 - b1_y1);
    float area2 = (b2_x2 - b2_x1) * (b2_y2 - b2_y1);

    // 3. Intersection
    float inter_x1 = fmaxf(b1_x1, b2_x1);
    float inter_y1 = fmaxf(b1_y1, b2_y1);
    float inter_x2 = fminf(b1_x2, b2_x2);
    float inter_y2 = fminf(b1_y2, b2_y2);

    float inter_w = fmaxf(inter_x2 - inter_x1, 0.0f);
    float inter_h = fmaxf(inter_y2 - inter_y1, 0.0f);
    float inter_area = inter_w * inter_h;

    // 4. Union
    float union_area = area1 + area2 - inter_area + eps;

    // 5. IoU
    float iou = inter_area / union_area;

    // 6. Enclosing Box
    float c_x1 = fminf(b1_x1, b2_x1);
    float c_y1 = fminf(b1_y1, b2_y1);
    float c_x2 = fmaxf(b1_x2, b2_x2);
    float c_y2 = fmaxf(b1_y2, b2_y2);

    float c_area = (c_x2 - c_x1) * (c_y2 - c_y1) + eps;

    // 7. GIoU Loss
    float giou = iou - (c_area - union_area) / c_area;
    
    output[idx] = 1.0f - giou;
}

torch::Tensor giou_loss_cuda_forward(
    const torch::Tensor& b1, 
    const torch::Tensor& b2,
    float eps,
    std::string reduction) 
{
    TORCH_CHECK(b1.is_cuda() && b2.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(b1.is_contiguous() && b2.is_contiguous(), "Inputs must be contiguous");
    TORCH_CHECK(b1.size(1) == 4 && b2.size(1) == 4, "Boxes must be (N, 4)");
    TORCH_CHECK(b1.size(0) == b2.size(0), "Batch size must match");

    const int N = b1.size(0);
    auto output = torch::empty({N}, b1.options());

    const int grid_size = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;

    giou_loss_kernel<<<grid_size, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        b1.data_ptr<float>(),
        b2.data_ptr<float>(),
        N,
        eps
    );

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

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self, eps=1e-7, reduction='none'):
        super(ModelNew, self).__init__()
        self.eps = eps
        self.reduction = reduction
        self.op = load_inline(
            name='giou_loss_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['giou_loss_cuda_forward'],
            verbose=False
        )

    def forward(self, b1: torch.Tensor, b2: torch.Tensor) -> torch.Tensor:
        return self.op.giou_loss_cuda_forward(
            b1.contiguous(), 
            b2.contiguous(), 
            self.eps, 
            self.reduction
        )