"""
Train the Random Forest crop recommendation/ranking classifier.

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

import argparse
from pathlib import Path

import joblib
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report
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 = "suitability_label"


def main() -> None:
    parser = argparse.ArgumentParser(description="Train Random Forest recommendation model")
    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), "rf_recommendation", minimum=100)

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

    class_counts = pd.Series(y).value_counts()
    logger.info(f"Suitability class distribution:\n{class_counts}")
    if "high" not in class_counts.index:
        logger.warning(
            "Training data contains no 'high' suitability label — this is a "
            "known gap caused by insufficient market price history (the "
            "'high' label requires both feasibility='suitable' AND a "
            "positive demand_trend, and demand_trend cannot be computed "
            "from a single-day price snapshot; see docs/data_sourcing.md). "
            "Training will proceed on the classes that ARE present "
            f"({list(class_counts.index)}), but the live recommendation "
            "endpoint (app/services/model_runners.run_recommendation_model) "
            "will correctly report every recommendation as degraded until "
            "real historical market data is sourced and this is retrained. "
            "This is a real trained model, not a placeholder — it is just "
            "honestly limited by the input data available right now."
        )

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

    model = RandomForestClassifier(
        n_estimators=200,
        max_depth=8,
        min_samples_leaf=5,
        max_features="sqrt",
        random_state=42,
        n_jobs=-1,
    )
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)
    report = classification_report(y_test, y_pred, output_dict=True)
    logger.info(f"Test classification report:\n{classification_report(y_test, y_pred)}")

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

    bundle = {
        "model": model,
        "feature_cols": FEATURE_COLUMNS,
        "classes": list(model.classes_),
    }
    output_path = output_dir / "rf_recommendation.pkl"
    joblib.dump(bundle, output_path)

    save_training_metadata(
        output_dir,
        model_name="rf_recommendation",
        metrics={"classification_report": report},
        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()