import csv
import json
from collections import Counter
from pathlib import Path

from ultralytics import YOLO

root = Path(r"D:\work\yolo\data\Smoke-Fire-Detection-YOLO")
weights = Path(r"D:\work\yolo\runs\detect\smoke_fire_yolo11s_v1-4\weights\best.pt")
audit_dir = Path(r"D:\work\yolo\runs\audit")
csv_output = audit_dir / "suspected_missing_labels.csv"
json_output = audit_dir / "suspected_missing_labels.json"
image_extensions = (".jpg", ".jpeg", ".png", ".bmp", ".webp")

images: list[str] = []
image_splits: dict[str, str] = {}
for split in ("train", "val", "test"):
    for label_path in sorted((root / split / "labels").glob("*.txt")):
        if label_path.read_text(encoding="utf-8", errors="replace").strip():
            continue
        for extension in image_extensions:
            image_path = root / split / "images" / f"{label_path.stem}{extension}"
            if image_path.is_file():
                normalized = str(image_path.resolve())
                images.append(normalized)
                image_splits[normalized.lower()] = split
                break

model = YOLO(str(weights))
rows: list[dict[str, object]] = []
chunk_size = 32
for offset in range(0, len(images), chunk_size):
    chunk = images[offset : offset + chunk_size]
    results = model.predict(source=chunk, imgsz=640, conf=0.25, device=0, batch=16, verbose=False)
    for source_path, result in zip(chunk, results, strict=True):
        boxes = result.boxes
        if boxes is None or len(boxes) == 0:
            continue
        image_path = source_path
        for class_id, confidence, xyxy in zip(boxes.cls.tolist(), boxes.conf.tolist(), boxes.xyxy.tolist()):
            class_index = int(class_id)
            rows.append({
                "split": image_splits.get(image_path.lower(), "unknown"),
                "image": image_path,
                "class_id": class_index,
                "class_name": model.names[class_index],
                "confidence": round(float(confidence), 6),
                "x1": round(float(xyxy[0]), 2),
                "y1": round(float(xyxy[1]), 2),
                "x2": round(float(xyxy[2]), 2),
                "y2": round(float(xyxy[3]), 2),
            })
    if offset and offset % 1024 == 0:
        print(f"scanned {offset}/{len(images)}", flush=True)

rows.sort(key=lambda row: float(row["confidence"]), reverse=True)
fieldnames = ["split", "image", "class_id", "class_name", "confidence", "x1", "y1", "x2", "y2"]
with csv_output.open("w", encoding="utf-8", newline="") as csv_file:
    writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
    writer.writeheader()
    writer.writerows(rows)

thresholds = {}
for threshold in (0.25, 0.5, 0.7, 0.85):
    selected = [row for row in rows if float(row["confidence"]) >= threshold]
    thresholds[str(threshold)] = {
        "detections": len(selected),
        "images": len({row["image"] for row in selected}),
        "smoke_images": len({row["image"] for row in selected if row["class_name"] == "smoke"}),
        "fire_images": len({row["image"] for row in selected if row["class_name"] == "fire"}),
        "by_split": dict(Counter(str(row["split"]) for row in selected)),
    }
report = {
    "empty_label_images": len(images),
    "thresholds": thresholds,
    "top_candidates": rows[:50],
    "csv": str(csv_output),
}
json_output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
print(json.dumps(report, ensure_ascii=False, indent=2))