"""Inner-IoU / Focal-IoU — 针对小目标定位精度的损失函数

Inner-IoU: 用缩放后的内部框计算IoU，让模型更关注框的内部重叠
           对小目标特别有效（60.8%的目标面积<0.5%）
Focal-IoU: 对高IoU样本加大权重，推动模型精修框的位置

论文: Inner-IoU: More Effective Intersection over Union Loss (2023)

Patch方式: 替换 BboxLoss.forward，不改模型结构
"""
import math
import torch
import torch.nn.functional as F


def inner_bbox_iou(box1, box2, xywh=False, ratio=0.7, eps=1e-7):
    """Inner-IoU: 用缩放后的内部框计算CIoU。

    Args:
        box1, box2: 预测框和目标框 (xyxy格式)
        ratio: 内部框缩放比例 (0.5~0.9, 越小越关注中心)
        eps: 防止除零
    """
    if xywh:
        (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1)
        w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
        b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_
        b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
    else:
        b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1)
        b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, -1)
        w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
        w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps

    # Inner boxes: 以中心为基准缩放
    cx1, cy1 = (b1_x1 + b1_x2) / 2, (b1_y1 + b1_y2) / 2
    cx2, cy2 = (b2_x1 + b2_x2) / 2, (b2_y1 + b2_y2) / 2

    inner_w1, inner_h1 = w1 * ratio, h1 * ratio
    inner_w2, inner_h2 = w2 * ratio, h2 * ratio

    inner_b1_x1 = cx1 - inner_w1 / 2
    inner_b1_x2 = cx1 + inner_w1 / 2
    inner_b1_y1 = cy1 - inner_h1 / 2
    inner_b1_y2 = cy1 + inner_h1 / 2

    inner_b2_x1 = cx2 - inner_w2 / 2
    inner_b2_x2 = cx2 + inner_w2 / 2
    inner_b2_y1 = cy2 - inner_h2 / 2
    inner_b2_y2 = cy2 + inner_h2 / 2

    # Inner intersection
    inner_inter = (
        (inner_b1_x2.minimum(inner_b2_x2) - inner_b1_x1.maximum(inner_b2_x1)).clamp_(0)
        * (inner_b1_y2.minimum(inner_b2_y2) - inner_b1_y1.maximum(inner_b2_y1)).clamp_(0)
    )

    # Inner union
    inner_union = inner_w1 * inner_h1 + inner_w2 * inner_h2 - inner_inter + eps

    # Inner IoU
    inner_iou = inner_inter / inner_union

    # CIoU penalties (用原始框计算)
    cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1)
    ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1)
    c2 = cw.pow(2) + ch.pow(2) + eps
    rho2 = (
        (b2_x1 + b2_x2 - b1_x1 - b1_x2).pow(2)
        + (b2_y1 + b2_y2 - b1_y1 - b1_y2).pow(2)
    ) / 4

    v = (4 / math.pi ** 2) * ((w2 / h2).atan() - (w1 / h1).atan()).pow(2)
    with torch.no_grad():
        alpha = v / (1 - inner_iou + v + eps)

    return inner_iou - (rho2 / c2 + v * alpha)


def focal_iou_weight(iou, gamma=0.5):
    """Focal weight: 对高IoU样本加大权重。"""
    return iou.detach().pow(gamma)


def _inner_iou_forward(self, pred_dist, pred_bboxes, anchor_points,
                        target_bboxes, target_scores, target_scores_sum,
                        fg_mask, imgsz, stride):
    """替换 BboxLoss.forward，使用 Inner-IoU + Focal weighting。"""
    from ultralytics.utils.tal import bbox2dist

    weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)

    # Inner-IoU (ratio=0.7 对小目标效果好)
    iou = inner_bbox_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask],
                          xywh=False, ratio=0.7)

    # Focal weighting: 高IoU样本获得更大权重
    focal_w = focal_iou_weight(iou.clamp(0, 1), gamma=0.5)
    loss_iou = ((1.0 - iou) * weight * focal_w).sum() / target_scores_sum

    # DFL loss (不变)
    if self.dfl_loss:
        target_ltrb = bbox2dist(anchor_points, target_bboxes, self.dfl_loss.reg_max - 1)
        loss_dfl = self.dfl_loss(pred_dist[fg_mask].view(-1, self.dfl_loss.reg_max),
                                  target_ltrb[fg_mask]) * weight
        loss_dfl = loss_dfl.sum() / target_scores_sum
    else:
        target_ltrb = bbox2dist(anchor_points, target_bboxes)
        target_ltrb = target_ltrb * stride
        target_ltrb[..., 0::2] /= imgsz[1]
        target_ltrb[..., 1::2] /= imgsz[0]
        pred_dist_s = pred_dist * stride
        pred_dist_s[..., 0::2] /= imgsz[1]
        pred_dist_s[..., 1::2] /= imgsz[0]
        loss_dfl = (
            F.l1_loss(pred_dist_s[fg_mask], target_ltrb[fg_mask], reduction="none")
            .mean(-1, keepdim=True) * weight
        )
        loss_dfl = loss_dfl.sum() / target_scores_sum

    return loss_iou, loss_dfl


def patch_inner_iou():
    """Patch BboxLoss 使用 Inner-IoU + Focal。"""
    import ultralytics.utils.loss as los
    los.BboxLoss.forward = _inner_iou_forward
    print("[Loss] Patched BboxLoss → Inner-IoU (ratio=0.7) + Focal (gamma=0.5)")