"""
Faster R-CNN 밑바닥 구현 — [4/5] 손실 계산
=============================================
Faster R-CNN은 두 곳에서 손실이 발생한다.

  RPN 손실 : 앵커가 객체인가(분류) + 앵커→정답 박스 보정(회귀)
  RoI 손실 : 후보영역의 클래스(분류) + 클래스별 박스 보정(회귀)

핵심은 "타깃 할당(target assignment)":
  - 각 앵커/후보영역이 어떤 정답 박스를 담당하는지 IoU로 정한다.
  - IoU 높으면 positive(객체), 낮으면 negative(배경), 애매하면 무시.
  - 그 다음 positive에만 회귀 손실, pos+neg에 분류 손실을 건다.
"""

import torch
import torch.nn.functional as F

from box_utils import box_iou, encode_boxes, decode_boxes, clip_boxes


def _sample(pos_idx, neg_idx, num, pos_frac):
    """pos/neg를 정해진 개수·비율로 무작위 샘플링(클래스 불균형 방지)."""
    num_pos = min(int(num * pos_frac), pos_idx.numel())
    num_neg = min(num - num_pos, neg_idx.numel())
    pos = pos_idx[torch.randperm(pos_idx.numel())[:num_pos]]
    neg = neg_idx[torch.randperm(neg_idx.numel())[:num_neg]]
    return pos, neg


# ---------------------------------------------------------------
# RPN 손실
# ---------------------------------------------------------------
def rpn_loss(rpn_logits, rpn_deltas, anchors, gt_boxes, img_hw,
             pos_iou=0.7, neg_iou=0.3, num_samples=256, pos_frac=0.5):
    device = rpn_logits.device

    # 이미지 밖으로 나간 앵커는 학습에서 제외
    inside = ((anchors[:, 0] >= 0) & (anchors[:, 1] >= 0) &
              (anchors[:, 2] <= img_hw[1]) & (anchors[:, 3] <= img_hw[0]))
    idx_inside = torch.where(inside)[0]
    anc = anchors[idx_inside]

    labels = torch.full((anc.shape[0],), -1, dtype=torch.float32, device=device)  # -1=무시

    if gt_boxes.numel() > 0:
        ious = box_iou(anc, gt_boxes)          # [A, G]
        max_iou, argmax = ious.max(dim=1)       # 각 앵커의 최고 IoU 정답
        labels[max_iou < neg_iou] = 0           # 배경
        labels[max_iou >= pos_iou] = 1          # 객체
        # 각 정답 박스에 대해 IoU 최대인 앵커도 강제로 positive
        gt_best = ious.argmax(dim=0)
        labels[gt_best] = 1
        matched_gt = gt_boxes[argmax]
    else:
        labels[:] = 0
        matched_gt = torch.zeros_like(anc)

    pos = torch.where(labels == 1)[0]
    neg = torch.where(labels == 0)[0]
    pos, neg = _sample(pos, neg, num_samples, pos_frac)
    samp = torch.cat([pos, neg])

    # --- 분류 손실 (객체/배경) ---
    logits_inside = rpn_logits[idx_inside]
    cls_loss = F.binary_cross_entropy_with_logits(
        logits_inside[samp], labels[samp])

    # --- 회귀 손실 (positive 앵커만) ---
    if pos.numel() > 0:
        deltas_inside = rpn_deltas[idx_inside]
        reg_targets = encode_boxes(matched_gt[pos], anc[pos])
        reg_loss = F.smooth_l1_loss(deltas_inside[pos], reg_targets,
                                    beta=1.0 / 9.0)
    else:
        reg_loss = torch.tensor(0.0, device=device)

    return cls_loss + reg_loss


# ---------------------------------------------------------------
# RoI 타깃 할당 (후보영역 → 학습 샘플)
# ---------------------------------------------------------------
def assign_roi_targets(proposals, gt_boxes, gt_labels,
                       pos_iou=0.5, neg_iou_hi=0.5, neg_iou_lo=0.0,
                       num_samples=128, pos_frac=0.25):
    """후보영역에 클래스 라벨과 회귀 타깃을 붙이고 샘플링."""
    device = proposals.device
    # 정답 박스도 후보에 추가(학습 초기 안정화)
    if gt_boxes.numel() > 0:
        proposals = torch.cat([proposals, gt_boxes], dim=0)

    if gt_boxes.numel() == 0:
        # 정답이 없으면 전부 배경
        n = min(num_samples, proposals.shape[0])
        sel = proposals[:n]
        labels = torch.zeros((n,), dtype=torch.int64, device=device)
        reg_t = torch.zeros((n, 4), device=device)
        return sel, labels, reg_t, torch.zeros((n,), dtype=torch.bool, device=device)

    ious = box_iou(proposals, gt_boxes)
    max_iou, argmax = ious.max(dim=1)
    gt_for_prop = gt_labels[argmax]
    matched_gt = gt_boxes[argmax]

    labels = torch.zeros_like(gt_for_prop)          # 0=배경 기본
    pos_mask = max_iou >= pos_iou
    labels[pos_mask] = gt_for_prop[pos_mask]         # 객체 클래스 부여

    pos = torch.where(pos_mask)[0]
    neg = torch.where((max_iou < neg_iou_hi) & (max_iou >= neg_iou_lo))[0]
    pos, neg = _sample(pos, neg, num_samples, pos_frac)
    samp = torch.cat([pos, neg])

    sel_prop = proposals[samp]
    sel_labels = labels[samp]
    reg_targets = encode_boxes(matched_gt[samp], sel_prop)
    is_pos = torch.zeros((samp.numel(),), dtype=torch.bool, device=device)
    is_pos[:pos.numel()] = True
    return sel_prop, sel_labels, reg_targets, is_pos


# ---------------------------------------------------------------
# RoI 손실
# ---------------------------------------------------------------
def roi_loss(head, feat, stride, proposals, gt_boxes, gt_labels):
    sel_prop, labels, reg_targets, is_pos = assign_roi_targets(
        proposals, gt_boxes, gt_labels)

    cls_logits, reg = head(feat, sel_prop, stride)  # [S,C], [S,C*4]

    # --- 분류 손실 (전체 샘플) ---
    cls_loss = F.cross_entropy(cls_logits, labels)

    # --- 회귀 손실 (positive만, 해당 클래스의 4개 좌표만) ---
    if is_pos.sum() > 0:
        S = reg.shape[0]
        reg = reg.reshape(S, -1, 4)
        pos_idx = torch.where(is_pos)[0]
        pos_labels = labels[pos_idx]
        reg_pos = reg[pos_idx, pos_labels]           # 해당 클래스 좌표
        reg_loss = F.smooth_l1_loss(reg_pos, reg_targets[pos_idx], beta=1.0)
    else:
        reg_loss = torch.tensor(0.0, device=feat.device)

    return cls_loss + reg_loss