test_api_connection.py

"""
Test API Connection with Trained Model

Verify that:
1. Model loads correctly
2. predict_soh() works
3. API endpoints work
4. End-to-end prediction flow
"""

import pandas as pd
import numpy as np
import pickle
import json
from predict_soh import predict_soh
from api_service import get_model_info_internal

def test_model_loading():
    """Test that model loads correctly."""
    print("="*80)
    print("TEST 1: MODEL LOADING")
    print("="*80)
    
    try:
        with open('./models/nasa_soh_model.pkl', 'rb') as f:
            model_data = pickle.load(f)
        
        model = model_data['model']
        feature_cols = model_data['feature_cols']
        metrics = model_data['metrics']
        
        print(f"✅ Model loaded successfully")
        print(f"   Model type: {type(model).__name__}")
        print(f"   Feature count: {len(feature_cols)}")
        print(f"   Validation R²: {metrics['val_r2']:.6f}")
        print(f"   Validation MAE: {metrics['val_mae']:.6f}")
        print(f"   OOB Score: {metrics.get('oob_score', 'N/A')}")
        
        return True, model_data
    except Exception as e:
        print(f"❌ Model loading failed: {e}")
        return False, None


def test_prediction_function():
    """Test predict_soh() function."""
    print("\n" + "="*80)
    print("TEST 2: PREDICTION FUNCTION")
    print("="*80)
    
    try:
        # Load test data (B0007, first 100 cycles)
        test_data = pd.read_csv('./data/nasa_converted/B0007.csv')
        test_subset = test_data[test_data['cycle_index'] <= 100].copy()
        
        print(f"Test data: {len(test_subset)} cycles")
        print(f"Columns: {test_subset.columns.tolist()}")
        
        # Run prediction
        predictions = predict_soh(
            test_subset,
            model_path='./models/nasa_soh_model.pkl',
            pipeline_path='./models/feature_pipeline_causal.pkl'
        )
        
        print(f"\n✅ Prediction successful!")
        print(f"   Current SOH: {predictions['current_soh']['point_estimate']:.4f}")
        print(f"   Health Score: {predictions['health_score']:.2f}/100")
        print(f"   Remaining Life: {predictions['remaining_life']['years']:.2f} years")
        print(f"   Degradation Rate: {predictions['degradation_rate']['per_year_pct']:.2f}%/year")
        print(f"   Data Quality: {predictions['data_quality']['quality_score']:.2f}")
        print(f"   Confidence: {predictions['confidence']:.4f}")
        
        # Test forecasts
        print(f"\n   Forecasts:")
        for key, forecast in predictions['forecasts'].items():
            cycles = key.replace('soh_at_', '').replace('_cycles', '')
            print(f"     {cycles} cycles: {forecast['point_estimate']:.4f} [{forecast['lower_95']:.4f}, {forecast['upper_95']:.4f}]")
        
        return True, predictions
    except Exception as e:
        print(f"❌ Prediction failed: {e}")
        import traceback
        traceback.print_exc()
        return False, None


def test_api_info():
    """Test API model info endpoint."""
    print("\n" + "="*80)
    print("TEST 3: API MODEL INFO")
    print("="*80)
    
    try:
        info = get_model_info_internal()
        
        print(f"✅ Model info retrieved successfully")
        print(f"   Model Type: {info['model_type']}")
        print(f"   Training Data: {info['training_data']}")
        print(f"   Validation Data: {info['validation_data']}")
        print(f"   R² Score: {info['metrics']['r2']:.6f}")
        print(f"   MAE: {info['metrics']['mae']:.6f}")
        print(f"   RMSE: {info['metrics']['rmse']:.6f}")
        print(f"   Feature Count: {info['feature_count']}")
        print(f"   Version: {info['version']}")
        
        return True, info
    except Exception as e:
        print(f"❌ API info failed: {e}")
        import traceback
        traceback.print_exc()
        return False, None


def test_full_pipeline():
    """Test complete prediction pipeline."""
    print("\n" + "="*80)
    print("TEST 4: FULL PREDICTION PIPELINE")
    print("="*80)
    
    try:
        # Simulate BMS data (new battery, not in training set)
        bms_data = pd.DataFrame({
            'cycle_index': list(range(1, 51)),
            'voltage': np.linspace(4.2, 3.8, 50) + np.random.normal(0, 0.02, 50),
            'current': np.random.choice([1.5, -1.5], 50),
            'temperature': 25 + np.random.normal(0, 1, 50),
            'capacity': np.linspace(2.0, 1.9, 50) + np.random.normal(0, 0.01, 50)
        })
        
        print(f"Simulated BMS data: {len(bms_data)} cycles")
        
        predictions = predict_soh(
            bms_data,
            model_path='./models/nasa_soh_model.pkl',
            pipeline_path='./models/feature_pipeline_causal.pkl',
            initial_capacity=2.0
        )
        
        print(f"\n✅ Full pipeline test successful!")
        print(f"   Current Cycle: {predictions['current_cycle']}")
        print(f"   Current SOH: {predictions['current_soh']['point_estimate']:.4f}")
        print(f"   95% Interval: [{predictions['current_soh']['lower_95']:.4f}, {predictions['current_soh']['upper_95']:.4f}]")
        print(f"   Health Score: {predictions['health_score']:.2f}/100")
        print(f"   RUL (cycles): {predictions['remaining_life']['cycles']:.1f}")
        print(f"   RUL (years): {predictions['remaining_life']['years']:.2f}")
        
        return True
    except Exception as e:
        print(f"❌ Full pipeline test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    print("="*80)
    print("EXPERT MODEL - API CONNECTION TESTS")
    print("="*80)
    
    results = {}
    
    # Test 1: Model loading
    success, model_data = test_model_loading()
    results['model_loading'] = success
    
    # Test 2: Prediction function
    success, predictions = test_prediction_function()
    results['prediction_function'] = success
    
    # Test 3: API info
    success, api_info = test_api_info()
    results['api_info'] = success
    
    # Test 4: Full pipeline
    success = test_full_pipeline()
    results['full_pipeline'] = success
    
    # Summary
    print("\n" + "="*80)
    print("TEST SUMMARY")
    print("="*80)
    
    for test_name, success in results.items():
        status = "✅ PASS" if success else "❌ FAIL"
        print(f"  {test_name}: {status}")
    
    all_passed = all(results.values())
    
    if all_passed:
        print("\n" + "="*80)
        print("✅ ALL TESTS PASSED - MODEL IS PRODUCTION-READY")
        print("="*80)
        print("\nNext steps:")
        print("  1. Start API service: python api_service.py")
        print("  2. API will be available at: http://localhost:8000")
        print("  3. Test endpoint: POST /api/predict")
        print("  4. Get model info: GET /api/model-info")
    else:
        print("\n" + "="*80)
        print("⚠️  SOME TESTS FAILED - CHECK ERRORS ABOVE")
        print("="*80)