"""Run YOLOv11-seg inference on tuber images and emit per-image .txt files.

Default predict parameters: conf=0.25, imgsz=1280. Output .txt files
follow the YOLO-seg format consumed by the morphometrics stage.

Single-trial usage:
    uv run python scripts/run_yolo_inference.py \\
        --weights path/to/bestModel.pt \\
        --images path/to/2025NCPT \\
        --out-dir outputs/preds/NCPT

All three trials usage:
    uv run python scripts/run_yolo_inference.py \\
        --weights path/to/bestModel.pt \\
        --images-root path/to/rawData/images \\
        --out-root outputs/preds
"""

from __future__ import annotations

import argparse
from pathlib import Path

from potatopheno.segmentation import load_model, predict_directory

TRIALS = ("NCPT", "ADC", "NCR")


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--weights", type=Path, required=True, help="Path to bestModel.pt.")

    single = p.add_argument_group("single-trial mode")
    single.add_argument(
        "--images",
        type=Path,
        help="Directory of input images (single trial).",
    )
    single.add_argument(
        "--out-dir",
        type=Path,
        help="Directory to write .txt files into.",
    )

    multi = p.add_argument_group("all-trials mode")
    multi.add_argument(
        "--images-root",
        type=Path,
        help="Parent directory with 2025<TRIAL>/ subdirs.",
    )
    multi.add_argument(
        "--out-root",
        type=Path,
        help="Parent directory; <TRIAL>/preds/ will be created under it.",
    )
    multi.add_argument(
        "--trials",
        type=lambda s: tuple(t.strip() for t in s.split(",") if t.strip()),
        default=TRIALS,
        help=f"Comma-separated trial names (default: {','.join(TRIALS)}).",
    )

    p.add_argument("--conf", type=float, default=0.25, help="Confidence threshold.")
    p.add_argument("--imgsz", type=int, default=1280, help="Inference image size.")
    p.add_argument("--chunk-size", type=int, default=15, help="Images per streaming chunk.")
    return p.parse_args()


def main() -> None:
    args = _parse_args()

    single_mode = args.images is not None and args.out_dir is not None
    multi_mode = args.images_root is not None and args.out_root is not None
    if single_mode == multi_mode:
        raise SystemExit(
            "Pass EITHER (--images and --out-dir) OR (--images-root and --out-root), not both."
        )

    model = load_model(args.weights)
    predict_kwargs = {"conf": args.conf, "imgsz": args.imgsz, "verbose": False}

    if single_mode:
        predict_directory(
            model=model,
            input_dir=args.images,
            output_dir=args.out_dir,
            chunk_size=args.chunk_size,
            **predict_kwargs,
        )
        return

    for trial in args.trials:
        predict_directory(
            model=model,
            input_dir=args.images_root / f"2025{trial}",
            output_dir=args.out_root / trial / "preds",
            chunk_size=args.chunk_size,
            **predict_kwargs,
        )


if __name__ == "__main__":
    main()