"""用 YOLO 格式数据集训练 torchvision Faster R-CNN。"""

from __future__ import annotations

import argparse
import time
from pathlib import Path

import cv2
import numpy as np
import torch
import yaml
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

from detectors import build_faster_rcnn, resolve_device

ROOT = Path(__file__).resolve().parent


class YoloDetDataset(Dataset):
    """读取 Ultralytics YOLO txt 标签（class cx cy w h，归一化）。"""

    def __init__(self, img_dir: Path, label_dir: Path, imgsz: int = 640) -> None:
        self.img_paths = sorted(
            [p for p in img_dir.iterdir() if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".bmp"}]
        )
        self.label_dir = label_dir
        self.imgsz = int(imgsz)

    def __len__(self) -> int:
        return len(self.img_paths)

    def __getitem__(self, idx: int):
        path = self.img_paths[idx]
        bgr = cv2.imread(str(path))
        if bgr is None:
            raise RuntimeError(f"无法读取图像: {path}")
        h0, w0 = bgr.shape[:2]
        scale = self.imgsz / max(h0, w0)
        nh, nw = max(1, int(round(h0 * scale))), max(1, int(round(w0 * scale)))
        bgr = cv2.resize(bgr, (nw, nh), interpolation=cv2.INTER_LINEAR)
        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        image = torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0

        label_path = self.label_dir / f"{path.stem}.txt"
        boxes = []
        labels = []
        if label_path.exists():
            for line in label_path.read_text(encoding="utf-8").splitlines():
                parts = line.strip().split()
                if len(parts) < 5:
                    continue
                cls_id = int(float(parts[0]))
                cx, cy, bw, bh = map(float, parts[1:5])
                x1 = (cx - bw / 2.0) * nw
                y1 = (cy - bh / 2.0) * nh
                x2 = (cx + bw / 2.0) * nw
                y2 = (cy + bh / 2.0) * nh
                x1 = float(np.clip(x1, 0, nw - 1))
                y1 = float(np.clip(y1, 0, nh - 1))
                x2 = float(np.clip(x2, 0, nw - 1))
                y2 = float(np.clip(y2, 0, nh - 1))
                if x2 <= x1 or y2 <= y1:
                    continue
                boxes.append([x1, y1, x2, y2])
                # Faster R-CNN 前景从 1 开始
                labels.append(cls_id + 1)

        if boxes:
            boxes_t = torch.tensor(boxes, dtype=torch.float32)
            labels_t = torch.tensor(labels, dtype=torch.int64)
        else:
            boxes_t = torch.zeros((0, 4), dtype=torch.float32)
            labels_t = torch.zeros((0,), dtype=torch.int64)

        target = {
            "boxes": boxes_t,
            "labels": labels_t,
            "image_id": torch.tensor([idx]),
            "area": (boxes_t[:, 3] - boxes_t[:, 1]) * (boxes_t[:, 2] - boxes_t[:, 0])
            if len(boxes_t)
            else torch.zeros((0,), dtype=torch.float32),
            "iscrowd": torch.zeros((len(boxes_t),), dtype=torch.int64),
        }
        return image, target


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


def load_data_yaml(path: Path) -> dict:
    data = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
    root = Path(data["path"])
    if not root.is_absolute():
        root = (path.parent / root).resolve()
    return {
        "root": root,
        "train": root / data.get("train", "images/train"),
        "val": root / data.get("val", "images/val"),
        "names": data.get("names", {0: "object"}),
    }


def load_cfg(path: Path) -> dict:
    cfg = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
    return cfg


@torch.inference_mode()
def evaluate(model, loader, device) -> dict[str, float]:
    model.eval()
    total = 0
    hit = 0
    score_sum = 0.0
    for images, targets in loader:
        images = [img.to(device) for img in images]
        outputs = model(images)
        for out, tgt in zip(outputs, targets):
            total += 1
            gt = tgt["boxes"]
            if len(gt) == 0:
                continue
            if len(out["boxes"]) == 0:
                continue
            # 简化指标：最高分框与任一 GT IoU>0.5 记命中
            scores = out["scores"].detach().cpu()
            boxes = out["boxes"].detach().cpu()
            best = int(torch.argmax(scores))
            score_sum += float(scores[best])
            pb = boxes[best]
            for gb in gt:
                xx1 = max(float(pb[0]), float(gb[0]))
                yy1 = max(float(pb[1]), float(gb[1]))
                xx2 = min(float(pb[2]), float(gb[2]))
                yy2 = min(float(pb[3]), float(gb[3]))
                inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
                area_p = max(0.0, float(pb[2] - pb[0])) * max(0.0, float(pb[3] - pb[1]))
                area_g = max(0.0, float(gb[2] - gb[0])) * max(0.0, float(gb[3] - gb[1]))
                union = area_p + area_g - inter + 1e-6
                if inter / union >= 0.5:
                    hit += 1
                    break
    return {
        "val_images": float(total),
        "recall_proxy": hit / max(1, total),
        "avg_top_score": score_sum / max(1, total),
    }


def train_one_epoch(model, optimizer, loader, device, print_freq: int) -> float:
    model.train()
    loss_meter = 0.0
    n = 0
    pbar = tqdm(loader, desc="train", leave=False)
    for i, (images, targets) in enumerate(pbar):
        images = [img.to(device) for img in images]
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())
        optimizer.zero_grad(set_to_none=True)
        losses.backward()
        optimizer.step()
        loss_meter += float(losses.detach().cpu())
        n += 1
        if (i + 1) % print_freq == 0:
            pbar.set_postfix(loss=f"{loss_meter / n:.4f}")
    return loss_meter / max(1, n)


def main() -> None:
    parser = argparse.ArgumentParser(description="Train Faster R-CNN on YOLO-format dataset")
    parser.add_argument("--cfg", type=str, default=str(ROOT / "faster_rcnn.yaml"))
    args = parser.parse_args()

    cfg = load_cfg(Path(args.cfg))
    data_yaml = Path(cfg.get("data", "01data.yaml"))
    if not data_yaml.is_absolute():
        data_yaml = ROOT / data_yaml
    data = load_data_yaml(data_yaml)

    train_img = Path(data["train"])
    val_img = Path(data["val"])
    # YOLO layout: images/train -> labels/train
    train_lbl = train_img.parent.parent / "labels" / train_img.name
    val_lbl = val_img.parent.parent / "labels" / val_img.name
    if not train_lbl.exists():
        train_lbl = Path(str(train_img).replace("images", "labels"))
    if not val_lbl.exists():
        val_lbl = Path(str(val_img).replace("images", "labels"))

    imgsz = int(cfg.get("imgsz", 640))
    train_set = YoloDetDataset(train_img, train_lbl, imgsz=imgsz)
    val_set = YoloDetDataset(val_img, val_lbl, imgsz=imgsz)
    if len(train_set) == 0:
        raise RuntimeError(f"训练集为空: {train_img}")

    device = resolve_device(cfg.get("device", "0"))
    num_classes = int(cfg.get("num_classes", 2))
    model = build_faster_rcnn(num_classes, pretrained_backbone=bool(cfg.get("pretrained", True)))
    model.to(device)

    resume = str(cfg.get("resume") or "").strip()
    start_epoch = 1
    if resume:
        ckpt = torch.load(resume, map_location="cpu", weights_only=False)
        model.load_state_dict(ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt)
        start_epoch = int(ckpt.get("epoch", 0)) + 1 if isinstance(ckpt, dict) else 1
        print(f"已从 {resume} 恢复，下一轮 epoch={start_epoch}")

    workers = int(cfg.get("num_workers", 2))
    batch = int(cfg.get("batch_size", 2))
    train_loader = DataLoader(
        train_set,
        batch_size=batch,
        shuffle=True,
        num_workers=workers,
        collate_fn=collate_fn,
        pin_memory=device.type == "cuda",
    )
    val_loader = DataLoader(
        val_set,
        batch_size=1,
        shuffle=False,
        num_workers=max(0, workers // 2),
        collate_fn=collate_fn,
    )

    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(
        params,
        lr=float(cfg.get("lr", 0.005)),
        momentum=float(cfg.get("momentum", 0.9)),
        weight_decay=float(cfg.get("weight_decay", 0.0005)),
    )
    lr_scheduler = torch.optim.lr_scheduler.StepLR(
        optimizer,
        step_size=int(cfg.get("step_size", 12)),
        gamma=float(cfg.get("gamma", 0.1)),
    )

    out_dir = ROOT / cfg.get("project", "runs/faster_rcnn") / cfg.get("name", "exp")
    if out_dir.exists() and not bool(cfg.get("exist_ok", True)):
        raise RuntimeError(f"输出目录已存在: {out_dir}")
    out_dir.mkdir(parents=True, exist_ok=True)
    weights_dir = out_dir / "weights"
    weights_dir.mkdir(exist_ok=True)

    epochs = int(cfg.get("epochs", 30))
    print_freq = int(cfg.get("print_freq", 20))
    eval_every = int(cfg.get("eval_every", 1))
    best_score = -1.0

    print(f"device={device}  train={len(train_set)}  val={len(val_set)}  out={out_dir}")
    for epoch in range(start_epoch, epochs + 1):
        t0 = time.perf_counter()
        avg_loss = train_one_epoch(model, optimizer, train_loader, device, print_freq)
        lr_scheduler.step()
        metrics = {"loss": avg_loss}
        if epoch % eval_every == 0 and len(val_set) > 0:
            metrics.update(evaluate(model, val_loader, device))
        elapsed = time.perf_counter() - t0
        print(
            f"epoch {epoch}/{epochs}  loss={avg_loss:.4f}  "
            f"recall_proxy={metrics.get('recall_proxy', float('nan')):.3f}  "
            f"time={elapsed:.1f}s"
        )

        ckpt = {
            "epoch": epoch,
            "model": model.state_dict(),
            "num_classes": num_classes,
            "names": data["names"],
            "cfg": cfg,
            "metrics": metrics,
        }
        torch.save(ckpt, weights_dir / "last.pth")
        if epoch % 5 == 0:
            torch.save(ckpt, weights_dir / f"epoch{epoch}.pth")
        score = float(metrics.get("recall_proxy", 0.0))
        if score >= best_score:
            best_score = score
            torch.save(ckpt, weights_dir / "best.pth")
            print(f"  saved best.pth (recall_proxy={best_score:.3f})")

    print(f"训练完成。权重目录: {weights_dir}")


if __name__ == "__main__":
    main()