"""Run the eye open/closed model on one or more images."""
import sys
from pathlib import Path

from ultralytics import YOLO


def main():
    project_root = Path(__file__).resolve().parent.parent
    weight_candidates = [
        project_root / "weights" / "best.pt",
        project_root / "runs" / "classify" / "runs_cls" / "eye_open_closed_cpu" / "weights" / "best.pt",
        project_root / "runs" / "classify" / "runs_cls" / "eye_open_closed_cpu" / "weights" / "last.pt",
    ]
    weights = next((p for p in weight_candidates if p.is_file()), None)
    if weights is None:
        print("Weights not found. Put best.pt in weights/ or runs/.../weights/ (from model team).")
        sys.exit(1)

    if len(sys.argv) < 2:
        print("Usage: python scripts/predict_image.py <image1> [image2 ...]")
        print("Example: python scripts/predict_image.py path/to/image.png")
        sys.exit(0)

    model = YOLO(str(weights))
    names = model.names

    for path in sys.argv[1:]:
        p = Path(path)
        if not p.is_file():
            print(p, "- file not found")
            continue
        try:
            results = model.predict(str(p), imgsz=224, device="cpu", verbose=False)
        except Exception as e:
            print(p, "- error:", e)
            continue
        if not results:
            print(p, "- no result")
            continue
        r = results[0]
        top_idx = int(r.probs.top1)
        conf = float(r.probs.top1conf)
        label = names[top_idx]
        print(f"{p.name}: {label} ({conf:.2%})")


if __name__ == "__main__":
    main()