"""Camera 01 录像：直接 YOLO 与共享实时缩图管线的一致性抽样验证。

手动运行：python tests/validate_yolo_replay.py <video> <weights> [sample_count] [device]
"""
import sys

import cv2
from ultralytics import YOLO


def iou(a, b):
    x1, y1, x2, y2 = max(a[0], b[0]), max(a[1], b[1]), min(a[2], b[2]), min(a[3], b[3])
    inter = max(0, x2 - x1) * max(0, y2 - y1)
    union = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
    return inter / union if union > 0 else 0


def boxes(result):
    if result.boxes is None:
        return []
    return [list(map(float, row)) for row in result.boxes.xyxy.cpu().numpy()]


def main(video_path, weights_path, sample_count=5, device="cpu"):
    model = YOLO(weights_path)
    names = model.names if isinstance(model.names, dict) else dict(enumerate(model.names))
    person_ids = [cid for cid, name in names.items() if name == "person"]
    cap = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    indices = [int((i + 1) * total / (sample_count + 1)) for i in range(sample_count)]
    direct_count = matched_count = 0
    matched_ious = []
    for index in indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index)
        ok, frame = cap.read()
        if not ok:
            continue
        h, w = frame.shape[:2]
        precision = 16 if str(device).lower() != "cpu" else None
        direct = boxes(model.predict(frame, imgsz=640, conf=.3, iou=.5,
                                     classes=person_ids, device=device,
                                     quantize=precision, verbose=False)[0])
        scale = min(1.0, 2560.0 / max(w, h))
        small = (frame if scale == 1.0 else
                 cv2.resize(frame, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA))
        ok, encoded = cv2.imencode(".jpg", small, [cv2.IMWRITE_JPEG_QUALITY, 95])
        realtime = boxes(model.predict(cv2.imdecode(encoded, cv2.IMREAD_COLOR), imgsz=640,
                                       conf=.3, iou=.5, classes=person_ids,
                                       device=device, quantize=precision,
                                       verbose=False)[0]) if ok else []
        realtime = [[v / scale for v in box] for box in realtime]
        unused = set(range(len(realtime)))
        for box in direct:
            direct_count += 1
            candidates = [(iou(box, realtime[j]), j) for j in unused]
            best = max(candidates, default=(0, None))
            if best[0] >= .8:
                matched_count += 1
                matched_ious.append(best[0])
                unused.remove(best[1])
        print("frame=%d direct_person=%d realtime_person=%d" % (index, len(direct), len(realtime)))
    cap.release()
    recall = matched_count / direct_count if direct_count else 1.0
    mean_iou = sum(matched_ious) / len(matched_ious) if matched_ious else 0.0
    print("matched=%d/%d recall=%.3f mean_iou=%.3f" %
          (matched_count, direct_count, recall, mean_iou))
    return 0 if recall >= .95 and (not matched_ious or min(matched_ious) >= .8) else 1


if __name__ == "__main__":
    if len(sys.argv) < 3:
        raise SystemExit("usage: validate_yolo_replay.py VIDEO WEIGHTS [SAMPLES] [DEVICE]")
    raise SystemExit(main(sys.argv[1], sys.argv[2], int(sys.argv[3]) if len(sys.argv) > 3 else 5,
                          sys.argv[4] if len(sys.argv) > 4 else "cpu"))