evaluate.py

"""
Evaluation Script for Anomaly Detection Model.

Evaluates model performance on test set and calculates metrics.
"""

import argparse
from pathlib import Path

import numpy as np
import torch
from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    roc_auc_score,
    confusion_matrix,
)
from loguru import logger

from src.data.data_loader import create_dataloaders
from src.models.autoencoder import VibrationAutoencoder
from src.utils.config_loader import load_config


def calculate_reconstruction_errors(
    model: torch.nn.Module,
    dataloader: torch.utils.data.DataLoader,
    device: str = "cuda",
) -> tuple:
    """Calculate reconstruction errors for all samples."""
    model.eval()
    model.to(device)
    
    all_errors = []
    all_labels = []
    
    with torch.no_grad():
        for batch in dataloader:
            if isinstance(batch, (list, tuple)) and len(batch) == 2:
                data, labels = batch
                all_labels.extend(labels.cpu().numpy())
            else:
                data = batch
            
            data = data.to(device)
            
            # Forward pass
            reconstruction = model(data)
            
            # Calculate reconstruction error (MSE per sample)
            errors = torch.mean((data - reconstruction) ** 2, dim=1)
            all_errors.extend(errors.cpu().numpy())
    
    return np.array(all_errors), np.array(all_labels) if all_labels else None


def find_optimal_threshold(
    errors: np.ndarray,
    labels: np.ndarray,
    method: str = "f1",
) -> float:
    """Find optimal threshold based on labels."""
    thresholds = np.percentile(errors, np.arange(0, 100, 1))
    
    best_threshold = 0
    best_score = 0
    
    for threshold in thresholds:
        predictions = (errors > threshold).astype(int)
        
        if method == "f1":
            score = f1_score(labels, predictions)
        elif method == "accuracy":
            score = accuracy_score(labels, predictions)
        else:
            score = f1_score(labels, predictions)
        
        if score > best_score:
            best_score = score
            best_threshold = threshold
    
    return best_threshold


def calculate_threshold_statistical(
    errors: np.ndarray,
    sigma_multiplier: float = 3.0,
) -> float:
    """Calculate threshold using statistical method (μ + k*σ)."""
    mean = np.mean(errors)
    std = np.std(errors)
    threshold = mean + sigma_multiplier * std
    return threshold


def evaluate_model(
    model: torch.nn.Module,
    dataloader: torch.utils.data.DataLoader,
    threshold: float,
    device: str = "cuda",
) -> dict:
    """Evaluate model and calculate metrics."""
    logger.info("Calculating reconstruction errors...")
    errors, labels = calculate_reconstruction_errors(model, dataloader, device)
    
    # Predictions
    predictions = (errors > threshold).astype(int)
    
    metrics = {
        "num_samples": len(errors),
        "threshold": threshold,
        "mean_error": float(np.mean(errors)),
        "std_error": float(np.std(errors)),
        "min_error": float(np.min(errors)),
        "max_error": float(np.max(errors)),
    }
    
    if labels is not None:
        # Classification metrics
        metrics.update({
            "accuracy": accuracy_score(labels, predictions),
            "precision": precision_score(labels, predictions, zero_division=0),
            "recall": recall_score(labels, predictions, zero_division=0),
            "f1_score": f1_score(labels, predictions, zero_division=0),
        })
        
        # Confusion matrix
        cm = confusion_matrix(labels, predictions)
        tn, fp, fn, tp = cm.ravel() if cm.size == 4 else (0, 0, 0, 0)
        
        metrics.update({
            "true_negatives": int(tn),
            "false_positives": int(fp),
            "false_negatives": int(fn),
            "true_positives": int(tp),
            "specificity": tn / (tn + fp) if (tn + fp) > 0 else 0,
            "false_positive_rate": fp / (fp + tn) if (fp + tn) > 0 else 0,
            "false_negative_rate": fn / (fn + tp) if (fn + tp) > 0 else 0,
        })
        
        # ROC AUC
        try:
            metrics["roc_auc"] = roc_auc_score(labels, errors)
        except:
            metrics["roc_auc"] = None
    
    return metrics


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description="Evaluate anomaly detection model")
    
    parser.add_argument(
        "--checkpoint",
        type=str,
        required=True,
        help="Path to model checkpoint",
    )
    
    parser.add_argument(
        "--config",
        type=str,
        default="config/config.yaml",
        help="Path to configuration file",
    )
    
    parser.add_argument(
        "--threshold",
        type=float,
        default=None,
        help="Anomaly threshold (if None, calculates from data)",
    )
    
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        choices=["cuda", "cpu"],
        help="Device to use",
    )
    
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output file for metrics (.json)",
    )
    
    return parser.parse_args()


def main():
    """Main evaluation function."""
    args = parse_args()
    
    # Load configuration
    logger.info(f"Loading configuration from {args.config}")
    config = load_config(args.config)
    
    # Device setup
    device = args.device
    if device == "cuda" and not torch.cuda.is_available():
        logger.warning("CUDA not available, using CPU")
        device = "cpu"
    
    # Load model
    logger.info(f"Loading model from {args.checkpoint}")
    checkpoint = torch.load(args.checkpoint, map_location=device)
    
    if "model" in checkpoint:
        model = checkpoint["model"]
    else:
        # Load architecture
        model = VibrationAutoencoder()
        model.load_state_dict(checkpoint["model_state_dict"])
    
    model.eval()
    
    # Load data
    logger.info("Loading test data...")
    _, _, test_loader, preprocessor = create_dataloaders(config)
    
    # Calculate or use provided threshold
    if args.threshold is None:
        logger.info("Calculating optimal threshold...")
        errors, labels = calculate_reconstruction_errors(model, test_loader, device)
        
        if labels is not None:
            threshold = find_optimal_threshold(errors, labels, method="f1")
            logger.info(f"Optimal threshold (F1): {threshold:.6f}")
        else:
            threshold = calculate_threshold_statistical(errors, sigma_multiplier=3.0)
            logger.info(f"Statistical threshold (μ + 3σ): {threshold:.6f}")
    else:
        threshold = args.threshold
        logger.info(f"Using provided threshold: {threshold:.6f}")
    
    # Evaluate
    logger.info("Evaluating model...")
    metrics = evaluate_model(model, test_loader, threshold, device)
    
    # Print results
    logger.info("\n" + "=" * 50)
    logger.info("EVALUATION RESULTS")
    logger.info("=" * 50)
    
    logger.info(f"\nDataset Statistics:")
    logger.info(f"  Samples: {metrics['num_samples']}")
    logger.info(f"  Mean Error: {metrics['mean_error']:.6f}")
    logger.info(f"  Std Error: {metrics['std_error']:.6f}")
    logger.info(f"  Min Error: {metrics['min_error']:.6f}")
    logger.info(f"  Max Error: {metrics['max_error']:.6f}")
    
    if "accuracy" in metrics:
        logger.info(f"\nClassification Metrics:")
        logger.info(f"  Accuracy: {metrics['accuracy']:.4f}")
        logger.info(f"  Precision: {metrics['precision']:.4f}")
        logger.info(f"  Recall: {metrics['recall']:.4f}")
        logger.info(f"  F1-Score: {metrics['f1_score']:.4f}")
        logger.info(f"  Specificity: {metrics['specificity']:.4f}")
        
        if metrics["roc_auc"] is not None:
            logger.info(f"  ROC-AUC: {metrics['roc_auc']:.4f}")
        
        logger.info(f"\nConfusion Matrix:")
        logger.info(f"  True Negatives: {metrics['true_negatives']}")
        logger.info(f"  False Positives: {metrics['false_positives']}")
        logger.info(f"  False Negatives: {metrics['false_negatives']}")
        logger.info(f"  True Positives: {metrics['true_positives']}")
        logger.info(f"  FPR: {metrics['false_positive_rate']:.4f}")
        logger.info(f"  FNR: {metrics['false_negative_rate']:.4f}")
    
    logger.info("=" * 50 + "\n")
    
    # Save metrics
    if args.output:
        import json
        with open(args.output, "w") as f:
            json.dump(metrics, f, indent=2)
        logger.info(f"Metrics saved to {args.output}")


if __name__ == "__main__":
    main()