"""Run M1, M2, M3, and RF under leave-one-trial-out yield prediction.

Adds the M3 structural features, runs LOTO for each of the four models,
and writes the per-fold metrics, per-plot predictions, and the cross-fold
summary.

Usage:
    uv run python scripts/run_loto_models.py \\
        --input outputs/df_model_yield_by_plot.csv \\
        --out-dir outputs/loto/
"""

from __future__ import annotations

import argparse
from pathlib import Path

import pandas as pd

from potatopheno.evaluate import (
    run_loto_linear,
    run_loto_rf,
    summarize_across_trials,
)
from potatopheno.models import (
    PREDICTORS_M1,
    PREDICTORS_M2,
    PREDICTORS_M3,
    PREDICTORS_RF,
    RESPONSE,
    RF_KWARGS,
    add_model_features,
)


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument(
        "--input",
        type=Path,
        required=True,
        help="Plot-level model dataset (output of build_plot_model_dataset.py).",
    )
    p.add_argument(
        "--out-dir",
        type=Path,
        required=True,
        help="Directory to write the metrics, predictions, and summary CSVs.",
    )
    return p.parse_args()


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

    df = pd.read_csv(args.input)
    df = add_model_features(df)

    runs = [
        ("M1", run_loto_linear, PREDICTORS_M1, {}),
        ("M2", run_loto_linear, PREDICTORS_M2, {}),
        ("M3", run_loto_linear, PREDICTORS_M3, {}),
        ("RF", run_loto_rf, PREDICTORS_RF, RF_KWARGS),
    ]

    pred_parts: list[pd.DataFrame] = []
    metric_parts: list[pd.DataFrame] = []
    for name, runner, predictors, extra in runs:
        df_pred, df_metrics = runner(df, predictors, RESPONSE, name, **extra)
        pred_parts.append(df_pred)
        metric_parts.append(df_metrics)

    df_pred_all = pd.concat(pred_parts, ignore_index=True)
    df_metrics_all = (
        pd.concat(metric_parts, ignore_index=True)
        .sort_values(["model", "test_trial"])
        .reset_index(drop=True)
    )
    df_summary = summarize_across_trials(df_metrics_all).sort_values("mean_MAPE")

    df_metrics_all.to_csv(args.out_dir / "loto_metrics_by_fold.csv", index=False)
    df_pred_all.to_csv(args.out_dir / "loto_predictions_by_plot.csv", index=False)
    df_summary.to_csv(args.out_dir / "loto_summary_by_model.csv", index=False)

    print("\nLOTO metrics by fold:\n")
    print(df_metrics_all.round(3).to_string(index=False))
    print("\nLOTO summary across folds (mean of trial-level metrics):\n")
    print(df_summary.round(3).to_string(index=False))
    print(f"\n[DONE] wrote 3 CSVs → {args.out_dir}")


if __name__ == "__main__":
    main()