src/hyperparameter_tuner.py

"""
Baby Cry AI - Hyperparameter Tuning
Uses Optuna for systematic hyperparameter optimization
"""

import os
import sys
import numpy as np
import optuna
from sklearn.model_selection import cross_val_score
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
import warnings
warnings.filterwarnings('ignore')

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

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


class HyperparameterTuner:
    """Hyperparameter tuning for models"""
    
    def __init__(self, data_dir="../data"):
        """
        Initialize tuner
        
        Args:
            data_dir: Path to data directory
        """
        self.data_dir = data_dir
        self.processor = AudioProcessor()
        self.best_params = {}
        self.best_score = 0
    
    def tune_random_forest(self, X, y, n_trials=50):
        """
        Tune Random Forest hyperparameters
        
        Args:
            X: Feature matrix
            y: Labels
            n_trials: Number of optimization trials
        """
        print("🔍 Tuning Random Forest hyperparameters...")
        
        def objective(trial):
            # Suggest hyperparameters
            n_estimators = trial.suggest_int('n_estimators', 50, 300, step=50)
            max_depth = trial.suggest_int('max_depth', 5, 30, step=5)
            min_samples_split = trial.suggest_int('min_samples_split', 2, 10)
            min_samples_leaf = trial.suggest_int('min_samples_leaf', 1, 5)
            max_features = trial.suggest_categorical('max_features', ['sqrt', 'log2', None])
            
            # Create model
            model = RandomForestClassifier(
                n_estimators=n_estimators,
                max_depth=max_depth,
                min_samples_split=min_samples_split,
                min_samples_leaf=min_samples_leaf,
                max_features=max_features,
                class_weight='balanced',
                random_state=42,
                n_jobs=-1
            )
            
            # Scale features
            scaler = StandardScaler()
            X_scaled = scaler.fit_transform(X)
            
            # Cross-validation score
            scores = cross_val_score(model, X_scaled, y, cv=5, scoring='accuracy', n_jobs=-1)
            return scores.mean()
        
        study = optuna.create_study(direction='maximize')
        study.optimize(objective, n_trials=n_trials, show_progress_bar=True)
        
        self.best_params['random_forest'] = study.best_params
        self.best_score = study.best_value
        
        print(f"✅ Best Random Forest score: {study.best_value:.4f}")
        print(f"📊 Best parameters: {study.best_params}")
        
        return study.best_params, study.best_value
    
    def tune_neural_network(self, X, y, n_trials=20):
        """
        Tune Neural Network hyperparameters
        
        Args:
            X: Mel-spectrograms
            y: Labels
            n_trials: Number of optimization trials
        """
        print("🔍 Tuning Neural Network hyperparameters...")
        
        def objective(trial):
            # Suggest hyperparameters
            dropout_rate = trial.suggest_float('dropout_rate', 0.3, 0.7)
            learning_rate = trial.suggest_loguniform('learning_rate', 1e-5, 1e-2)
            batch_size = trial.suggest_categorical('batch_size', [16, 32, 64])
            num_conv_filters_1 = trial.suggest_int('num_conv_filters_1', 16, 64, step=16)
            num_conv_filters_2 = trial.suggest_int('num_conv_filters_2', 32, 128, step=32)
            num_dense_units = trial.suggest_int('num_dense_units', 128, 512, step=128)
            
            # Create and train model
            model = NeuralModel()
            model.input_shape = X[0].shape
            
            # Build custom model with suggested parameters
            import tensorflow as tf
            from tensorflow import keras
            from tensorflow.keras import layers
            
            num_classes = len(np.unique(y))
            from sklearn.preprocessing import LabelEncoder
            le = LabelEncoder()
            y_encoded = le.fit_transform(y)
            
            nn_model = keras.Sequential([
                layers.Conv2D(num_conv_filters_1, (3, 3), activation='relu', input_shape=X[0].shape),
                layers.BatchNormalization(),
                layers.MaxPooling2D((2, 2)),
                layers.Dropout(dropout_rate * 0.5),
                
                layers.Conv2D(num_conv_filters_2, (3, 3), activation='relu'),
                layers.BatchNormalization(),
                layers.MaxPooling2D((2, 2)),
                layers.Dropout(dropout_rate),
                
                layers.GlobalAveragePooling2D(),
                layers.Dense(num_dense_units, activation='relu'),
                layers.BatchNormalization(),
                layers.Dropout(dropout_rate),
                layers.Dense(num_classes, activation='softmax')
            ])
            
            nn_model.compile(
                optimizer=keras.optimizers.Adam(learning_rate=learning_rate),
                loss='sparse_categorical_crossentropy',
                metrics=['accuracy']
            )
            
            # Train with early stopping
            from sklearn.model_selection import train_test_split
            X_train, X_val, y_train, y_val = train_test_split(
                X, y_encoded, test_size=0.2, random_state=42, stratify=y_encoded
            )
            
            from tensorflow.keras import callbacks
            early_stop = callbacks.EarlyStopping(
                monitor='val_loss',
                patience=5,
                restore_best_weights=True,
                verbose=0
            )
            
            try:
                nn_model.fit(
                    X_train, y_train,
                    batch_size=batch_size,
                    epochs=20,
                    validation_data=(X_val, y_val),
                    callbacks=[early_stop],
                    verbose=0
                )
                
                # Evaluate
                val_loss, val_acc = nn_model.evaluate(X_val, y_val, verbose=0)
                return val_acc
            except Exception as e:
                print(f"  ⚠️  Trial failed: {e}")
                return 0.0
        
        study = optuna.create_study(direction='maximize')
        study.optimize(objective, n_trials=n_trials, show_progress_bar=True)
        
        self.best_params['neural_network'] = study.best_params
        self.best_score = study.best_value
        
        print(f"✅ Best Neural Network score: {study.best_value:.4f}")
        print(f"📊 Best parameters: {study.best_params}")
        
        return study.best_params, study.best_value
    
    def tune_audio_processor(self, file_paths, labels, n_trials=30):
        """
        Tune audio processing parameters
        
        Args:
            file_paths: List of audio file paths
            labels: Corresponding labels
            n_trials: Number of optimization trials
        """
        print("🔍 Tuning Audio Processor hyperparameters...")
        
        def objective(trial):
            # Suggest preprocessing parameters
            highpass_cutoff = trial.suggest_int('highpass_cutoff', 50, 200, step=50)
            trim_top_db = trial.suggest_int('trim_top_db', 20, 40, step=5)
            n_mfcc = trial.suggest_int('n_mfcc', 10, 20, step=2)
            
            # Create processor with suggested parameters
            processor = AudioProcessor()
            # Note: These would need to be configurable in AudioProcessor
            # For now, we'll use a simplified approach
            
            # Extract features and evaluate
            try:
                features_list = []
                labels_list = []
                
                for file_path, label in zip(file_paths[:50], labels[:50]):  # Limit for speed
                    features = processor.extract_features_from_file(str(file_path))
                    if features is not None:
                        features_list.append(list(features.values()))
                        labels_list.append(label)
                
                if len(features_list) < 10:
                    return 0.0
                
                X = np.array(features_list)
                y = np.array(labels_list)
                
                # Quick evaluation with simple model
                from sklearn.ensemble import RandomForestClassifier
                from sklearn.model_selection import cross_val_score
                from sklearn.preprocessing import StandardScaler
                
                scaler = StandardScaler()
                X_scaled = scaler.fit_transform(X)
                
                model = RandomForestClassifier(n_estimators=50, random_state=42, n_jobs=-1)
                scores = cross_val_score(model, X_scaled, y, cv=3, scoring='accuracy', n_jobs=-1)
                
                return scores.mean()
            except Exception as e:
                return 0.0
        
        study = optuna.create_study(direction='maximize')
        study.optimize(objective, n_trials=n_trials, show_progress_bar=True)
        
        self.best_params['audio_processor'] = study.best_params
        self.best_score = study.best_value
        
        print(f"✅ Best Audio Processor score: {study.best_value:.4f}")
        print(f"📊 Best parameters: {study.best_params}")
        
        return study.best_params, study.best_value
    
    def get_best_params(self):
        """Get best parameters found"""
        return self.best_params
    
    def save_results(self, output_path="hyperparameter_tuning_results.json"):
        """Save tuning results"""
        import json
        from datetime import datetime
        
        results = {
            'timestamp': datetime.now().isoformat(),
            'best_params': self.best_params,
            'best_score': self.best_score
        }
        
        with open(output_path, 'w') as f:
            json.dump(results, f, indent=2)
        
        print(f"💾 Results saved to: {output_path}")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='Hyperparameter tuning')
    parser.add_argument('--data-dir', type=str, default='../data',
                       help='Path to data directory')
    parser.add_argument('--model', type=str, choices=['rf', 'nn', 'both'],
                       default='both', help='Model to tune')
    parser.add_argument('--n-trials', type=int, default=50,
                       help='Number of optimization trials')
    
    args = parser.parse_args()
    
    tuner = HyperparameterTuner(data_dir=args.data_dir)
    
    if args.model in ['rf', 'both']:
        # Load data for Random Forest
        baseline_model = BaselineModel()
        X, y = baseline_model.load_data_from_directory(args.data_dir, balance_data=True)
        if X is not None and y is not None:
            tuner.tune_random_forest(X, y, n_trials=args.n_trials)
    
    if args.model in ['nn', 'both']:
        # Load data for Neural Network
        neural_model = NeuralModel()
        X, y = neural_model.load_data_from_directory(args.data_dir, balance_data=True)
        if X is not None and y is not None:
            tuner.tune_neural_network(X, y, n_trials=min(args.n_trials, 20))  # Limit NN trials
    
    tuner.save_results()