"""Leave-one-trial-out evaluation for the yield-prediction models.

Each runner returns:
- df_pred: one row per plot per fold with predicted, error, abs_error,
  fold_test_trial, and model columns appended.
- df_metrics: one row per (model, test_trial) with R2, RMSE, MAE, MAPE,
  BIAS. BIAS uses the convention mean(obs - pred), so positive bias
  means under-prediction.
"""

from __future__ import annotations

from collections.abc import Iterable

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score


def _fold_metrics(
    model_name: str,
    test_trial: str,
    y_test: np.ndarray | pd.Series,
    pred: np.ndarray,
    n_test: int,
) -> dict:
    return {
        "model": model_name,
        "test_trial": test_trial,
        "n_test_plots": n_test,
        "R2": r2_score(y_test, pred),
        "RMSE": float(np.sqrt(mean_squared_error(y_test, pred))),
        "MAE": mean_absolute_error(y_test, pred),
        "MAPE": float(np.mean(np.abs((y_test - pred) / y_test)) * 100),
        "BIAS": float(np.mean(y_test - pred)),
    }


def _annotate_predictions(
    test: pd.DataFrame,
    pred: np.ndarray,
    response: str,
    test_trial: str,
    model_name: str,
) -> pd.DataFrame:
    out = test.copy()
    out["predicted"] = pred
    out["error"] = out[response] - out["predicted"]
    out["abs_error"] = out["error"].abs()
    out["fold_test_trial"] = test_trial
    out["model"] = model_name
    return out


def run_loto_linear(
    df: pd.DataFrame,
    predictors: Iterable[str],
    response: str,
    model_name: str,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """LOTO over trials with a fresh `LinearRegression` per fold."""
    predictors = list(predictors)
    df = df.copy()
    trials = sorted(df["trial"].dropna().unique())

    pred_rows: list[pd.DataFrame] = []
    metric_rows: list[dict] = []
    for test_trial in trials:
        train = df[df["trial"] != test_trial].copy()
        test = df[df["trial"] == test_trial].copy()
        train = train.dropna(subset=predictors + [response])
        test = test.dropna(subset=predictors + [response])

        X_train, y_train = train[predictors], train[response]
        X_test, y_test = test[predictors], test[response]

        model = LinearRegression()
        model.fit(X_train, y_train)
        pred = model.predict(X_test)

        pred_rows.append(_annotate_predictions(test, pred, response, test_trial, model_name))
        metric_rows.append(_fold_metrics(model_name, test_trial, y_test, pred, len(test)))

    df_pred = pd.concat(pred_rows, axis=0, ignore_index=True)
    df_metrics = pd.DataFrame(metric_rows)
    return df_pred, df_metrics


def run_loto_rf(
    df: pd.DataFrame,
    predictors: Iterable[str],
    response: str,
    model_name: str,
    **rf_kwargs,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """LOTO over trials with a fresh `RandomForestRegressor` per fold."""
    predictors = list(predictors)
    df = df.copy()
    trials = sorted(df["trial"].dropna().unique())

    pred_rows: list[pd.DataFrame] = []
    metric_rows: list[dict] = []
    for test_trial in trials:
        train = df[df["trial"] != test_trial].copy()
        test = df[df["trial"] == test_trial].copy()
        train = train.dropna(subset=predictors + [response])
        test = test.dropna(subset=predictors + [response])

        X_train, y_train = train[predictors], train[response]
        X_test, y_test = test[predictors], test[response]

        model = RandomForestRegressor(**rf_kwargs)
        model.fit(X_train, y_train)
        pred = model.predict(X_test)

        pred_rows.append(_annotate_predictions(test, pred, response, test_trial, model_name))
        metric_rows.append(_fold_metrics(model_name, test_trial, y_test, pred, len(test)))

    df_pred = pd.concat(pred_rows, axis=0, ignore_index=True)
    df_metrics = pd.DataFrame(metric_rows)
    return df_pred, df_metrics


def summarize_across_trials(df_metrics: pd.DataFrame) -> pd.DataFrame:
    """Mean of each metric across LOTO folds, one row per model."""
    return (
        df_metrics.groupby("model", as_index=False)
        .agg(
            n_test_plots=("n_test_plots", "sum"),
            mean_R2=("R2", "mean"),
            mean_RMSE=("RMSE", "mean"),
            mean_MAE=("MAE", "mean"),
            mean_MAPE=("MAPE", "mean"),
            mean_BIAS=("BIAS", "mean"),
        )
    )