"""Predictor sets and feature engineering for the yield-prediction models.

M1, M2, M3 are linear regressions on increasingly rich plot-level
features. RF is a Random Forest baseline that uses a seven-feature
predictor set including squared and total-area aggregate features.
"""

from __future__ import annotations

import pandas as pd

PREDICTORS_M1: tuple[str, ...] = ("total_area_in2",)
PREDICTORS_M2: tuple[str, ...] = (
    "area_B_in2",
    "area_A1_in2",
    "area_A2_in2",
    "area_OV_in2",
)
PREDICTORS_M3: tuple[str, ...] = (
    "area_B_in2",
    "area_A1_in2",
    "area_A2_in2",
    "area_OV_in2",
    "mean_area",
)
# Wider feature set for the Random Forest baseline: four class areas plus
# mean_area, total_area_sq, and total_area_in2.
PREDICTORS_RF: tuple[str, ...] = (
    "area_B_in2",
    "area_A1_in2",
    "area_A2_in2",
    "area_OV_in2",
    "mean_area",
    "total_area_sq",
    "total_area_in2",
)

RESPONSE: str = "obs_sum_B_A1_A2_OV"

RF_KWARGS: dict = {
    "n_estimators": 500,
    "random_state": 42,
    "max_depth": None,
    "min_samples_split": 2,
    "min_samples_leaf": 1,
    "n_jobs": -1,
}


def add_model_features(df: pd.DataFrame) -> pd.DataFrame:
    """Recompute total_area_in2 from class areas, add mean_area and total_area_sq.

    total_area_in2 is overwritten with the sum of the four class-area
    columns (this can differ from the upstream aggregate if any tubers
    have a NaN class). mean_area = total_area_in2 / n_tubers, and
    total_area_sq = total_area_in2 ** 2.
    """
    df = df.copy()
    df["total_area_in2"] = df[list(PREDICTORS_M2)].sum(axis=1)
    df["mean_area"] = df["total_area_in2"] / df["n_tubers"]
    df["total_area_sq"] = df["total_area_in2"] ** 2
    return df