"""
Faster R-CNN 밑바닥 구현 — [2/5] 박스 연산 유틸
==================================================
객체탐지의 수학적 핵심이 모두 여기에 있다.

  1) 앵커(anchor) 생성      : 격자마다 여러 크기·비율의 기준 박스를 깐다
  2) IoU                    : 두 박스가 얼마나 겹치는가
  3) 박스 인코딩/디코딩     : (박스 → 회귀 타깃) / (예측값 → 박스)
  4) NMS                    : 겹치는 중복 예측을 제거

이 파일만 이해하면 Faster R-CNN의 절반을 이해한 것이다.
"""

import torch


# ---------------------------------------------------------------
# 1) 앵커 생성
# ---------------------------------------------------------------
def generate_anchors(base_size=16, ratios=(0.5, 1.0, 2.0),
                     scales=(8, 16, 32)):
    """
    한 격자점(cell)에 놓을 기준 앵커들을 만든다.
    ratios(가로세로비) × scales(크기) 조합 → 보통 9개 앵커.

    반환: [num_anchors, 4] 형태의 (x1,y1,x2,y2), 중심이 원점 기준.
    """
    anchors = []
    for scale in scales:
        area = (base_size * scale) ** 2
        for ratio in ratios:
            # 넓이는 유지하고 가로세로비만 바꾼다
            w = round((area / ratio) ** 0.5)
            h = round(w * ratio)
            anchors.append([-w / 2, -h / 2, w / 2, h / 2])
    return torch.tensor(anchors, dtype=torch.float32)


def shift_anchors(base_anchors, feat_h, feat_w, stride):
    """
    기준 앵커를 특징맵 전체 격자에 복제·이동시켜
    모든 위치의 앵커를 만든다.

    feat_h, feat_w : 특징맵 크기
    stride         : 원본 이미지 대비 특징맵 축소 배율(예: 16)
    반환: [feat_h*feat_w*num_anchors, 4] (원본 이미지 좌표계)
    """
    # 각 격자점의 이미지상 중심 좌표
    shift_x = (torch.arange(feat_w) + 0.5) * stride
    shift_y = (torch.arange(feat_h) + 0.5) * stride
    sy, sx = torch.meshgrid(shift_y, shift_x, indexing="ij")
    shifts = torch.stack([sx.reshape(-1), sy.reshape(-1),
                          sx.reshape(-1), sy.reshape(-1)], dim=1)  # [K,4]

    # [K,1,4] + [1,A,4] → [K,A,4] → [K*A,4]
    anchors = shifts[:, None, :] + base_anchors[None, :, :]
    return anchors.reshape(-1, 4)


# ---------------------------------------------------------------
# 2) IoU (Intersection over Union)
# ---------------------------------------------------------------
def box_iou(boxes1, boxes2):
    """
    [N,4], [M,4] → [N,M] IoU 행렬.
    IoU = 교집합 넓이 / 합집합 넓이. 0(안 겹침)~1(완전 일치).
    """
    area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
    area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])

    lt = torch.max(boxes1[:, None, :2], boxes2[None, :, :2])  # 교집합 좌상단
    rb = torch.min(boxes1[:, None, 2:], boxes2[None, :, 2:])  # 교집합 우하단
    wh = (rb - lt).clamp(min=0)
    inter = wh[:, :, 0] * wh[:, :, 1]

    union = area1[:, None] + area2[None, :] - inter
    return inter / union.clamp(min=1e-6)


# ---------------------------------------------------------------
# 3) 박스 인코딩 / 디코딩
# ---------------------------------------------------------------
def encode_boxes(gt, anchors):
    """
    정답 박스(gt)를 앵커 기준 회귀 타깃 (dx,dy,dw,dh)으로 변환.
    네트워크는 절대 좌표가 아니라 "앵커로부터의 상대 변형"을 배운다.
    """
    aw = anchors[:, 2] - anchors[:, 0]
    ah = anchors[:, 3] - anchors[:, 1]
    ax = anchors[:, 0] + 0.5 * aw
    ay = anchors[:, 1] + 0.5 * ah

    gw = gt[:, 2] - gt[:, 0]
    gh = gt[:, 3] - gt[:, 1]
    gx = gt[:, 0] + 0.5 * gw
    gy = gt[:, 1] + 0.5 * gh

    dx = (gx - ax) / aw
    dy = (gy - ay) / ah
    dw = torch.log(gw / aw)
    dh = torch.log(gh / ah)
    return torch.stack([dx, dy, dw, dh], dim=1)


def decode_boxes(deltas, anchors):
    """
    네트워크가 예측한 (dx,dy,dw,dh)를 실제 박스 좌표로 복원.
    encode_boxes의 역연산.
    """
    aw = anchors[:, 2] - anchors[:, 0]
    ah = anchors[:, 3] - anchors[:, 1]
    ax = anchors[:, 0] + 0.5 * aw
    ay = anchors[:, 1] + 0.5 * ah

    dx, dy, dw, dh = deltas[:, 0], deltas[:, 1], deltas[:, 2], deltas[:, 3]
    # dw,dh 폭주 방지 클램프
    dw = torch.clamp(dw, max=4.135)
    dh = torch.clamp(dh, max=4.135)

    px = dx * aw + ax
    py = dy * ah + ay
    pw = torch.exp(dw) * aw
    ph = torch.exp(dh) * ah

    x1 = px - 0.5 * pw
    y1 = py - 0.5 * ph
    x2 = px + 0.5 * pw
    y2 = py + 0.5 * ph
    return torch.stack([x1, y1, x2, y2], dim=1)


def clip_boxes(boxes, img_h, img_w):
    """박스를 이미지 경계 안으로 자른다."""
    boxes[:, 0].clamp_(min=0, max=img_w)
    boxes[:, 1].clamp_(min=0, max=img_h)
    boxes[:, 2].clamp_(min=0, max=img_w)
    boxes[:, 3].clamp_(min=0, max=img_h)
    return boxes


# ---------------------------------------------------------------
# 4) NMS (Non-Maximum Suppression)
# ---------------------------------------------------------------
def nms(boxes, scores, iou_thresh=0.7):
    """
    점수 높은 박스부터 남기고, 그와 많이 겹치는 박스는 제거.
    torchvision.ops.nms 를 써도 되지만, 원리 학습용으로 직접 구현.
    반환: 남길 인덱스.
    """
    if boxes.numel() == 0:
        return torch.empty((0,), dtype=torch.int64)

    order = scores.argsort(descending=True)
    keep = []
    while order.numel() > 0:
        i = order[0].item()
        keep.append(i)
        if order.numel() == 1:
            break
        ious = box_iou(boxes[i].unsqueeze(0), boxes[order[1:]]).squeeze(0)
        # 임계값 이하만 남긴다
        order = order[1:][ious <= iou_thresh]
    return torch.tensor(keep, dtype=torch.int64)