#!/usr/bin/env python3
"""Audit an existing still collection with the anime-face detector:
how many saved images actually contain a qualifying face?

Useful for measuring how much junk the old person-proxy pipeline left
behind, and as a QC pass over any facesaver output.

    python audit_stills.py -I <stills_dir> --gpu 1 [--write-report report.csv]
"""

import argparse
import csv
import os
import sys


def parse_args():
    p = argparse.ArgumentParser(
        description=(__doc__ or "Audit stills for anime faces.").split("\n")[0])
    p.add_argument("-I", "--input-dir", required=True,
                   help="Directory of still images to audit.")
    p.add_argument("--model",
                   default=os.path.join(
                       os.path.dirname(os.path.abspath(__file__)),
                       "models", "yolov8x6_animeface.pt"))
    p.add_argument("--gpu", default="1",
                   help="GPU index, 'cpu', or 'auto'.")
    p.add_argument("--conf", type=float, default=0.3)
    p.add_argument("--imgsz", type=int, default=1280)
    p.add_argument("--batch", type=int, default=8)
    p.add_argument("--write-report", default=None,
                   help="Optional CSV: filename, n_faces, max_w, max_h, "
                        "ok (1/0).")
    return p.parse_args()


def main():
    args = parse_args()
    if args.gpu.lower() == "cpu":
        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
        device = "cpu"
    elif args.gpu.lower() != "auto":
        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
        device = "cuda"
    else:
        device = None

    import cv2
    import numpy as np
    from ultralytics import YOLO

    exts = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
    files = sorted(f for f in os.listdir(args.input_dir)
                   if f.lower().endswith(exts))
    if not files:
        print("No images found.")
        sys.exit(1)
    print(f"Auditing {len(files)} images with "
          f"{os.path.basename(args.model)} on {device or 'auto'}")

    model = YOLO(args.model)

    rows = []
    no_face = 0
    # stream in chunks so 1000s of 1080p pngs don't eat RAM
    for i in range(0, len(files), args.batch):
        chunk = files[i:i + args.batch]
        imgs, kept = [], []
        for f in chunk:
            img = cv2.imread(os.path.join(args.input_dir, f))
            if img is None:
                print(f"  WARN unreadable: {f}")
                continue
            imgs.append(img)
            kept.append(f)
        results = model.predict(imgs, conf=args.conf, imgsz=args.imgsz,
                                device=device, verbose=False)
        for f, res in zip(kept, results):
            n, mw, mh = 0, 0, 0
            if res.boxes is not None and len(res.boxes):
                xyxy = res.boxes.xyxy.cpu().numpy()
                n = len(xyxy)
                for (x1, y1, x2, y2) in xyxy:
                    mw = max(mw, int(x2 - x1))
                    mh = max(mh, int(y2 - y1))
            ok = 1 if n > 0 else 0
            if not ok:
                no_face += 1
                print(f"  NO FACE: {f}")
            rows.append((f, n, mw, mh, ok))
        done = min(i + args.batch, len(files))
        print(f"  {done}/{len(files)}", flush=True)

    print()
    print(f"Total: {len(rows)}   with face: {len(rows) - no_face}   "
          f"NO face: {no_face} ({100.0 * no_face / max(1, len(rows)):.1f}%)")

    if args.write_report:
        with open(args.write_report, "w", newline="") as fh:
            w = csv.writer(fh)
            w.writerow(["filename", "n_faces", "max_w", "max_h", "ok"])
            w.writerows(rows)
        print(f"Report: {args.write_report}")


if __name__ == "__main__":
    main()