import csv
from pathlib import Path

import cv2
from ultralytics import YOLO

weights = Path(r"D:\work\yolo\runs\detect\smoke_fire_yolo11s_v1-4\weights\best.pt")
candidates_path = Path(r"D:\work\yolo\runs\audit\suspected_missing_labels.csv")
output_dir = Path(r"D:\work\yolo\runs\audit\missing_label_review_ge_0.50")
threshold = 0.50
class_names = ["smoke", "fire"]

with candidates_path.open(encoding="utf-8", newline="") as csv_file:
    candidates = list(csv.DictReader(csv_file))
image_paths = []
for row in candidates:
    if float(row["confidence"]) < threshold:
        continue
    image_path = Path(row["image"])
    if image_path not in image_paths:
        image_paths.append(image_path)

model = YOLO(str(weights))
output_dir.mkdir(parents=True, exist_ok=True)
manifest = []
for index, image_path in enumerate(image_paths, start=1):
    image = cv2.imread(str(image_path))
    if image is None:
        continue
    height, width = image.shape[:2]
    label_path = image_path.parents[1] / "labels" / f"{image_path.stem}.txt"
    labeled = image.copy()
    label_lines = [line for line in label_path.read_text(encoding="utf-8").splitlines() if line.strip()] if label_path.is_file() else []
    for line in label_lines:
        values = line.split()
        class_id = int(float(values[0]))
        x_center, y_center, box_width, box_height = map(float, values[1:])
        x1 = int((x_center - box_width / 2) * width)
        y1 = int((y_center - box_height / 2) * height)
        x2 = int((x_center + box_width / 2) * width)
        y2 = int((y_center + box_height / 2) * height)
        cv2.rectangle(labeled, (x1, y1), (x2, y2), (0, 220, 0), 3)
        cv2.putText(labeled, class_names[class_id], (x1, max(28, y1 - 7)), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 220, 0), 2)

    predicted = image.copy()
    result = model.predict(str(image_path), imgsz=640, conf=threshold, device=0, verbose=False)[0]
    predictions = []
    for box in result.boxes:
        class_id = int(box.cls.item())
        confidence = float(box.conf.item())
        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
        predictions.append({"class": class_names[class_id], "confidence": round(confidence, 6), "xyxy": [x1, y1, x2, y2]})
        cv2.rectangle(predicted, (x1, y1), (x2, y2), (0, 0, 255), 3)
        cv2.putText(predicted, f"{class_names[class_id]} {confidence:.3f}", (x1, max(28, y1 - 7)), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)

    canvas = cv2.hconcat([labeled, predicted])
    cv2.rectangle(canvas, (0, 0), (canvas.shape[1], 58), (255, 255, 255), -1)
    cv2.putText(canvas, f"Original label: {image_path.name}", (20, 38), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 120, 0), 2)
    cv2.putText(canvas, "YOLO11s prediction", (width + 20, 38), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 180), 2)
    output_path = output_dir / f"{index:03d}_{image_path.stem}.jpg"
    if not cv2.imwrite(str(output_path), canvas):
        raise OSError(output_path)
    manifest.append({"source": str(image_path), "label": str(label_path), "review_image": str(output_path), "predictions": predictions})
    print(f"rendered {index}/{len(image_paths)} {image_path.name}", flush=True)

manifest_path = output_dir / "manifest.json"
import json
manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
print(f"completed={len(manifest)} manifest={manifest_path}")