evaluate_models.py

#!/usr/bin/env python3
"""
Model Evaluation Script for Metal Detector AI

Provides comprehensive evaluation metrics including:
- Cross-validation scores
- Confusion matrix
- Precision/Recall/F1 scores
- ROC curves
- Feature importance analysis
"""

import argparse
import numpy as np
from pathlib import Path
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import (
    confusion_matrix, classification_report, roc_curve, auc,
    precision_recall_curve, average_precision_score
)
from sklearn.model_selection import cross_val_score, StratifiedKFold
from sklearn.preprocessing import label_binarize
import json
import pandas as pd
from datetime import datetime
from rich.console import Console
from rich.table import Table
from rich.panel import Panel

from src.ml.advanced_classifier import AdvancedMetalClassifier
from src.ml.classifier import MetalClassifier

console = Console()

class ModelEvaluator:
    """Comprehensive model evaluation tools."""
    
    def __init__(self, model_path: Path = Path("models/advanced")):
        self.model_path = model_path
        self.results_dir = Path("evaluation_results")
        self.results_dir.mkdir(exist_ok=True)
        
        # Create subdirectories
        self.plots_dir = self.results_dir / "plots"
        self.plots_dir.mkdir(exist_ok=True)
        
        self.timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    def evaluate_advanced_model(self, data_dir: Path):
        """
        Evaluates ensemble model with metrics, plots, and feature analysis [ref: sklearn.metrics].
        """
        console.print(Panel.fit("🔬 Advanced Model Evaluation", style="bold blue"))
        
        # Load classifier
        classifier = AdvancedMetalClassifier(model_path=self.model_path)
        
        if not classifier._model_exists():
            console.print("❌ No trained models found! Train models first.")
            return
        
        classifier.load_model()
        console.print("✅ Models loaded successfully!")
        
        # Prepare test data
        console.print("\n📊 Preparing evaluation data...")
        audio_data, labels, features, _ = classifier.prepare_training_data(data_dir)
        
        if len(audio_data) == 0:
            console.print("❌ No evaluation data found!")
            return
        
        # Encode labels
        y_true = classifier.label_encoder.transform(labels)
        classes = classifier.label_encoder.classes_
        n_classes = len(classes)
        
        console.print(f"Found {len(audio_data)} samples across {n_classes} classes")
        
        # Get predictions
        console.print("\n🔮 Getting predictions...")
        y_pred = []
        y_proba = []
        
        with console.status("Classifying samples..."):
            for i, audio in enumerate(audio_data):
                # Get predictions from each model
                pred_cnn, prob_cnn = classifier._predict_cnn(audio)
                pred_trans, prob_trans = classifier._predict_transformer(audio)
                pred_trad, prob_trad = classifier._predict_traditional(audio)
                
                # Ensemble prediction (same as in classify_audio_file)
                ensemble_prob = (prob_cnn * 0.4 + prob_trans * 0.4 + prob_trad * 0.2)
                ensemble_pred = np.argmax(ensemble_prob)
                
                y_pred.append(ensemble_pred)
                y_proba.append(ensemble_prob)
                
                if (i + 1) % 50 == 0:
                    console.print(f"Processed {i + 1}/{len(audio_data)} samples")
        
        y_pred = np.array(y_pred)
        y_proba = np.array(y_proba)
        
        # Generate evaluation metrics
        self._generate_classification_report(y_true, y_pred, classes, "advanced")
        self._generate_confusion_matrix(y_true, y_pred, classes, "advanced")
        self._generate_roc_curves(y_true, y_proba, classes, "advanced")
        self._analyze_feature_importance(classifier, features, labels, "advanced")
        
        # Save results summary
        results = {
            "timestamp": self.timestamp,
            "model_type": "advanced_ensemble",
            "n_samples": len(y_true),
            "n_classes": n_classes,
            "classes": classes.tolist(),
            "accuracy": float((y_pred == y_true).mean()),
            "results_dir": str(self.results_dir)
        }
        
        results_file = self.results_dir / f"evaluation_summary_{self.timestamp}.json"
        with open(results_file, 'w') as f:
            json.dump(results, f, indent=2)
        
        console.print(f"\n✅ Evaluation complete! Results saved to: {self.results_dir}")
    
    def evaluate_baseline_model(self, data_dir: Path):
        """
        Evaluates traditional ML model including cross-validation [ref: sklearn.model_selection].
        """
        console.print(Panel.fit("🔬 Baseline Model Evaluation", style="bold yellow"))
        
        # Load classifier
        classifier = MetalClassifier(model_path=Path("models/basic"))
        
        if not classifier.model_exists():
            console.print("❌ No trained baseline model found! Train model first.")
            return
        
        classifier.load_model()
        console.print("✅ Baseline model loaded successfully!")
        
        # Prepare test data
        console.print("\n📊 Preparing evaluation data...")
        X, y, label_names = classifier.prepare_training_data(data_dir)
        
        if len(X) == 0:
            console.print("❌ No evaluation data found!")
            return
        
        # Scale features
        X_scaled = classifier.scaler.transform(X)
        y_encoded = classifier.label_encoder.transform(y)
        classes = classifier.label_encoder.classes_
        
        console.print(f"Found {len(X)} samples across {len(classes)} classes")
        
        # Get predictions
        y_pred = classifier.model.predict(X_scaled)
        y_proba = classifier.model.predict_proba(X_scaled)
        
        # Generate evaluation metrics
        self._generate_classification_report(y_encoded, y_pred, classes, "baseline")
        self._generate_confusion_matrix(y_encoded, y_pred, classes, "baseline")
        self._generate_roc_curves(y_encoded, y_proba, classes, "baseline")
        self._analyze_baseline_feature_importance(classifier, X, y, "baseline")
        
        # Cross-validation
        self._perform_cross_validation(classifier.model, X_scaled, y_encoded, "baseline")
        
        console.print(f"\n✅ Baseline evaluation complete! Results saved to: {self.results_dir}")
    
    def _generate_classification_report(self, y_true, y_pred, classes, model_name):
        """
        Computes and displays classification metrics, saves JSON report.
        """
        console.print("\n📋 Classification Report:")
        
        report = classification_report(y_true, y_pred, target_names=classes, output_dict=True)
        
        # Display in console
        table = Table(title="Per-Class Metrics")
        table.add_column("Class", style="cyan")
        table.add_column("Precision", style="green")
        table.add_column("Recall", style="yellow")
        table.add_column("F1-Score", style="magenta")
        table.add_column("Support", style="blue")
        
        for class_name in classes:
            metrics = report[class_name]
            table.add_row(
                class_name,
                f"{metrics['precision']:.3f}",
                f"{metrics['recall']:.3f}",
                f"{metrics['f1-score']:.3f}",
                str(int(metrics['support']))
            )
        
        console.print(table)
        
        # Overall metrics
        console.print(f"\n🎯 Overall Accuracy: {report['accuracy']:.3f}")
        console.print(f"📊 Macro Avg F1: {report['macro avg']['f1-score']:.3f}")
        console.print(f"📊 Weighted Avg F1: {report['weighted avg']['f1-score']:.3f}")
        
        # Save report
        report_file = self.results_dir / f"{model_name}_classification_report_{self.timestamp}.json"
        with open(report_file, 'w') as f:
            json.dump(report, f, indent=2)
    
    def _generate_confusion_matrix(self, y_true, y_pred, classes, model_name):
        """
        Creates and saves confusion matrix heatmap using seaborn.
        """
        console.print("\n🔲 Generating Confusion Matrix...")
        
        cm = confusion_matrix(y_true, y_pred)
        
        plt.figure(figsize=(10, 8))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                    xticklabels=classes, yticklabels=classes)
        plt.title(f'Confusion Matrix - {model_name.title()} Model')
        plt.xlabel('Predicted Label')
        plt.ylabel('True Label')
        plt.tight_layout()
        
        plot_file = self.plots_dir / f"{model_name}_confusion_matrix_{self.timestamp}.png"
        plt.savefig(plot_file, dpi=300, bbox_inches='tight')
        plt.close()
        
        console.print(f"💾 Confusion matrix saved to: {plot_file}")
    
    def _generate_roc_curves(self, y_true, y_proba, classes, model_name):
        """
        Plots multi-class ROC curves and saves figure.
        """
        console.print("\n📈 Generating ROC Curves...")
        
        # Binarize labels for multi-class ROC
        y_true_bin = label_binarize(y_true, classes=range(len(classes)))
        
        # Compute ROC curve for each class
        plt.figure(figsize=(10, 8))
        
        for i, class_name in enumerate(classes):
            fpr, tpr, _ = roc_curve(y_true_bin[:, i], y_proba[:, i])
            roc_auc = auc(fpr, tpr)
            plt.plot(fpr, tpr, label=f'{class_name} (AUC = {roc_auc:.3f})')
        
        plt.plot([0, 1], [0, 1], 'k--', label='Random Classifier')
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title(f'ROC Curves - {model_name.title()} Model')
        plt.legend(loc="lower right")
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        
        plot_file = self.plots_dir / f"{model_name}_roc_curves_{self.timestamp}.png"
        plt.savefig(plot_file, dpi=300, bbox_inches='tight')
        plt.close()
        
        console.print(f"💾 ROC curves saved to: {plot_file}")
    
    def _analyze_feature_importance(self, classifier, features, labels, model_name):
        """
        Extracts and plots top feature importances from RandomForest component.
        """
        console.print("\n🔍 Analyzing Feature Importance...")
        
        # For traditional ensemble component
        if hasattr(classifier, 'traditional_ensemble'):
            # Get feature importance from Random Forest in the ensemble
            rf_model = classifier.traditional_ensemble.estimators_[0]
            if hasattr(rf_model, 'feature_importances_'):
                importances = rf_model.feature_importances_
                
                # Create feature names (you might want to update this based on actual features)
                n_features = len(importances)
                feature_names = [f'Feature_{i}' for i in range(n_features)]
                
                # Sort features by importance
                indices = np.argsort(importances)[::-1][:20]  # Top 20 features
                
                plt.figure(figsize=(12, 6))
                plt.bar(range(len(indices)), importances[indices])
                plt.xticks(range(len(indices)), [feature_names[i] for i in indices], rotation=45, ha='right')
                plt.title(f'Top 20 Feature Importances - {model_name.title()} Model (Traditional Component)')
                plt.xlabel('Features')
                plt.ylabel('Importance')
                plt.tight_layout()
                
                plot_file = self.plots_dir / f"{model_name}_feature_importance_{self.timestamp}.png"
                plt.savefig(plot_file, dpi=300, bbox_inches='tight')
                plt.close()
                
                console.print(f"💾 Feature importance plot saved to: {plot_file}")
    
    def _analyze_baseline_feature_importance(self, classifier, X, y, model_name):
        """
        Plots top feature importances for baseline model.
        """
        if hasattr(classifier.model, 'feature_importances_'):
            importances = classifier.model.feature_importances_
            feature_names = classifier.feature_names
            
            # Sort features by importance
            indices = np.argsort(importances)[::-1][:20]  # Top 20 features
            
            plt.figure(figsize=(12, 6))
            plt.bar(range(len(indices)), importances[indices])
            plt.xticks(range(len(indices)), [feature_names[i] for i in indices], rotation=45, ha='right')
            plt.title(f'Top 20 Feature Importances - {model_name.title()} Model')
            plt.xlabel('Features')
            plt.ylabel('Importance')
            plt.tight_layout()
            
            plot_file = self.plots_dir / f"{model_name}_feature_importance_{self.timestamp}.png"
            plt.savefig(plot_file, dpi=300, bbox_inches='tight')
            plt.close()
            
            console.print(f"💾 Feature importance plot saved to: {plot_file}")
    
    def _perform_cross_validation(self, model, X, y, model_name):
        """
        Runs stratified k-fold CV with multiple metrics, saves results.
        """
        console.print("\n🔄 Performing Cross-Validation...")
        
        cv_folds = min(5, len(np.unique(y)))
        if cv_folds < 2:
            console.print("⚠️ Not enough samples for cross-validation")
            return
        
        skf = StratifiedKFold(n_splits=cv_folds, shuffle=True, random_state=42)
        
        # Different scoring metrics
        scoring_metrics = ['accuracy', 'f1_macro', 'precision_macro', 'recall_macro']
        
        cv_results = {}
        for metric in scoring_metrics:
            scores = cross_val_score(model, X, y, cv=skf, scoring=metric)
            cv_results[metric] = {
                'mean': float(scores.mean()),
                'std': float(scores.std()),
                'scores': scores.tolist()
            }
            console.print(f"{metric}: {scores.mean():.3f} (+/- {scores.std() * 2:.3f})")
        
        # Save CV results
        cv_file = self.results_dir / f"{model_name}_cv_results_{self.timestamp}.json"
        with open(cv_file, 'w') as f:
            json.dump(cv_results, f, indent=2)

def main():
    """Main evaluation function."""
    parser = argparse.ArgumentParser(description="Evaluate Metal Detector AI models")
    parser.add_argument("--data-dir", type=str, default="data",
                       help="Directory containing evaluation data")
    parser.add_argument("--model", type=str, choices=['advanced', 'baseline', 'all'], 
                       default='advanced', help="Which model to evaluate")
    parser.add_argument("--model-dir", type=str, help="Custom model directory")
    
    args = parser.parse_args()
    
    data_dir = Path(args.data_dir)
    if not data_dir.exists():
        console.print(f"❌ Data directory not found: {data_dir}")
        return
    
    # Initialize evaluator
    model_path = Path(args.model_dir) if args.model_dir else Path("models/advanced")
    evaluator = ModelEvaluator(model_path)
    
    # Run evaluation
    if args.model in ['advanced', 'all']:
        evaluator.evaluate_advanced_model(data_dir)
    
    if args.model in ['baseline', 'all']:
        evaluator.evaluate_baseline_model(data_dir)
    
    console.print("\n🎉 Evaluation complete!")

if __name__ == "__main__":
    main()