"""
Train the XGBoost crop feasibility classifier.

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

import argparse
from pathlib import Path

import joblib
import xgboost as xgb
from sklearn.metrics import classification_report
from sklearn.model_selection import cross_val_score, train_test_split
from sklearn.preprocessing import LabelEncoder

import pandas as pd

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

logger = get_logger(__name__)

# Must exactly match app/services/feature_builder.py FEATURE_COLUMNS order.
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 = "feasibility_label"


def main() -> None:
    parser = argparse.ArgumentParser(description="Train XGBoost feasibility classifier")
    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: lowered from the production default of 200 to 100
    # for this initial Tamil Nadu pilot run (159 rows available after the
    # merge — see docs/data_sourcing.md and the project chat history for
    # context). This is a deliberate, visible decision, not a silent
    # weakening of the safety check. Models trained at this volume are
    # for proving the pipeline mechanically, NOT for client-facing
    # predictions — raise this back to 200+ once real district-level
    # data and a fuller crop/year range are sourced.
    check_minimum_samples(len(df), "xgb_feasibility", minimum=100)

    X = df[FEATURE_COLUMNS].values
    le = LabelEncoder()
    y = le.fit_transform(df[TARGET_COLUMN])

    class_counts = pd.Series(df[TARGET_COLUMN]).value_counts()
    logger.info(f"Class distribution:\n{class_counts}")
    if class_counts.min() < 20:
        logger.warning(
            f"Smallest class has only {class_counts.min()} samples. "
            f"Predictions for this class will be less reliable — collect "
            f"more labeled examples for it if possible."
        )

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

    model = xgb.XGBClassifier(
        n_estimators=300,
        max_depth=5,
        learning_rate=0.05,
        subsample=0.8,
        colsample_bytree=0.8,
        random_state=42,
        eval_metric="mlogloss",
    )

    model.fit(
        X_train, y_train,
        eval_set=[(X_test, y_test)],
        verbose=False,
    )

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

    cv_scores = cross_val_score(model, X, y, cv=5, scoring="accuracy")
    logger.info(f"5-fold CV accuracy: {cv_scores.mean():.3f} +/- {cv_scores.std():.3f}")

    if cv_scores.mean() < 0.6:
        logger.warning(
            f"Cross-validated accuracy is {cv_scores.mean():.1%}, which is low. "
            f"Review feature quality and label correctness before deploying "
            f"this model to production — it may not provide reliable guidance."
        )

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

    bundle = {"model": model, "label_encoder": le, "feature_cols": FEATURE_COLUMNS}
    output_path = output_dir / "xgb_feasibility.pkl"
    joblib.dump(bundle, output_path)

    save_training_metadata(
        output_dir,
        model_name="xgb_feasibility",
        metrics={
            "classification_report": report,
            "cv_accuracy_mean": float(cv_scores.mean()),
            "cv_accuracy_std": float(cv_scores.std()),
        },
        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()