"""
Train the Decision Tree trend analysis classifier.

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

import argparse
from pathlib import Path

import joblib
from sklearn.metrics import accuracy_score, classification_report
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, export_text

import pandas as pd

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

logger = get_logger(__name__)

# Deliberately a different, smaller feature set than the main FEATURE_COLUMNS —
# trend analysis is about multi-year patterns, not the current month's weather.
# app/services/model_runners.run_trend_model passes the full feature_vector
# and relies on this model having been trained on a vector of the SAME
# length and order as app/services/feature_builder.FEATURE_COLUMNS, so we
# train on the full set here too for v1 simplicity. If you want the
# original smaller research feature set described in the design docs,
# update both this script and app/services/feature_builder.py together.
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 = "trend_label"


def main() -> None:
    parser = argparse.ArgumentParser(description="Train Decision Tree trend 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])
    check_minimum_samples(len(df), "dt_trend", minimum=150)

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

    class_counts = pd.Series(y).value_counts()
    logger.info(f"Trend class distribution:\n{class_counts}")

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

    model = DecisionTreeClassifier(
        max_depth=6,
        min_samples_split=20,
        min_samples_leaf=10,
        random_state=42,
    )
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    report = classification_report(y_test, y_pred, output_dict=True)
    logger.info(f"Test accuracy: {accuracy:.3f}")
    logger.info(f"Decision tree structure:\n{export_text(model, feature_names=FEATURE_COLUMNS)}")

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

    bundle = {
        "model": model,
        "feature_cols": FEATURE_COLUMNS,
        "source_feature_cols": FEATURE_COLUMNS,
    }
    output_path = output_dir / "dt_trend.pkl"
    joblib.dump(bundle, output_path)

    save_training_metadata(
        output_dir,
        model_name="dt_trend",
        metrics={"accuracy": float(accuracy), "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()