"""Faster R-CNN val集评估 —— 输出 mAP@50 / mAP@50-95 / P / R

用法：
  python Scripts/eval_fasterrcnn.py
"""
import torch
from torch.utils.data import DataLoader
from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchmetrics.detection.mean_ap import MeanAveragePrecision
from pathlib import Path
from PIL import Image
import torchvision.transforms.functional as F
from torch.utils.data import Dataset

DATA_ROOT = Path("Data/Detection_dataset")
CKPT = Path("Detection_experiments/compare_fasterrcnn/best.pt")
NUM_CLASSES = 2
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
CONF_THRESH = 0.5


class GoatDataset(Dataset):
    def __init__(self, split: str):
        img_dir = DATA_ROOT / "images" / split
        lbl_dir = DATA_ROOT / "labels" / split
        self.samples = []
        for img_path in sorted(img_dir.glob("*.jpg")):
            lbl_path = lbl_dir / (img_path.stem + ".txt")
            if lbl_path.exists():
                self.samples.append((img_path, lbl_path))

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        img_path, lbl_path = self.samples[idx]
        img = Image.open(img_path).convert("RGB")
        w, h = img.size
        img_tensor = F.to_tensor(img)

        boxes, labels = [], []
        for line in lbl_path.read_text().splitlines():
            parts = line.strip().split()
            if len(parts) != 5:
                continue
            _, cx, cy, bw, bh = map(float, parts)
            x1 = (cx - bw / 2) * w
            y1 = (cy - bh / 2) * h
            x2 = (cx + bw / 2) * w
            y2 = (cy + bh / 2) * h
            boxes.append([x1, y1, x2, y2])
            labels.append(1)

        target = {
            "boxes": torch.tensor(boxes, dtype=torch.float32) if boxes else torch.zeros((0, 4), dtype=torch.float32),
            "labels": torch.tensor(labels, dtype=torch.int64) if labels else torch.zeros(0, dtype=torch.int64),
        }
        return img_tensor, target


def collate_fn(batch):
    return tuple(zip(*batch))


def build_model():
    from torchvision.models import ResNet50_Weights
    model = fasterrcnn_resnet50_fpn_v2(weights=None, weights_backbone=ResNet50_Weights.IMAGENET1K_V2)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, NUM_CLASSES)
    return model


def main():
    model = build_model().to(DEVICE)
    model.load_state_dict(torch.load(CKPT, map_location=DEVICE))
    model.eval()

    val_ds = GoatDataset("val")
    val_loader = DataLoader(val_ds, batch_size=1, shuffle=False,
                            num_workers=2, collate_fn=collate_fn)

    metric = MeanAveragePrecision(iou_type="bbox", class_metrics=False)

    print(f"Evaluating on {len(val_ds)} val images...")
    with torch.no_grad():
        for imgs, targets in val_loader:
            imgs = [img.to(DEVICE) for img in imgs]
            outputs = model(imgs)

            preds = []
            for out in outputs:
                keep = out["scores"] >= CONF_THRESH
                preds.append({
                    "boxes":  out["boxes"][keep].cpu(),
                    "scores": out["scores"][keep].cpu(),
                    "labels": out["labels"][keep].cpu(),
                })

            gts = []
            for t in targets:
                gts.append({
                    "boxes":  t["boxes"],
                    "labels": t["labels"],
                })

            metric.update(preds, gts)

    result = metric.compute()
    map50_95 = result["map"].item()
    map50    = result["map_50"].item()
    mar      = result["mar_100"].item()

    print(f"\n{'='*45}")
    print(f"  Faster R-CNN (ImageNet backbone, 30 epochs)")
    print(f"{'='*45}")
    print(f"  mAP@50      : {map50:.3f}")
    print(f"  mAP@50-95   : {map50_95:.3f}")
    print(f"  mAR@100     : {mar:.3f}")
    print(f"  conf_thresh : {CONF_THRESH}")
    print(f"{'='*45}")


if __name__ == "__main__":
    main()