"""Eval harness: compare predicted YOLO-seg labels against a gold set.

Reports detection precision/recall/F1 per class (bbox IoU >= threshold to match)
and mean mask IoU over matched instances. Operates on label files only (no source
images needed), so it can score any pre-labeler / model against the frozen gold set.

    python -m src.pipeline.evaluate --pred OUT/labels --gold GOLD/labels [--names dataset.yaml]
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Dict, List

from src.eval.metrics import aggregate, evaluate_image, full_report
from src.eval.yolo_labels import load_class_names, read_kept_outlines, read_label_file


def _stems(label_dir: Path) -> Dict[str, Path]:
    return {p.stem: p for p in label_dir.glob("*.txt")}


def evaluate_dirs(pred_dir, gold_dir, iou_thr: float = 0.5,
                  names: List[str] | None = None) -> Dict:
    pred_dir, gold_dir = Path(pred_dir), Path(gold_dir)
    gold_files = _stems(gold_dir)
    pred_files = _stems(pred_dir)
    # outlines a model proposed and the reviewer kept are not judged (see evaluate_image)
    kept = read_kept_outlines(gold_dir.parent)

    # Gold defines the evaluation set; a missing prediction counts as all-FN.
    per_image = []
    missing_pred = 0
    ignored = 0
    for stem, gpath in sorted(gold_files.items()):
        gold = read_label_file(gpath)
        ppath = pred_files.get(stem)
        pred = read_label_file(ppath) if ppath else []
        if ppath is None:
            missing_pred += 1
        ign = kept.get(stem, ())
        ignored += len(ign)
        per_image.append(evaluate_image(gold, pred, iou_thr, ignore=ign))

    report = full_report(aggregate(per_image), names)
    report["dataset"] = {
        "gold_images": len(gold_files),
        "pred_images": len(pred_files),
        "missing_predictions": missing_pred,
        "iou_threshold": iou_thr,
        "ignored_objects": ignored,
    }
    return report


def _print_report(report: Dict) -> None:
    ds = report["dataset"]
    print(f"\nGold images: {ds['gold_images']}  |  predictions: {ds['pred_images']}  "
          f"|  missing: {ds['missing_predictions']}  |  match IoU >= {ds['iou_threshold']}")
    header = f"{'class':<14}{'P':>7}{'R':>7}{'F1':>7}{'maskIoU':>9}{'TP':>6}{'FP':>6}{'FN':>6}"
    print("\n" + header)
    print("-" * len(header))

    def row(name, r):
        print(f"{name:<14}{r['precision']:>7.3f}{r['recall']:>7.3f}{r['f1']:>7.3f}"
              f"{r['mean_mask_iou']:>9.3f}{r['tp']:>6}{r['fp']:>6}{r['fn']:>6}")

    for name, r in report["per_class"].items():
        row(name, r)
    print("-" * len(header))
    row("OVERALL", report["overall"])


def main() -> int:
    ap = argparse.ArgumentParser(description="Evaluate predicted YOLO-seg labels vs. a gold set")
    ap.add_argument("--pred", required=True, help="Directory of predicted .txt labels")
    ap.add_argument("--gold", required=True, help="Directory of gold .txt labels")
    ap.add_argument("--iou", type=float, default=0.5, help="bbox IoU threshold to match instances")
    ap.add_argument("--names", default=None, help="Optional dataset.yaml for class names")
    ap.add_argument("--report", default=None, help="Optional path to write the JSON report")
    args = ap.parse_args()

    names = load_class_names(args.names) if args.names else []
    report = evaluate_dirs(args.pred, args.gold, iou_thr=args.iou, names=names)
    _print_report(report)
    if args.report:
        Path(args.report).write_text(json.dumps(report, indent=2), encoding="utf-8")
        print(f"\nWrote {args.report}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())