"""Hunt label errors in the dataset with the trained model + FiftyOne (plan: data audit).

Where the model confidently disagrees with the ground truth, the *label* is
often what's wrong — especially for the NO-* violation classes, where missed
annotations are easiest to make. This script loads a split, attaches model
predictions, computes FiftyOne's mistakenness score, and opens the FiftyOne app
sorted so the most suspicious annotations come first.

Needs a local dataset copy (scripts/download_data.py) and the `audit` group:

    uv run --group audit python scripts/audit_labels.py                  # train split, opens the app
    uv run --group audit python scripts/audit_labels.py --split valid
    uv run --group audit python scripts/audit_labels.py --no-app --top 50

Review workflow: walk the sorted samples in the app, fix bad boxes/classes in
your labeling tool (Roboflow/CVAT) — FiftyOne is the *finder*, not the editor.
Expect a meaningful hit-rate in the top few hundred; stop when flags stop
being real errors.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import fiftyone as fo
import fiftyone.brain as fob
from ultralytics import YOLO

WEIGHTS = "runs/detect/yolov8n_v1_train/weights/best.pt"
DATA_DIR = "work"


def attach_predictions(dataset: fo.Dataset, weights: str, conf: float, imgsz: int) -> None:
    model = YOLO(weights)
    names = model.names
    with fo.ProgressBar() as pb:
        for sample in pb(dataset):
            r = model.predict(sample.filepath, conf=conf, imgsz=imgsz, verbose=False)[0]
            dets = []
            for box in r.boxes:
                x1, y1, x2, y2 = box.xyxyn[0].tolist()
                dets.append(
                    fo.Detection(
                        label=names[int(box.cls)],
                        bounding_box=[x1, y1, x2 - x1, y2 - y1],
                        confidence=float(box.conf),
                    )
                )
            sample["predictions"] = fo.Detections(detections=dets)
            sample.save()


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--data-dir", default=DATA_DIR)
    p.add_argument("--weights", default=WEIGHTS)
    p.add_argument("--split", choices=["train", "valid", "test"], default="train")
    p.add_argument("--conf", type=float, default=0.15, help="low on purpose — misses need candidates to compare")
    p.add_argument("--imgsz", type=int, default=640)
    p.add_argument("--top", type=int, default=25, help="how many most-suspicious samples to print")
    p.add_argument("--no-app", action="store_true", help="print the ranking only, don't launch the FiftyOne app")
    args = p.parse_args()

    if not (Path(args.data_dir) / args.split / "images").is_dir():
        raise SystemExit(f"{args.data_dir}/{args.split}/images not found — fetch with scripts/download_data.py")

    name = f"ppe-audit-{args.split}"
    if fo.dataset_exists(name):
        fo.delete_dataset(name)  # fresh predictions each run beat stale cached ones
    dataset = fo.Dataset.from_dir(
        name=name,
        dataset_type=fo.types.YOLOv5Dataset,
        dataset_dir=args.data_dir,
        yaml_path="data.yaml",
        split=args.split,
        label_field="ground_truth",
    )
    print(f"Loaded {len(dataset)} {args.split} samples")

    print("Attaching model predictions ...")
    attach_predictions(dataset, args.weights, args.conf, args.imgsz)

    print("Computing mistakenness ...")
    fob.compute_mistakenness(dataset, "predictions", label_field="ground_truth")

    view = dataset.sort_by("mistakenness", reverse=True)
    print(f"\nTop {args.top} most-suspicious samples (highest mistakenness first):")
    for sample in view[: args.top]:
        print(f"  {sample.mistakenness:.3f}  {Path(sample.filepath).name}")

    if not args.no_app:
        session = fo.launch_app(view)
        session.wait()  # keep the app alive until the user closes it


if __name__ == "__main__":
    main()