"""
Faster R-CNN 밑바닥 구현 — 추론(inference) 스크립트
=====================================================
학습된 frcnn.pth 로 임의의 이미지에서 객체를 탐지하고,
박스 + 클래스명 + 점수를 그려서 저장한다.

실행 예:
  # 이미지 한 장
  python infer.py --ckpt frcnn.pth --image test.jpg

  # 폴더 안 모든 이미지
  python infer.py --ckpt frcnn.pth --image_dir ./samples --out_dir ./results

  # 점수 임계값 조정(기본 0.5)
  python infer.py --ckpt frcnn.pth --image test.jpg --score_thresh 0.7

주의:
  - train.py, model.py 등과 같은 폴더에서 실행할 것.
  - 학습과 동일한 리사이즈/정규화를 적용해야 결과가 정상.
"""

import os
import argparse

import torch
from PIL import Image, ImageDraw, ImageFont
import torchvision.transforms.functional as F

from model import FasterRCNN
from dataset import NUM_CLASSES, VOC_CLASSES


# 클래스별 색상(20개) — 시각적으로 구분되도록 HSV 분할
def _class_colors():
    import colorsys
    colors = []
    for i in range(len(VOC_CLASSES)):
        h = i / len(VOC_CLASSES)
        r, g, b = colorsys.hsv_to_rgb(h, 0.75, 0.95)
        colors.append((int(r * 255), int(g * 255), int(b * 255)))
    return colors

CLASS_COLORS = _class_colors()


def preprocess(pil_img, min_size=600, max_size=1000):
    """학습과 동일한 리사이즈 + 정규화. 원본 복원용 scale도 반환."""
    w, h = pil_img.size
    short, long = min(w, h), max(w, h)
    scale = min_size / short
    if long * scale > max_size:
        scale = max_size / long
    new_w, new_h = int(round(w * scale)), int(round(h * scale))

    resized = pil_img.resize((new_w, new_h), Image.BILINEAR)
    t = F.to_tensor(resized)
    t = F.normalize(t, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    return t, scale


def draw_detections(pil_img, boxes, labels, scores, scale):
    """탐지 결과를 원본 이미지 좌표로 되돌려 박스와 라벨을 그린다."""
    draw = ImageDraw.Draw(pil_img)
    try:
        font = ImageFont.truetype("arial.ttf", 16)
    except Exception:
        font = ImageFont.load_default()

    for box, label, score in zip(boxes, labels, scores):
        # 모델은 리사이즈된 좌표를 출력 → 원본 크기로 되돌림(÷scale)
        x1, y1, x2, y2 = (box / scale).tolist()
        cls_idx = int(label) - 1  # 0=배경 제외
        if cls_idx < 0 or cls_idx >= len(VOC_CLASSES):
            continue
        name = VOC_CLASSES[cls_idx]
        color = CLASS_COLORS[cls_idx]

        # 박스
        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)

        # 라벨 배경 + 텍스트
        text = f"{name} {score:.2f}"
        tb = draw.textbbox((x1, y1), text, font=font)
        draw.rectangle([tb[0], tb[1], tb[2], tb[3]], fill=color)
        draw.text((x1, y1), text, fill="white", font=font)

    return pil_img


@torch.no_grad()
def infer_image(model, img_path, device, score_thresh):
    pil = Image.open(img_path).convert("RGB")
    tensor, scale = preprocess(pil)
    tensor = tensor.to(device).unsqueeze(0)

    det = model(tensor)  # eval 모드 → {boxes, labels, scores}
    keep = det["scores"] >= score_thresh
    boxes = det["boxes"][keep].cpu()
    labels = det["labels"][keep].cpu()
    scores = det["scores"][keep].cpu()

    result = draw_detections(pil, boxes, labels, scores, scale)
    return result, len(boxes)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True, help="학습된 가중치 (frcnn.pth)")
    ap.add_argument("--image", help="단일 이미지 경로")
    ap.add_argument("--image_dir", help="이미지 폴더 경로")
    ap.add_argument("--out_dir", default="./results", help="결과 저장 폴더")
    ap.add_argument("--score_thresh", type=float, default=0.5,
                    help="이 점수 이상만 표시")
    args = ap.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("device:", device)

    model = FasterRCNN(NUM_CLASSES).to(device)
    model.load_state_dict(torch.load(args.ckpt, map_location=device))
    model.eval()
    print(f"모델 로드 완료: {args.ckpt}")

    os.makedirs(args.out_dir, exist_ok=True)

    # 처리할 이미지 목록 구성
    targets = []
    if args.image:
        targets.append(args.image)
    if args.image_dir:
        for fn in os.listdir(args.image_dir):
            if fn.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
                targets.append(os.path.join(args.image_dir, fn))

    if not targets:
        print("이미지를 지정하세요: --image 또는 --image_dir")
        return

    for path in targets:
        result, n = infer_image(model, path, device, args.score_thresh)
        out_path = os.path.join(args.out_dir, "det_" + os.path.basename(path))
        result.save(out_path)
        print(f"  {os.path.basename(path)}: {n}개 탐지 → {out_path}")

    print(f"완료. 결과는 {args.out_dir} 폴더에 저장됨.")


if __name__ == "__main__":
    main()