"""
Train the XGBoost yield prediction regressor.

Usage:
    python -m ml_pipeline.training.train_xgb_yield \
        --input data/processed/final_feature_matrix.csv \
        --output-dir data/models/v1
"""

import argparse
from pathlib import Path

import joblib
import numpy as np
import xgboost as xgb
from sklearn.metrics import mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split

import pandas as pd

from ml_pipeline.training.training_utils import check_minimum_samples, get_logger, save_training_metadata

logger = get_logger(__name__)

FEATURE_COLUMNS = [
    "rainfall_mm", "temp_c", "humidity",
    "season_index", "rainfall_anomaly",
    "rainfall_ma3", "rainfall_lag1",
    "water_req_mm", "min_temp_c", "max_temp_c",
    "drought_tolerant", "growth_days",
]
TARGET_COLUMN = "yield_ton_ha"


def main() -> None:
    parser = argparse.ArgumentParser(description="Train XGBoost yield regressor")
    parser.add_argument("--input", required=True)
    parser.add_argument("--output-dir", required=True)
    args = parser.parse_args()

    df = pd.read_csv(args.input)

    missing_cols = set(FEATURE_COLUMNS + [TARGET_COLUMN]) - set(df.columns)
    if missing_cols:
        raise ValueError(f"Input is missing required columns: {missing_cols}")

    df = df.dropna(subset=FEATURE_COLUMNS + [TARGET_COLUMN])
    # PILOT THRESHOLD: see same note in train_xgb_feasibility.py — lowered
    # from 200 to 100 for this initial small-data pilot run, deliberately
    # and visibly, not silently. Raise back before any client-facing use.
    check_minimum_samples(len(df), "xgb_yield", minimum=100)

    X = df[FEATURE_COLUMNS].values
    y = df[TARGET_COLUMN].values

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    model = xgb.XGBRegressor(
        n_estimators=300,
        max_depth=5,
        learning_rate=0.05,
        subsample=0.8,
        random_state=42,
    )
    model.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)

    y_pred = model.predict(X_test)
    mae = mean_absolute_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    logger.info(f"Test MAE: {mae:.3f} tonnes/ha, R2: {r2:.3f}")

    if r2 < 0.5:
        logger.warning(
            f"R2 of {r2:.2f} means the model explains less than half the "
            f"variance in yield. Treat predictions as directional, not precise, "
            f"until more training data or better features are available."
        )

    # Bootstrap ensemble for uncertainty: train N models on resampled data,
    # std of their predictions becomes a real uncertainty estimate rather
    # than the placeholder heuristic the live API falls back to if this
    # artifact is absent.
    logger.info("Training bootstrap ensemble for uncertainty estimation...")
    n_bootstrap = 10
    bootstrap_models = []
    rng = np.random.default_rng(42)

    for i in range(n_bootstrap):
        idx = rng.choice(len(X_train), size=len(X_train), replace=True)
        bm = xgb.XGBRegressor(
            n_estimators=300, max_depth=5, learning_rate=0.05,
            subsample=0.8, random_state=i,
        )
        bm.fit(X_train[idx], y_train[idx])
        bootstrap_models.append(bm)

    bootstrap_preds = np.array([bm.predict(X_test) for bm in bootstrap_models])
    avg_uncertainty = float(np.mean(np.std(bootstrap_preds, axis=0)))
    logger.info(f"Average bootstrap uncertainty (std across ensemble): {avg_uncertainty:.3f} tonnes/ha")

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    bundle = {
        "model": model,
        "feature_cols": FEATURE_COLUMNS,
        "bootstrap_models": bootstrap_models,  # used by app for real uncertainty
                                                 # if the inference code is upgraded
                                                 # beyond the v1 placeholder heuristic
    }
    output_path = output_dir / "xgb_yield.pkl"
    joblib.dump(bundle, output_path)

    save_training_metadata(
        output_dir,
        model_name="xgb_yield",
        metrics={
            "mae_ton_ha": float(mae),
            "r2": float(r2),
            "avg_bootstrap_uncertainty_ton_ha": avg_uncertainty,
        },
        feature_cols=FEATURE_COLUMNS,
        n_train=len(X_train),
        n_test=len(X_test),
    )

    logger.info(f"Saved model bundle to {output_path}")


if __name__ == "__main__":
    main()