"""Wise-IoU v1 动态聚焦权重 — 正确的 ultralytics 8.4.46 集成方式

ultralytics 8.4.46 的 BboxLoss.forward 直接用 iou 标量：
  loss_iou = ((1.0 - iou) * weight).sum() / target_scores_sum
不支持元组返回。因此 patch BboxLoss.forward，在 weight 中融入 β。
"""
import torch
import torch.nn.functional as F


def _wise_beta(box1, box2, eps=1e-7):
    """计算宽高比偏差动态聚焦系数 β，归一化使 batch 均值≈1。"""
    b1_x1, b1_y1, b1_x2, b1_y2 = box1.unbind(-1)
    b2_x1, b2_y1, b2_x2, b2_y2 = box2.unbind(-1)
    w1 = (b1_x2 - b1_x1).clamp(min=eps)
    h1 = (b1_y2 - b1_y1).clamp(min=eps)
    w2 = (b2_x2 - b2_x1).clamp(min=eps)
    h2 = (b2_y2 - b2_y1).clamp(min=eps)
    rw = (w1 / w2).clamp(0.1, 10)
    rh = (h1 / h2).clamp(0.1, 10)
    raw = (rw - 1) ** 2 + (rh - 1) ** 2
    return (raw / (raw.mean() + eps)).detach().unsqueeze(-1)  # (N, 1)


def patch_wise_iou():
    """Patch BboxLoss.forward 以融入 Wise-IoU β 动态权重。

    只在 BboxLoss.forward 中把 weight 乘以 β，其余逻辑不变。
    不替换 bbox_iou，依然使用 CIoU 作为 iou 值（更稳定）。
    """
    import ultralytics.utils.loss as los
    from ultralytics.utils.loss import bbox_iou, bbox2dist, DFLoss
    from ultralytics.utils.tal import dist2bbox

    def _wise_forward(
        self,
        pred_dist,
        pred_bboxes,
        anchor_points,
        target_bboxes,
        target_scores,
        target_scores_sum,
        fg_mask,
        imgsz,
        stride,
    ):
        weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)
        beta = _wise_beta(pred_bboxes[fg_mask], target_bboxes[fg_mask])
        weighted = weight * beta                                    # Wise-IoU 动态权重
        iou = bbox_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask], xywh=False, CIoU=True)
        loss_iou = ((1.0 - iou) * weighted).sum() / target_scores_sum

        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 = pred_dist * stride
            pred_dist[..., 0::2] /= imgsz[1]
            pred_dist[..., 1::2] /= imgsz[0]
            loss_dfl = (
                F.l1_loss(pred_dist[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

    los.BboxLoss.forward = _wise_forward
    print("[WiseIoU] BboxLoss.forward 已替换为 Wise-IoU v1 动态权重版本")