import torch
import torch.nn as nn
import torch.nn.functional as F


B, COORD = 32, 4
EPS = 1e-6


class IoULossBatch(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        px1, py1, px2, py2 = pred.split(1, dim=1)
        tx1, ty1, tx2, ty2 = target.split(1, dim=1)

        # 1. 计算交集区域坐标
        ix1 = torch.max(px1, tx1)
        iy1 = torch.max(py1, ty1)
        ix2 = torch.min(px2, tx2)
        iy2 = torch.min(py2, ty2)

        # 2. 计算交集区域面积 (确保宽度/高度非负)
        iw = torch.max(ix2 - ix1, torch.tensor(0.0).to(pred.device))
        ih = torch.max(iy2 - iy1, torch.tensor(0.0).to(pred.device))
        intersection = iw * ih

        # 3. 计算预测框和目标框的面积
        area_p = (px2 - px1) * (py2 - py1)
        area_t = (tx2 - tx1) * (ty2 - ty1)

        # 4. 计算并集区域面积: Union = Area_p + Area_t - Intersection
        union = area_p + area_t - intersection
        
        # 5. 计算 IoU
        iou = intersection / (union + EPS)

        # 6. 计算 IoU Loss (1 - IoU)
        loss = 1.0 - iou
        
        return loss.mean()


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.op = IoULossBatch()

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        return self.op(pred, target)


def get_inputs():
    torch.manual_seed(42)
    # 随机生成坐标 (0, 100) 范围的基数
    base = torch.rand(B, COORD) * 100
    
    # 构建预测框 (确保 x1 < x2 且 y1 < y2)
    pred = torch.empty_like(base)
    pred[:, 0] = torch.min(base[:, 0], base[:, 2])
    pred[:, 1] = torch.min(base[:, 1], base[:, 3])
    pred[:, 2] = torch.max(base[:, 0], base[:, 2])
    pred[:, 3] = torch.max(base[:, 1], base[:, 3])
    
    # 构建目标框
    target = torch.empty_like(base)
    target[:, 0] = torch.min(base[:, 0] + 5, base[:, 2] + 5)
    target[:, 1] = torch.min(base[:, 1] + 5, base[:, 3] + 5)
    target[:, 2] = torch.max(base[:, 0] + 5, base[:, 2] + 5)
    target[:, 3] = torch.max(base[:, 1] + 5, base[:, 3] + 5)

    return [pred, target]


def get_init_inputs():
    return []