#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
人体身份重识别（Person ReID）对比学习训练脚本（单机单文件）。
- Dataset：Market-1501 风格文件名解析 pid/cam；
- PK 采样器：每个 batch 内 P 个身份 × 每张 K 样本；
- 模型：ImageNet ResNet50 骨干 + GAP + 512 维嵌入 + ID 分类头；
- 损失：Batch-Hard Triplet + ID（交叉熵）联合；
"""
from __future__ import annotations

import argparse
import random
import re
from pathlib import Path
from typing import Iterator

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader, Dataset, Sampler
from torchvision import models, transforms

# ---------------------------------------------------------------------------
# 文件名解析：例如 24502_c1_s1_00001.jpg → pid=24502，cam_id=1
# ---------------------------------------------------------------------------

_NAME_RE = re.compile(
    r"^(?P<pid>\d+)_c(?P<cam>\d+)_",
    flags=re.I,
)


def parse_market1501_style_name(stem: str) -> tuple[int | None, int | None]:
    """从「不含后缀」的文件名 stem 中提取身份 ID、机位 ID。"""
    m = _NAME_RE.match(stem)
    if not m:
        return None, None
    return int(m.group("pid")), int(m.group("cam"))


class DoctorReIDDataset(Dataset):
    """医生 ReID：解析 pid/cam，将 pid 重映射到 0..num_classes-1。"""

    def __init__(self, image_root: Path, augment: bool) -> None:
        self.image_root = Path(image_root).resolve()
        exts = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
        paths = sorted(
            p
            for p in self.image_root.rglob("*")
            if p.is_file() and p.suffix.lower() in exts
        )
        pid_raw_list: list[int] = []
        cam_raw_list: list[int] = []
        valid_paths: list[Path] = []
        for p in paths:
            pid_raw, cam_raw = parse_market1501_style_name(p.stem)
            if pid_raw is None:
                continue
            valid_paths.append(p)
            pid_raw_list.append(pid_raw)
            cam_raw_list.append(int(cam_raw))

        unique_pids = sorted(set(pid_raw_list))
        self.pid_to_label = {pid: i for i, pid in enumerate(unique_pids)}
        self.labels: list[int] = [self.pid_to_label[r] for r in pid_raw_list]
        self.cam_ids: list[int] = cam_raw_list
        self.paths = valid_paths

        if len(self.paths) == 0:
            raise RuntimeError(f"目录下未发现有效图像: {self.image_root}")

        # Resize(128,256) 在 torchvision 中为 (height, width) → (256,128) 常见于 ReID。
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        if augment:
            self.transform = transforms.Compose(
                [
                    transforms.Resize((256, 128)),
                    transforms.RandomHorizontalFlip(p=0.5),
                    transforms.ColorJitter(
                        brightness=0.15, contrast=0.15, saturation=0.15, hue=0.02
                    ),
                    transforms.ToTensor(),
                    normalize,
                    # RandomErasure 需作用在 Tensor 上，故放在 ToTensor 之后；
                    transforms.RandomErasing(p=0.5, scale=(0.02, 0.25), ratio=(0.3, 3.3)),
                ]
            )
        else:
            self.transform = transforms.Compose(
                [
                    transforms.Resize((256, 128)),
                    transforms.ToTensor(),
                    normalize,
                ]
            )

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

    def __getitem__(self, idx: int):
        img = Image.open(self.paths[idx]).convert("RGB")
        x = self.transform(img)
        y = torch.tensor(self.labels[idx], dtype=torch.long)
        cam = torch.tensor(self.cam_ids[idx], dtype=torch.long)
        return x, y, cam


class RandomIdentityPKSampler(Sampler[int]):
    """
    PK / Random Identity Sampler（Random Identity Sampling）

    【原理简述】对比学习三元组需在 batch 内出现「同类多图」才能把 anchor / positive 找出来。
    PK 策略：在每个 mini-batch 中固定结构为：
        - 先随机选 **P** 个不同身份；
        - 每个身份若无放回地抽 **K** 张图像（样本不足则从该身份放回随机抽满 K）；
            则 batch_size = P * K。
    这样保证了每个身份在同一个 batch 中至少有若干张可用于 Triplet，
    Batch-Hard 才能选「最难的正对 / 最难的负样本」。
    【注意】P 不能超过数据集中可用身份总数；若 K > 该类张数则用放回采样。
    """

    def __init__(
        self,
        labels: list[int],
        p: int,
        k: int,
        *,
        seed: int = 0,
        length: int | None = None,
    ) -> None:
        super().__init__()
        self.labels_np = np.array(labels)
        self.p = int(p)
        self.k = int(k)
        self.seed = seed
        self.epoch = 0

        self.identities = sorted(np.unique(self.labels_np).tolist())
        self.num_identities = len(self.identities)

        idx_by_label: dict[int, list[int]] = {}
        for i, lbl in enumerate(self.labels_np.tolist()):
            idx_by_label.setdefault(int(lbl), []).append(i)
        self.idx_by_label = {k_: np.array(v) for k_, v in idx_by_label.items()}

        if self.num_identities == 0:
            raise RuntimeError("PKSampler: 无任何身份类别")

        if self.p > self.num_identities:
            raise ValueError(
                f"P={self.p} 大于数据集身份数 {self.num_identities}。"
                " 请减小 --p。"
            )

        # 每个 epoch 内「迭代步数」：batch 数（每次 batch 含 P×K 张图）
        self.labels_flat = labels
        if length is None:
            # 约按「扫过一遍身份组合」的量级设一个稳定值
            self.num_batches_est = max(32, min(200, len(self.labels_flat) // max(1, self.p * self.k)))
        else:
            self.num_batches_est = int(length)

        self.batch_size = self.p * self.k

    def __len__(self) -> int:
        return self.num_batches_est * self.batch_size

    def set_epoch(self, epoch: int) -> None:
        self.epoch = epoch

    def __iter__(self) -> Iterator[int]:
        rng = np.random.RandomState((self.epoch * 9973 + self.seed) & 0xFFFFFFFF)
        ids_arr = np.asarray(self.identities, dtype=np.int64)

        for _batch_id in range(self.num_batches_est):
            # 一步：无放回抽 P 个身份（若身份不够则允许放回，实际 5 类且 P≤5 时总是无放回）
            if ids_arr.size >= self.p:
                chosen = rng.choice(ids_arr, size=self.p, replace=False)
            else:
                chosen = rng.choice(ids_arr, size=self.p, replace=True)

            out: list[int] = []
            for pid_pick in chosen.tolist():
                pool = self.idx_by_label[int(pid_pick)]
                if pool.size >= self.k:
                    idx_pick = rng.choice(pool, size=self.k, replace=False)
                else:
                    idx_pick = rng.choice(pool, size=self.k, replace=True)
                out.extend(int(t) for t in idx_pick)

            yield from out


def batch_hard_triplet_loss(embeddings: torch.Tensor, labels: torch.Tensor, margin: float) -> torch.Tensor:
    """
    Batch-Hard Triplet Loss（Hermans ECCV17 一类的标准形式）

    对每个 anchor（batch 里的每个样本 i）在同一 batch 中选：
      - Positive：与它同身份的样本 j 里面，距离**最大**的那一个（ hardest positive ）；
      - Negative：与它不同身份的样本里，距离**最小**的那一个（ hardest negative ）；
      单项损失通常为 relu(d_pos_hard - d_neg_hard + margin) 对每个有效 anchor 求平均。
    【数学含义】拉大「最难正对」相对「最易负样本」的间隔。
    【实现】使用欧氏距离；若某样本在 batch 内无「异类」（极少见）或无「同类第二张」则不参与均值。
    """
    dist = torch.cdist(embeddings.float(), embeddings.float(), p=2.0).clamp(min=1e-8)

    bs = dist.size(0)
    lbl = labels.long()
    same = lbl.unsqueeze(1).eq(lbl.unsqueeze(0))
    losses: list[torch.Tensor] = []

    for i in range(bs):
        same_pos = same[i].clone()
        same_pos[i] = False
        if not same_pos.any():
            continue
        hardest_pos = dist[i][same_pos].max()

        neg_mask = ~same[i]
        neg_mask[i] = False
        if not neg_mask.any():
            continue
        hardest_neg = dist[i][neg_mask].min()

        losses.append(F.relu(hardest_pos - hardest_neg + margin))

    if not losses:
        return embeddings.sum() * 0.0
    return torch.stack(losses).mean()


class ReIDEmbedModel(nn.Module):
    """ResNet50 预训练 backbone（去掉 fc）→ BN 512 嵌入 → logits（num_classes）。"""

    def __init__(self, num_classes: int, feat_dim: int = 512) -> None:
        super().__init__()
        backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
        self.backbone = nn.Sequential(*list(backbone.children())[:-2])  # 到 GAP 前，输出 [B,2048,7,7]
        self.gap = nn.AdaptiveAvgPool2d((1, 1))
        self.bottleneck = nn.Sequential(
            nn.Linear(2048, feat_dim),
            nn.BatchNorm1d(feat_dim),
            nn.ReLU(inplace=True),
        )
        self.classifier = nn.Linear(feat_dim, num_classes)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        h = self.backbone(x)
        h = self.gap(h)
        h = h.view(h.size(0), -1)
        emb = self.bottleneck(h)
        logits = self.classifier(emb)
        return emb, logits


def collate_fn_pk(batch):
    xs, ys, cams = zip(*batch, strict=True)
    return torch.stack(xs, dim=0), torch.stack(ys, dim=0), torch.stack(cams, dim=0)


def train_one_epoch(
    model: nn.Module,
    loader: DataLoader,
    optim: torch.optim.Optimizer,
    device: torch.device,
    margin: float,
    triplet_w: float,
    id_w: float,
) -> tuple[float, float, float]:
    model.train()
    sum_t = 0.0
    sum_id = 0.0
    n = 0
    ce = nn.CrossEntropyLoss()

    for x, y, _cam in loader:
        x = x.to(device, non_blocking=True)
        y = y.to(device, non_blocking=True)
        emb, logits = model(x)
        emb = F.normalize(emb, p=2, dim=1)
        loss_t = batch_hard_triplet_loss(emb, y, margin=margin)
        loss_id = ce(logits, y)
        loss = triplet_w * loss_t + id_w * loss_id

        optim.zero_grad()
        loss.backward()
        optim.step()

        bs = x.size(0)
        sum_t += loss_t.detach().item() * bs
        sum_id += loss_id.detach().item() * bs
        n += bs

    return sum_t / max(n, 1), sum_id / max(n, 1)


@torch.no_grad()
def estimate_id_accuracy(model: nn.Module, loader: DataLoader, device: torch.device) -> float:
    model.eval()
    correct = 0
    total = 0
    for x, y, _ in loader:
        x = x.to(device)
        y = y.to(device)
        _, logits = model(x)
        pred = logits.argmax(dim=1)
        correct += int((pred == y).sum().item())
        total += y.numel()
    return correct / max(total, 1)


def resolve_image_root(cli: str | Path) -> Path:
    p = Path(cli).resolve()
    sub = p / "doctor_picture"
    if sub.is_dir():
        return sub
    if p.is_dir():
        jpgs = list(p.glob("*.jpg"))
        if len(jpgs) > 0:
            return p
    raise FileNotFoundError(f"未找到图像目录（可传 doctor_info_detect 或 doctor_picture）: {cli}")


def parse_args():
    ap = argparse.ArgumentParser(description="Doctor Re-ID PK Triplet + ID 训练")
    ap.add_argument(
        "--data-root",
        type=Path,
        default=Path(__file__).resolve().parent,
        help="含 doctor_picture 或图片根目录的路径",
    )
    ap.add_argument("--epochs", type=int, default=50, help="epoch 建议在 40–60（默认 50）")
    ap.add_argument(
        "--batch-p",
        type=int,
        default=5,
        help="PK 采样：每 batch 采样的身份数 P（将自动不大于身份总数，如默认 5 类）",
    )
    ap.add_argument(
        "--batch-k",
        type=int,
        default=8,
        help="PK 采样：每位身份抽样张数 K；batch_size=P×K",
    )
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--triplet-margin", type=float, default=0.3)
    ap.add_argument("--triplet-weight", type=float, default=1.0)
    ap.add_argument("--id-weight", type=float, default=2.0, help="ID Loss（CrossEntropy）权重")
    ap.add_argument("--workers", type=int, default=4)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument(
        "--save",
        type=Path,
        default=Path(__file__).resolve().parent / "doctor_reid_best.pth",
        help="最佳权重路径",
    )
    return ap.parse_args()


def main() -> None:
    args = parse_args()
    rng = random.Random(args.seed)
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    img_root = resolve_image_root(args.data_root)

    ds_train_aug = DoctorReIDDataset(img_root, augment=True)
    ds_eval = DoctorReIDDataset(img_root, augment=False)

    n = len(ds_train_aug)
    perm = list(range(n))
    rng.shuffle(perm)
    n_val = max(32, int(0.1 * n))
    val_ix = sorted(perm[:n_val])
    train_ix = perm[n_val:]

    labels_tr = [ds_train_aug.labels[i] for i in train_ix]

    num_classes = len(ds_train_aug.pid_to_label)
    p_eff = min(args.batch_p, len(set(labels_tr)))

    from torch.utils.data import Subset

    sampler = RandomIdentityPKSampler(
        labels_tr,
        p=p_eff,
        k=args.batch_k,
        seed=args.seed,
    )
    train_loader = DataLoader(
        Subset(ds_train_aug, train_ix),
        batch_size=sampler.batch_size,
        sampler=sampler,
        num_workers=args.workers,
        pin_memory=True,
        drop_last=True,
        collate_fn=collate_fn_pk,
    )
    val_loader = DataLoader(
        Subset(ds_eval, val_ix),
        batch_size=64,
        shuffle=False,
        num_workers=args.workers,
        collate_fn=collate_fn_pk,
    )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = ReIDEmbedModel(num_classes=num_classes, feat_dim=512).to(device)

    optim = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=5e-4)
    scheduler = CosineAnnealingLR(optim, T_max=args.epochs, eta_min=1e-6)

    best_acc = -1.0
    for epoch in range(1, args.epochs + 1):
        sampler.set_epoch(epoch)
        tr_t, tr_id = train_one_epoch(
            model,
            train_loader,
            optim,
            device,
            margin=args.triplet_margin,
            triplet_w=args.triplet_weight,
            id_w=args.id_weight,
        )
        scheduler.step()
        lr_now = optim.param_groups[0]["lr"]
        val_acc = estimate_id_accuracy(model, val_loader, device)

        print(
            f"epoch {epoch:03d}/{args.epochs} | "
            f"triplet {tr_t:.4f} | id_loss_ce {tr_id:.4f} | "
            f"lr {lr_now:.6f} | val_id_acc ~ {val_acc:.4f}"
        )

        if val_acc >= best_acc:
            best_acc = val_acc
            torch.save(
                {
                    "epoch": epoch,
                    "model_state": model.state_dict(),
                    "num_classes": num_classes,
                    "pid_to_label": ds_train_aug.pid_to_label,
                    "args": vars(args),
                },
                args.save,
            )
            print(f"[保存] checkpoint → {args.save} (best val_id_acc {best_acc:.4f})")

    print(f"训练结束。最佳 val_id_acc≈{best_acc:.4f}, 权重: {args.save}")


if __name__ == "__main__":
    main()