"""Fit M1, M2, M3 on the full plot-level dataset and save coefficients.

These in-sample fits are used by the demo app for single-image yield
prediction. The demo only has features from one uploaded image, so there
is no test fold; we want the best-known M3 mapping from (class areas,
mean_area) to predicted obs_sum_B_A1_A2_OV.

Usage:
    uv run python scripts/fit_final_models.py \\
        --input outputs/df_model_yield_by_plot.csv \\
        --out app/m3_coefficients.json
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score

from potatopheno.models import (
    PREDICTORS_M1,
    PREDICTORS_M2,
    PREDICTORS_M3,
    RESPONSE,
    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.")
    p.add_argument(
        "--out",
        type=Path,
        required=True,
        help="Path to write the JSON coefficients file (used by the Gradio app).",
    )
    return p.parse_args()


def _fit(df: pd.DataFrame, predictors: tuple[str, ...], name: str) -> dict:
    subset = df.dropna(subset=list(predictors) + [RESPONSE])
    X = subset[list(predictors)]
    y = subset[RESPONSE]
    model = LinearRegression()
    model.fit(X, y)
    pred = model.predict(X)
    return {
        "name": name,
        "predictors": list(predictors),
        "response": RESPONSE,
        "intercept": float(model.intercept_),
        "coefficients": {p: float(c) for p, c in zip(predictors, model.coef_)},
        "n_train_plots": int(len(subset)),
        "in_sample_r2": float(r2_score(y, pred)),
    }


def main() -> None:
    args = _parse_args()
    df = pd.read_csv(args.input)
    df = add_model_features(df)

    fits = {
        "M1": _fit(df, PREDICTORS_M1, "M1"),
        "M2": _fit(df, PREDICTORS_M2, "M2"),
        "M3": _fit(df, PREDICTORS_M3, "M3"),
    }

    args.out.parent.mkdir(parents=True, exist_ok=True)
    with open(args.out, "w") as f:
        json.dump(fits, f, indent=2)

    for name, fit in fits.items():
        print(
            f"{name}: n_train={fit['n_train_plots']:>4}  in-sample R²={fit['in_sample_r2']:.3f}  "
            f"intercept={fit['intercept']:+.3f}"
        )
    print(f"\n[DONE] wrote → {args.out}")


if __name__ == "__main__":
    main()