predict.py

"""
BMS Battery Condition Predictor
Loads trained model artifacts and provides prediction API for dashboard and hardware integration.

Usage:
    from predict import BatteryPredictor

    predictor = BatteryPredictor()
    condition, confidence = predictor.predict(
        voltage=3.7, current=1.2, temperature=28, soc=75, cycle_count=150
    )
    print(f"Condition: {condition}, Confidence: {confidence}")
"""

import os
import numpy as np
import pandas as pd
import joblib


class BatteryPredictor:
    """
    Battery condition classifier using trained ML model.

    Loads model artifacts (model, scaler, label encoder) and provides
    methods for single-sample and batch predictions.
    """

    def __init__(self, model_dir="."):
        """
        Initialize the predictor by loading saved model artifacts.

        Args:
            model_dir: Directory containing model .pkl files
        """
        self.model_dir = model_dir
        self.model = None
        self.scaler = None
        self.label_encoder = None
        self.app_model = None
        self.app_scaler = None
        self.app_label_encoder = None
        self._load_artifacts()

    def _load_artifacts(self):
        """Load model, scaler, and label encoder from disk."""
        # --- Condition model (required) ---
        model_path = os.path.join(self.model_dir, "bms_model.pkl")
        scaler_path = os.path.join(self.model_dir, "scaler.pkl")
        encoder_path = os.path.join(self.model_dir, "label_encoder.pkl")

        for path, name in [
            (model_path, "Model"),
            (scaler_path, "Scaler"),
            (encoder_path, "Label Encoder"),
        ]:
            if not os.path.exists(path):
                raise FileNotFoundError(
                    f"{name} not found at {path}. "
                    f"Run 'python train_model.py' first to train the model."
                )

        self.model = joblib.load(model_path)
        self.scaler = joblib.load(scaler_path)
        self.label_encoder = joblib.load(encoder_path)

        # --- Application model (optional) ---
        app_model_path = os.path.join(self.model_dir, "app_model.pkl")
        app_scaler_path = os.path.join(self.model_dir, "app_scaler.pkl")
        app_encoder_path = os.path.join(self.model_dir, "app_label_encoder.pkl")

        if all(os.path.exists(p) for p in [app_model_path, app_scaler_path, app_encoder_path]):
            self.app_model = joblib.load(app_model_path)
            self.app_scaler = joblib.load(app_scaler_path)
            self.app_label_encoder = joblib.load(app_encoder_path)

    def predict(self, voltage, current, temperature, soc=50.0, cycle_count=100, voltage_delta=0.0):
        """
        Predict battery condition from sensor values.

        Args:
            voltage: Battery voltage (V)
            current: Battery current (A)
            temperature: Battery temperature (°C)
            soc: State of charge (%)
            cycle_count: Number of charge cycles
            voltage_delta: Rate of voltage change

        Returns:
            tuple: (condition_label: str, confidence_dict: dict)
                - condition_label: "Normal", "Warning", or "Unsafe"
                - confidence_dict: Probability for each class
        """
        features = np.array([[voltage, current, temperature, soc, cycle_count, voltage_delta]])
        features_scaled = self.scaler.transform(features)

        prediction = self.model.predict(features_scaled)[0]
        probabilities = self.model.predict_proba(features_scaled)[0]

        condition_label = self.label_encoder.inverse_transform([prediction])[0]
        confidence_dict = {
            label: round(float(prob) * 100, 1)
            for label, prob in zip(self.label_encoder.classes_, probabilities)
        }

        return condition_label, confidence_dict

    def predict_application(self, voltage, current, temperature, soc=50.0, cycle_count=100, voltage_delta=0.0):
        """
        Predict suitable application category for a battery.

        Returns:
            tuple: (application_label: str, confidence_dict: dict)
                - application_label: "High-Performance", "Moderate-Load", "Low-Power", or "Unsuitable"
                - confidence_dict: Probability for each class
        """
        if self.app_model is None:
            return None, {}

        features = np.array([[voltage, current, temperature, soc, cycle_count, voltage_delta]])
        features_scaled = self.app_scaler.transform(features)

        prediction = self.app_model.predict(features_scaled)[0]
        probabilities = self.app_model.predict_proba(features_scaled)[0]

        app_label = self.app_label_encoder.inverse_transform([prediction])[0]
        confidence_dict = {
            label: round(float(prob) * 100, 1)
            for label, prob in zip(self.app_label_encoder.classes_, probabilities)
        }

        return app_label, confidence_dict

    def predict_batch(self, df):
        """
        Predict conditions for a batch of readings.

        Args:
            df: DataFrame with columns matching FEATURE_COLUMNS

        Returns:
            DataFrame with added 'predicted_condition' and confidence columns
        """
        from data_pipeline import FEATURE_COLUMNS

        available_features = [col for col in FEATURE_COLUMNS if col in df.columns]
        X = df[available_features].values.astype(np.float64)

        # Pad with zeros if some features are missing
        if X.shape[1] < 6:
            padding = np.zeros((X.shape[0], 6 - X.shape[1]))
            X = np.hstack([X, padding])

        X_scaled = self.scaler.transform(X)

        predictions = self.model.predict(X_scaled)
        probabilities = self.model.predict_proba(X_scaled)

        result = df.copy()
        result["predicted_condition"] = self.label_encoder.inverse_transform(predictions)

        for i, label in enumerate(self.label_encoder.classes_):
            result[f"confidence_{label}"] = (probabilities[:, i] * 100).round(1)

        return result


if __name__ == "__main__":
    print("🔋 BMS Battery Predictor — Self Test\n")

    try:
        predictor = BatteryPredictor()
        print("✅ Model artifacts loaded successfully\n")
    except FileNotFoundError as e:
        print(f"❌ {e}")
        print("   Run: python train_model.py")
        exit(1)

    # Test cases
    test_cases = [
        {"voltage": 3.7, "current": 1.0, "temperature": 25, "soc": 75, "cycle_count": 100, "expected": "Normal"},
        {"voltage": 3.1, "current": 3.0, "temperature": 45, "soc": 20, "cycle_count": 300, "expected": "Warning"},
        {"voltage": 2.6, "current": 4.5, "temperature": 65, "soc": 5, "cycle_count": 500, "expected": "Unsafe"},
        {"voltage": 4.35, "current": 0.5, "temperature": 30, "soc": 95, "cycle_count": 50, "expected": "Warning"},
        {"voltage": 4.6, "current": 4.2, "temperature": 60, "soc": 3, "cycle_count": 800, "expected": "Unsafe"},
    ]

    print(f"{'Test':>5s} | {'Voltage':>8s} {'Current':>8s} {'Temp':>6s} {'SoC':>5s} | {'Expected':>10s} → {'Predicted':>10s} | {'Confidence':>12s} | {'Pass':>4s}")
    print("-" * 95)

    all_passed = True
    for i, tc in enumerate(test_cases, 1):
        condition, conf = predictor.predict(
            voltage=tc["voltage"],
            current=tc["current"],
            temperature=tc["temperature"],
            soc=tc["soc"],
            cycle_count=tc["cycle_count"],
        )
        passed = condition == tc["expected"]
        if not passed:
            all_passed = False
        top_conf = max(conf.values())
        status = "✅" if passed else "⚠️"

        print(
            f"  {i:>3d} | {tc['voltage']:>7.2f}V {tc['current']:>7.2f}A {tc['temperature']:>5.1f}°C {tc['soc']:>4.1f}% "
            f"| {tc['expected']:>10s} → {condition:>10s} | {top_conf:>10.1f}% | {status}"
        )

    print(f"\n{'✅ All condition tests passed!' if all_passed else '⚠️  Some predictions did not match expected labels (may be acceptable depending on model confidence)'}")

    # --- Application suitability tests ---
    if predictor.app_model is not None:
        print("\n" + "=" * 95)
        print("🎯 APPLICATION SUITABILITY PREDICTIONS")
        print("=" * 95)

        app_tests = [
            {"voltage": 3.8, "current": 0.8, "temperature": 25, "soc": 90, "cycle_count": 80},
            {"voltage": 3.5, "current": 2.0, "temperature": 35, "soc": 55, "cycle_count": 350},
            {"voltage": 3.2, "current": 3.0, "temperature": 45, "soc": 20, "cycle_count": 600},
            {"voltage": 2.7, "current": 4.5, "temperature": 60, "soc": 5,  "cycle_count": 900},
        ]

        for i, tc in enumerate(app_tests, 1):
            app_label, app_conf = predictor.predict_application(
                voltage=tc["voltage"], current=tc["current"],
                temperature=tc["temperature"], soc=tc["soc"],
                cycle_count=tc["cycle_count"],
            )
            top_conf = max(app_conf.values())
            print(f"  {i:>3d} | {tc['voltage']:.2f}V {tc['current']:.1f}A {tc['temperature']}°C SoC:{tc['soc']}% Cycles:{tc['cycle_count']} → {app_label} ({top_conf:.1f}%)")
    else:
        print("\n⚠️  Application model not found. Run: python train_model.py --target application")