"""
Baby Cry AI - Baseline Model
Step 3: Build a quick baseline model for cry classification
"""

import numpy as np
import pandas as pd
import sys
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.metrics import classification_report, confusion_matrix
import joblib
import os
from pathlib import Path
import json

from audio_processor import AudioProcessor

class BaselineModel:
    def __init__(self, model_path="models/baseline_model.pkl"):
        self.model = RandomForestClassifier(
            n_estimators=100,
            random_state=42,
            max_depth=10,
            min_samples_split=5,
            class_weight='balanced'  # Handle imbalanced classes
        )
        self.scaler = StandardScaler()
        self.label_encoder = LabelEncoder()
        # Use fast_mode for faster feature extraction (skips expensive CQT features)
        self.audio_processor = AudioProcessor(fast_mode=True)
        self.model_path = model_path
        self.is_trained = False
        self.expected_feature_names = None  # Feature names the model was trained with
        
        # Create models directory if it doesn't exist
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
    
    def load_data_from_directory(self, data_dir, balance_data=True, max_per_category=None):
        """Load and process audio data from directory structure
        
        Args:
            data_dir: Path to data directory
            balance_data: If True, limit samples to match smallest category
            max_per_category: Optional max samples per category
        """
        print("📂 Loading data from directory...")
        
        data_path = Path(data_dir)
        if not data_path.exists():
            print(f"❌ Data directory not found: {data_dir}")
            return None, None
        
        categories = [d for d in os.listdir(data_path) if os.path.isdir(data_path / d)]
        print(f"📁 Found categories: {categories}")
        
        # First pass: count files per category
        category_files = {}
        for category in categories:
            category_path = data_path / category
            audio_files = [f for f in os.listdir(category_path) 
                          if f.endswith(('.wav', '.mp3', '.m4a', '.flac'))]
            category_files[category] = audio_files
        
        # Determine max files per category for balancing
        if balance_data:
            min_count = min(len(files) for files in category_files.values())
            samples_per_category = min_count
            print(f"⚖️  Balancing data: {samples_per_category} samples per category")
        else:
            samples_per_category = max(len(files) for files in category_files.values())
        
        if max_per_category:
            samples_per_category = min(samples_per_category, max_per_category)
        
        features_list = []
        labels_list = []
        
        for category in categories:
            audio_files = category_files[category][:samples_per_category]
            
            print(f"🎵 Processing {category}: {len(audio_files)} files")
            sys.stdout.flush()
            
            file_count = 0
            for file in audio_files:
                file_count += 1
                if file_count % 10 == 0:
                    print(f"   Processed {file_count}/{len(audio_files)} files...")
                    sys.stdout.flush()
                file_path = data_path / category / file
                
                try:
                    # Extract features
                    features = self.audio_processor.extract_features_from_file(str(file_path))
                    
                    if features is not None:
                        features_list.append(list(features.values()))
                        labels_list.append(category)
                        
                except Exception as e:
                    print(f"⚠️ Error processing {file}: {e}")
                    continue
        
        if not features_list:
            print("❌ No valid audio files found")
            return None, None
        
        # Convert to numpy arrays
        X = np.array(features_list)
        y = np.array(labels_list)
        
        print(f"✅ Loaded {len(features_list)} samples with {X.shape[1]} features")
        return X, y
    
    def train(self, X, y):
        """Train the baseline model"""
        print("🚀 Training baseline model...")
        sys.stdout.flush()
        
        if X is None or y is None:
            print("❌ No data to train on")
            sys.stdout.flush()
            return False
        
        # Encode labels
        y_encoded = self.label_encoder.fit_transform(y)
        
        # Split data
        X_train, X_test, y_train, y_test = train_test_split(
            X, y_encoded, test_size=0.2, random_state=42, stratify=y_encoded
        )
        
        # Scale features
        X_train_scaled = self.scaler.fit_transform(X_train)
        X_test_scaled = self.scaler.transform(X_test)
        
        # Train model
        print("⏳ Fitting Random Forest model (this may take 1-2 minutes)...")
        sys.stdout.flush()
        self.model.fit(X_train_scaled, y_train)
        print("✅ Model fitting completed!")
        sys.stdout.flush()
        
        # Evaluate
        train_score = self.model.score(X_train_scaled, y_train)
        test_score = self.model.score(X_test_scaled, y_test)
        
        print(f"📊 Training Accuracy: {train_score:.3f}")
        print(f"📊 Test Accuracy: {test_score:.3f}")
        sys.stdout.flush()
        
        # Cross-validation
        cv_scores = cross_val_score(self.model, X_train_scaled, y_train, cv=5)
        print(f"📊 Cross-validation: {cv_scores.mean():.3f} (+/- {cv_scores.std() * 2:.3f})")
        
        # Classification report
        y_pred = self.model.predict(X_test_scaled)
        class_names = self.label_encoder.classes_
        
        print("\n📋 Classification Report:")
        print(classification_report(y_test, y_pred, target_names=class_names))
        
        self.is_trained = True
        self.test_accuracy = test_score
        
        # Save model
        self.save_model()
        
        return test_score
    
    def predict(self, features):
        """Make prediction on new features"""
        if not self.is_trained:
            print("❌ Model not trained yet")
            return None, None
        
        # Convert dict to list, aligning with expected feature names
        if isinstance(features, dict):
            if self.expected_feature_names is not None:
                # Align features to match model's expected features
                aligned_features = []
                for feat_name in self.expected_feature_names:
                    if feat_name in features:
                        aligned_features.append(features[feat_name])
                    else:
                        # Feature missing - use 0 as default (could also use mean)
                        print(f"⚠️  Missing feature: {feat_name}, using 0")
                        aligned_features.append(0.0)
                features = aligned_features
            else:
                # Fallback: just use values in order (may cause mismatch)
                features = list(features.values())
        
        features = np.array(features).reshape(1, -1)
        
        # Check feature count matches
        # Try multiple ways to get expected feature count
        expected_count = None
        if hasattr(self.scaler, 'n_features_in_'):
            expected_count = self.scaler.n_features_in_
        elif hasattr(self.scaler, 'mean_') and self.scaler.mean_ is not None:
            expected_count = len(self.scaler.mean_)
        elif self.expected_feature_names:
            expected_count = len(self.expected_feature_names)
        
        if expected_count and features.shape[1] != expected_count:
            print(f"❌ Feature count mismatch: got {features.shape[1]}, expected {expected_count}")
            if self.expected_feature_names:
                print(f"   Expected features: {self.expected_feature_names[:10]}... (showing first 10)")
            return None, None
        
        features_scaled = self.scaler.transform(features)
        
        prediction = self.model.predict(features_scaled)[0]
        probability = self.model.predict_proba(features_scaled)[0]
        
        # Convert back to label
        label = self.label_encoder.inverse_transform([prediction])[0]
        confidence = np.max(probability)
        
        return label, confidence
    
    def predict_top_k(self, features, k=2):
        """Return the top-k (label, probability) pairs for a feature dict/list."""
        if not self.is_trained:
            print("❌ Model not trained yet")
            return None

        if isinstance(features, dict):
            if self.expected_feature_names is not None:
                features = [features.get(name, 0.0) for name in self.expected_feature_names]
            else:
                features = list(features.values())

        features = np.array(features).reshape(1, -1)
        features_scaled = self.scaler.transform(features)
        probs = self.model.predict_proba(features_scaled)[0]
        order = np.argsort(probs)[::-1][:k]
        labels = self.label_encoder.inverse_transform(order)
        return [(str(label), float(probs[i])) for label, i in zip(labels, order)]

    def predict_top_k_from_audio_file(self, file_path, k=2):
        """Top-k predictions straight from an audio file."""
        features = self.audio_processor.extract_features_from_file(file_path)
        if features is None:
            return None
        return self.predict_top_k(features, k=k)

    def predict_from_audio_file(self, file_path):
        """Predict from audio file"""
        features = self.audio_processor.extract_features_from_file(file_path)
        
        if features is None:
            return None, None
        
        return self.predict(features)
    
    def predict_from_audio_array(self, y, sr):
        """Predict from audio array"""
        features = self.audio_processor.extract_features_from_array(y, sr)
        
        if features is None:
            return None, None
        
        return self.predict(features)
    
    def save_model(self):
        """Save the trained model"""
        model_data = {
            'model': self.model,
            'scaler': self.scaler,
            'label_encoder': self.label_encoder,
            'feature_names': self.audio_processor.get_feature_names(),
            'is_trained': self.is_trained
        }
        
        joblib.dump(model_data, self.model_path)
        print(f"💾 Model saved to: {self.model_path}")
    
    def load_model(self):
        """Load a pre-trained model"""
        if not os.path.exists(self.model_path):
            print(f"❌ Model file not found: {self.model_path}")
            return False
        
        try:
            model_data = joblib.load(self.model_path)
            
            self.model = model_data['model']
            self.scaler = model_data['scaler']
            self.label_encoder = model_data['label_encoder']
            self.is_trained = model_data['is_trained']
            
            # Load expected feature names (for feature alignment)
            if 'feature_names' in model_data:
                self.expected_feature_names = model_data['feature_names']
                print(f"✅ Model loaded from: {self.model_path}")
                print(f"   Expected features: {len(self.expected_feature_names)}")
            else:
                # Fallback: try to infer from scaler
                if hasattr(self.scaler, 'n_features_in_'):
                    print(f"⚠️  Model doesn't have feature_names, using scaler dimension: {self.scaler.n_features_in_}")
                else:
                    print(f"⚠️  Model doesn't have feature_names, feature alignment may fail")
            
            return True
            
        except Exception as e:
            print(f"❌ Error loading model: {e}")
            return False
    
    def get_model_info(self):
        """Get information about the model"""
        if not self.is_trained:
            return {"status": "Not trained"}
        
        info = {
            "status": "Trained",
            "model_type": "Random Forest",
            "n_estimators": self.model.n_estimators,
            "feature_count": len(self.audio_processor.get_feature_names()),
            "classes": self.label_encoder.classes_.tolist(),
            "model_path": self.model_path
        }
        
        return info

# Example usage and testing
if __name__ == "__main__":
    print("🤖 Baseline Model Test")
    print("=" * 30)
    
    # Initialize model
    model = BaselineModel()
    
    # Try to load existing model first
    if model.load_model():
        print("✅ Loaded existing model")
        info = model.get_model_info()
        print(f"📊 Model info: {info}")
    else:
        print("🔄 No existing model found, training new one...")
        
        # Load data (you'll need to have data in the data/ directory)
        X, y = model.load_data_from_directory("data")
        
        if X is not None and y is not None:
            # Train model
            success = model.train(X, y)
            
            if success:
                print("✅ Model training completed!")
                info = model.get_model_info()
                print(f"📊 Model info: {info}")
            else:
                print("❌ Model training failed")
        else:
            print("❌ Could not load data for training")
            print("📝 Please ensure you have audio data in the data/ directory")
            print("   Expected structure:")
            print("   data/")
            print("   ├── hunger/")
            print("   ├── sleep/")
            print("   ├── discomfort/")
            print("   └── pain/")