train_model.py

"""
BMS AI Model Training
Trains a Random Forest classifier to predict battery condition or application suitability.

Usage:
    python train_model.py                              # Train condition model
    python train_model.py --target application         # Train application model
    python train_model.py --source battery_data.csv    # Train on specific CSV
"""

import argparse
import joblib
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix

from data_pipeline import load_data, preprocess, FEATURE_COLUMNS


def train(data_source="synthetic", num_samples=5000, target="condition"):
    """Train the BMS classifier and save model artifacts."""

    target_label = "Condition" if target == "condition" else "Application Suitability"
    print("=" * 60)
    print(f"  🔋 BMS AI Model Training — {target_label}")
    print("=" * 60)

    # --- 1. Load Data ---
    print(f"\n📦 Loading data (source: {data_source})...")
    df = load_data(data_source, num_samples=num_samples)
    print(f"   Dataset: {df.shape[0]} samples, {df.shape[1]} features")
    print(f"   {target_label} distribution:")
    for cond, count in df[target].value_counts().items():
        pct = count / len(df) * 100
        print(f"     {cond:>18s}: {count:>5d} ({pct:.1f}%)")

    # --- 2. Preprocess ---
    print("\n⚙️  Preprocessing...")
    X, y, label_encoder = preprocess(df, target=target)
    available_features = [col for col in FEATURE_COLUMNS if col in df.columns]
    print(f"   Features used: {available_features}")
    print(f"   Classes: {list(label_encoder.classes_)}")

    # --- 3. Train/Test Split ---
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    print(f"\n📊 Split: {len(X_train)} train / {len(X_test)} test")

    # --- 4. Scale Features ---
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # --- 5. Train Model ---
    print("\n🧠 Training Random Forest classifier...")
    model = RandomForestClassifier(
        n_estimators=200,
        max_depth=15,
        min_samples_split=5,
        min_samples_leaf=2,
        class_weight="balanced",
        random_state=42,
        n_jobs=-1,
    )
    model.fit(X_train_scaled, y_train)

    # --- 6. Evaluate ---
    y_pred = model.predict(X_test_scaled)
    accuracy = accuracy_score(y_test, y_pred)

    print(f"\n{'=' * 60}")
    print(f"  📈 MODEL EVALUATION RESULTS")
    print(f"{'=' * 60}")
    print(f"\n  Accuracy: {accuracy:.4f} ({accuracy * 100:.1f}%)")

    print(f"\n  Classification Report:")
    target_names = list(label_encoder.classes_)
    report = classification_report(y_test, y_pred, target_names=target_names)
    for line in report.split("\n"):
        print(f"    {line}")

    print(f"\n  Confusion Matrix:")
    cm = confusion_matrix(y_test, y_pred)
    print(f"    {'':>10s} ", end="")
    for name in target_names:
        print(f"{name:>10s}", end="")
    print()
    for i, row in enumerate(cm):
        print(f"    {target_names[i]:>10s} ", end="")
        for val in row:
            print(f"{val:>10d}", end="")
        print()

    # --- 7. Feature Importance ---
    print(f"\n  Feature Importance:")
    importances = model.feature_importances_
    sorted_idx = np.argsort(importances)[::-1]
    for idx in sorted_idx:
        bar = "█" * int(importances[idx] * 40)
        print(f"    {available_features[idx]:>16s}: {importances[idx]:.4f} {bar}")

    # --- 8. Save Artifacts ---
    if target == "condition":
        model_file, scaler_file, encoder_file = "bms_model.pkl", "scaler.pkl", "label_encoder.pkl"
    else:
        model_file, scaler_file, encoder_file = "app_model.pkl", "app_scaler.pkl", "app_label_encoder.pkl"

    print(f"\n💾 Saving model artifacts...")
    joblib.dump(model, model_file)
    joblib.dump(scaler, scaler_file)
    joblib.dump(label_encoder, encoder_file)
    print(f"   ✅ {model_file}")
    print(f"   ✅ {scaler_file}")
    print(f"   ✅ {encoder_file}")

    print(f"\n{'=' * 60}")
    print(f"  ✅ Training complete! Accuracy: {accuracy * 100:.1f}%")
    print(f"{'=' * 60}")

    return model, scaler, label_encoder, accuracy


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train BMS AI Model")
    parser.add_argument(
        "--source",
        default="synthetic",
        help="Data source: 'synthetic' or path to CSV file",
    )
    parser.add_argument(
        "--samples",
        type=int,
        default=5000,
        help="Number of synthetic samples (only used with synthetic source)",
    )
    parser.add_argument(
        "--target",
        default="condition",
        choices=["condition", "application"],
        help="Target to train: 'condition' or 'application'",
    )
    args = parser.parse_args()
    train(data_source=args.source, num_samples=args.samples, target=args.target)