utils.py

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Dict, List, Tuple, Any
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score,
    confusion_matrix, classification_report
)
import joblib


def compute_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, float]:
    return {
        'accuracy': accuracy_score(y_true, y_pred),
        'precision': precision_score(y_true, y_pred),
        'recall': recall_score(y_true, y_pred),
        'f1': f1_score(y_true, y_pred)
    }


def get_classification_report(y_true: np.ndarray, y_pred: np.ndarray,
                               target_names: List[str] = None) -> str:
    if target_names is None:
        target_names = ['Healthy', 'Parkinson']
    return classification_report(y_true, y_pred, target_names=target_names)


def plot_confusion_matrix(y_true: np.ndarray, y_pred: np.ndarray,
                          model_name: str, save_path: str,
                          target_names: List[str] = None) -> None:
    if target_names is None:
        target_names = ['Healthy', 'Parkinson']

    cm = confusion_matrix(y_true, y_pred)

    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=target_names, yticklabels=target_names,
                annot_kws={'size': 12})
    plt.title(f'Confusion Matrix - {model_name}', fontsize=14, fontweight='bold')
    plt.xlabel('Predicted Label', fontsize=12)
    plt.ylabel('True Label', fontsize=12)
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"  ✓ Confusion matrix saved: {save_path}")


def plot_accuracy_comparison(results: Dict[str, Dict[str, float]],
                             save_path: str, pca_suffix: str = "") -> None:
    model_names = list(results.keys())
    accuracies = [results[model]['accuracy'] for model in model_names]

    display_names = [name.replace('_no_pca', '').replace('_pca', '')
                     for name in model_names]

    plt.figure(figsize=(10, 6))
    bars = plt.bar(display_names, accuracies, color=['#3498db', '#2ecc71', '#e74c3c'])

    for bar, acc in zip(bars, accuracies):
        plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
                f'{acc:.4f}', ha='center', va='bottom', fontsize=11, fontweight='bold')

    plt.title(f'Model Accuracy Comparison {pca_suffix}'.strip(), fontsize=14, fontweight='bold')
    plt.xlabel('Model', fontsize=12)
    plt.ylabel('Accuracy', fontsize=12)
    plt.ylim(0, 1.1)
    plt.xticks(fontsize=10)
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"  ✓ Accuracy comparison saved: {save_path}")


def plot_pca_variance(pca, save_path: str) -> None:
    explained_variance = pca.explained_variance_ratio_
    cumulative_variance = np.cumsum(explained_variance)

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    ax1.bar(range(1, len(explained_variance) + 1), explained_variance,
            alpha=0.7, color='#3498db', edgecolor='black')
    ax1.set_xlabel('Principal Component', fontsize=12)
    ax1.set_ylabel('Explained Variance Ratio', fontsize=12)
    ax1.set_title('Individual Explained Variance', fontsize=13, fontweight='bold')
    ax1.grid(axis='y', alpha=0.3)

    ax2.plot(range(1, len(cumulative_variance) + 1), cumulative_variance,
             marker='o', markersize=8, linewidth=2, color='#e74c3c')
    ax2.axhline(y=0.95, color='green', linestyle='--', linewidth=2,
                label='95% Variance Threshold')
    ax2.axhline(y=0.90, color='orange', linestyle='--', linewidth=2,
                label='90% Variance Threshold')
    ax2.set_xlabel('Number of Principal Components', fontsize=12)
    ax2.set_ylabel('Cumulative Explained Variance', fontsize=12)
    ax2.set_title('Cumulative Explained Variance', fontsize=13, fontweight='bold')
    ax2.legend(fontsize=10)
    ax2.grid(alpha=0.3)

    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"  ✓ PCA variance plot saved: {save_path}")

    n_components_90 = np.argmax(cumulative_variance >= 0.90) + 1
    n_components_95 = np.argmax(cumulative_variance >= 0.95) + 1
    print(f"\n📊 PCA Variance Summary:")
    print(f"  • Components for 90% variance: {n_components_90}")
    print(f"  • Components for 95% variance: {n_components_95}")
    print(f"  • Total variance retained: {cumulative_variance[-1]:.4f}")


def plot_feature_importance(feature_names: List[str], importances: np.ndarray,
                            model_name: str, save_path: str,
                            top_n: int = 15) -> None:
    indices = np.argsort(importances)[::-1][:top_n]

    top_features = [feature_names[i] for i in indices]
    top_importances = importances[indices]

    plt.figure(figsize=(12, 8))
    colors = plt.cm.viridis(np.linspace(0.2, 0.8, len(top_features)))
    bars = plt.barh(range(len(top_features)), top_importances, color=colors)

    plt.yticks(range(len(top_features)), top_features, fontsize=9)
    plt.xlabel('Importance', fontsize=12)
    plt.title(f'Top {top_n} Feature Importances - {model_name}',
              fontsize=14, fontweight='bold')
    plt.gca().invert_yaxis()
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"  ✓ Feature importance plot saved: {save_path}")


def save_model(model: Any, filepath: str) -> None:
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    joblib.dump(model, filepath)
    print(f"  ✓ Model saved: {filepath}")


def load_model(filepath: str) -> Any:
    if not os.path.exists(filepath):
        raise FileNotFoundError(f"Model file not found: {filepath}")
    return joblib.load(filepath)


def save_scaler(scaler: Any, filepath: str) -> None:
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    joblib.dump(scaler, filepath)
    print(f"  ✓ Scaler saved: {filepath}")


def save_pca(pca: Any, filepath: str) -> None:
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    joblib.dump(pca, filepath)
    print(f"  ✓ PCA saved: {filepath}")


def save_results_to_csv(results: List[Dict[str, Any]], filepath: str) -> None:
    df = pd.DataFrame(results)

    column_order = ['model_name', 'pca', 'accuracy', 'precision', 'recall', 'f1']
    existing_columns = [col for col in column_order if col in df.columns]
    df = df[existing_columns]

    df.to_csv(filepath, index=False)
    print(f"\n📊 Results saved to: {filepath}")
    print(df.to_string(index=False))


def save_best_model(model_name: str, filepath: str) -> None:
    with open(filepath, 'w') as f:
        f.write(model_name)
    print(f"\n🏆 Best model saved: {filepath}")
    print(f"   Best Model: {model_name}")


def get_best_model(results: List[Dict[str, Any]], metric: str = 'accuracy') -> Tuple[str, float]:
    best = max(results, key=lambda x: x[metric])
    return best['model_name'], best[metric]


def ensure_directory(filepath: str) -> None:
    directory = os.path.dirname(filepath)
    if directory:
        os.makedirs(directory, exist_ok=True)


def get_wav_files(directory: str) -> List[str]:
    wav_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.lower().endswith('.wav'):
                wav_files.append(os.path.join(root, file))
    return sorted(wav_files)