import torch
import torch.nn as nn

BATCH_SIZE = 1024 * 1024 
SHAPE = (BATCH_SIZE, 4)

EPS_VALUE = 1e-7
REDUCTION_MODE = 'none' 

class GIoULoss(nn.Module):
    def __init__(self, eps=1e-7, reduction='none'):
        super(GIoULoss, self).__init__()
        self.eps = eps
        self.reduction = reduction

    def forward(self, b1: torch.Tensor, b2: torch.Tensor) -> torch.Tensor:
        b1_area = (b1[:, 2] - b1[:, 0]) * (b1[:, 3] - b1[:, 1])
        b2_area = (b2[:, 2] - b2[:, 0]) * (b2[:, 3] - b2[:, 1])

        inter_x1 = torch.max(b1[:, 0], b2[:, 0])
        inter_y1 = torch.max(b1[:, 1], b2[:, 1])
        inter_x2 = torch.min(b1[:, 2], b2[:, 2])
        inter_y2 = torch.min(b1[:, 3], b2[:, 3])

        inter_w = (inter_x2 - inter_x1).clamp(min=0)
        inter_h = (inter_y2 - inter_y1).clamp(min=0)
        inter_area = inter_w * inter_h

        union_area = b1_area + b2_area - inter_area + self.eps

        iou = inter_area / union_area

        c_x1 = torch.min(b1[:, 0], b2[:, 0])
        c_y1 = torch.min(b1[:, 1], b2[:, 1])
        c_x2 = torch.max(b1[:, 2], b2[:, 2])
        c_y2 = torch.max(b1[:, 3], b2[:, 3])

        c_area = (c_x2 - c_x1) * (c_y2 - c_y1) + self.eps

        giou = iou - (c_area - union_area) / c_area
        loss = 1.0 - giou

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, eps=1e-7, reduction='none'):
        super(Model, self).__init__()
        self.giou_loss = GIoULoss(eps=eps, reduction=reduction)
    
    def forward(self, b1: torch.Tensor, b2: torch.Tensor) -> torch.Tensor:
        return self.giou_loss(b1, b2)

def get_inputs():
    b1 = torch.randn(SHAPE, dtype=torch.float32)
    b2 = torch.randn(SHAPE, dtype=torch.float32)
    
    def make_valid_box(b):
        x_min, _ = torch.min(b[:, [0, 2]], dim=1, keepdim=True)
        x_max, _ = torch.max(b[:, [0, 2]], dim=1, keepdim=True)
        y_min, _ = torch.min(b[:, [1, 3]], dim=1, keepdim=True)
        y_max, _ = torch.max(b[:, [1, 3]], dim=1, keepdim=True)
        return torch.cat([x_min, y_min, x_max, y_max], dim=1).contiguous()

    return [make_valid_box(b1), make_valid_box(b2)]

def get_init_inputs():
    return [EPS_VALUE, REDUCTION_MODE]