"""Run the morphometrics stage end-to-end across all trials.

Reads YOLO-seg .txt predictions and the matching source images, computes
per-tuber polygon morphometrics, converts to real-scale (inches and
square inches), assigns USDA size class, and concatenates the three
trials into a single combined CSV.

Usage:
    uv run python scripts/extract_morphometrics.py \\
        --out outputs/df_class_real_scale_ALL.csv
"""

from __future__ import annotations

import argparse
from pathlib import Path

import pandas as pd

from potatopheno.calibration import BELT_WIDTH_IN_DEFAULT, apply_real_scale
from potatopheno.classification import assign_usda_class
from potatopheno.io import build_morphometrics_dataframe, plot_id_from_image_id

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


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument(
        "--preds-root",
        type=Path,
        required=True,
        help="Root of YOLO preds. Expects <trial>/preds/*.txt under it.",
    )
    p.add_argument(
        "--images-root",
        type=Path,
        required=True,
        help="Root of raw images. Expects 2025<TRIAL>/*.JPG under it.",
    )
    p.add_argument(
        "--out",
        type=Path,
        required=True,
        help="Path of the combined per-tuber CSV to write.",
    )
    p.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(
        "--belt-width-in",
        type=float,
        default=BELT_WIDTH_IN_DEFAULT,
        help="Conveyor belt width in inches used for pixel-to-inch calibration.",
    )
    return p.parse_args()


def main() -> None:
    args = _parse_args()
    args.out.parent.mkdir(parents=True, exist_ok=True)

    per_trial: list[pd.DataFrame] = []
    for trial in args.trials:
        preds_dir = args.preds_root / trial / "preds"
        images_dir = args.images_root / f"2025{trial}"

        df = build_morphometrics_dataframe(trial, images_dir, preds_dir)
        df["Plot"] = plot_id_from_image_id(df["image_id"])
        df = apply_real_scale(df, belt_width_in=args.belt_width_in)
        df = assign_usda_class(df)

        print(
            f"[{trial}] images={df['image_id'].nunique()} "
            f"instances={len(df)} "
            f"class_counts={df['class'].value_counts().to_dict()}"
        )
        per_trial.append(df)

    combined = pd.concat(per_trial, ignore_index=True)
    combined.to_csv(args.out, index=False)

    print(
        f"\n[DONE] wrote {len(combined)} rows across "
        f"{combined['trial'].nunique()} trials → {args.out}"
    )


if __name__ == "__main__":
    main()