src/models/ensemble_model.py

"""
Baby Cry AI - Ensemble Model
Combines Random Forest and Neural Network predictions
"""

import os
import sys
import numpy as np
from pathlib import Path
import warnings
warnings.filterwarnings('ignore')

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from models.baseline_model import BaselineModel
from models.neural_model import NeuralModel
from audio_processor import AudioProcessor


class EnsembleModel:
    """Ensemble model combining Random Forest and Neural Network"""
    
    def __init__(self, 
                 rf_model_path="models/baseline_model.pkl",
                 nn_model_path="models/neural_model.h5",
                 rf_weight=0.4,
                 nn_weight=0.6):
        """
        Initialize ensemble model
        
        Args:
            rf_model_path: Path to Random Forest model
            nn_model_path: Path to Neural Network model
            rf_weight: Weight for Random Forest predictions (0-1)
            nn_weight: Weight for Neural Network predictions (0-1)
        """
        self.rf_model = BaselineModel(model_path=rf_model_path)
        self.nn_model = NeuralModel(model_path=nn_model_path)
        self.audio_processor = AudioProcessor()
        
        # Normalize weights
        total_weight = rf_weight + nn_weight
        self.rf_weight = rf_weight / total_weight
        self.nn_weight = nn_weight / total_weight
        
        self.is_trained = False
        self.rf_loaded = False
        self.nn_loaded = False
    
    def load_models(self):
        """Load both models"""
        print("📦 Loading ensemble models...")
        
        # Load Random Forest
        if os.path.exists(self.rf_model.model_path):
            self.rf_loaded = self.rf_model.load_model()
            if self.rf_loaded:
                print("  ✅ Random Forest model loaded")
            else:
                print("  ⚠️  Failed to load Random Forest model")
        else:
            print(f"  ⚠️  Random Forest model not found: {self.rf_model.model_path}")
        
        # Load Neural Network
        if os.path.exists(self.nn_model.model_path):
            self.nn_loaded = self.nn_model.load_model()
            if self.nn_loaded:
                print("  ✅ Neural Network model loaded")
            else:
                print("  ⚠️  Failed to load Neural Network model")
        else:
            print(f"  ⚠️  Neural Network model not found: {self.nn_model.model_path}")
        
        self.is_trained = self.rf_loaded or self.nn_loaded
        
        if not self.is_trained:
            print("  ❌ No models loaded")
        else:
            print(f"  📊 Ensemble ready (RF: {self.rf_weight:.2f}, NN: {self.nn_weight:.2f})")
        
        return self.is_trained
    
    def predict(self, features_or_file, is_file_path=True):
        """
        Make ensemble prediction
        
        Args:
            features_or_file: Either file path or feature dict
            is_file_path: If True, treat as file path; else as features
        
        Returns:
            tuple: (prediction, confidence)
        """
        if not self.is_trained:
            if not self.load_models():
                return None, None
        
        predictions = {}
        confidences = {}
        
        # Random Forest prediction
        if self.rf_loaded:
            try:
                if is_file_path:
                    rf_pred, rf_conf = self.rf_model.predict_from_audio_file(features_or_file)
                else:
                    rf_pred, rf_conf = self.rf_model.predict(features_or_file)
                
                if rf_pred is not None:
                    predictions['rf'] = rf_pred
                    confidences['rf'] = rf_conf
            except Exception as e:
                print(f"  ⚠️  RF prediction error: {e}")
        
        # Neural Network prediction
        if self.nn_loaded:
            try:
                if is_file_path:
                    nn_pred, nn_conf = self.nn_model.predict_from_audio_file(features_or_file)
                else:
                    # For NN, we need mel-spectrogram, not features
                    # So we'll skip if not file path
                    nn_pred, nn_conf = None, None
                
                if nn_pred is not None:
                    predictions['nn'] = nn_pred
                    confidences['nn'] = nn_conf
            except Exception as e:
                print(f"  ⚠️  NN prediction error: {e}")
        
        # Combine predictions
        if not predictions:
            return None, None
        
        # Weighted voting
        if len(predictions) == 1:
            # Only one model available
            model_name = list(predictions.keys())[0]
            return predictions[model_name], confidences[model_name]
        
        # Both models available - weighted voting
        vote_counts = {}
        weighted_confidences = {}
        
        for model_name, pred in predictions.items():
            weight = self.rf_weight if model_name == 'rf' else self.nn_weight
            conf = confidences[model_name] * weight
            
            if pred not in vote_counts:
                vote_counts[pred] = 0
                weighted_confidences[pred] = 0
            
            vote_counts[pred] += weight
            weighted_confidences[pred] += conf
        
        # Get prediction with highest weighted vote
        ensemble_pred = max(vote_counts, key=vote_counts.get)
        ensemble_conf = weighted_confidences[ensemble_pred] / vote_counts[ensemble_pred]
        
        return ensemble_pred, ensemble_conf
    
    def predict_from_audio_file(self, file_path):
        """Predict from audio file"""
        return self.predict(file_path, is_file_path=True)
    
    def predict_from_features(self, features):
        """Predict from extracted features"""
        return self.predict(features, is_file_path=False)
    
    def get_model_info(self):
        """Get information about the ensemble"""
        info = {
            "status": "Ready" if self.is_trained else "Not ready",
            "model_type": "Ensemble (RF + NN)",
            "rf_loaded": self.rf_loaded,
            "nn_loaded": self.nn_loaded,
            "rf_weight": self.rf_weight,
            "nn_weight": self.nn_weight
        }
        
        if self.rf_loaded:
            rf_info = self.rf_model.get_model_info()
            info['rf_info'] = rf_info
        
        if self.nn_loaded:
            nn_info = self.nn_model.get_model_info()
            info['nn_info'] = nn_info
        
        return info


if __name__ == "__main__":
    print("🤖 Ensemble Model Test")
    print("=" * 50)
    
    # Initialize ensemble
    ensemble = EnsembleModel()
    
    # Load models
    if ensemble.load_models():
        info = ensemble.get_model_info()
        print(f"\n📊 Ensemble info:")
        for key, value in info.items():
            print(f"   {key}: {value}")
        
        # Test prediction (if test file exists)
        test_file = "../data/hunger/hunger_001.wav"
        if os.path.exists(test_file):
            print(f"\n🧪 Testing prediction on {test_file}...")
            pred, conf = ensemble.predict_from_audio_file(test_file)
            if pred:
                print(f"   Prediction: {pred}")
                print(f"   Confidence: {conf:.3f}")
    else:
        print("❌ Could not load models")