"""
Evaluation Metrics - PyTorch Neural Networks
Project: PyTorch Neural Networks
Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
Description: Evaluation metrics and visualization utilities
"""

import torch
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, classification_report, accuracy_score
import seaborn as sns


def calculate_accuracy(y_true, y_pred):
    """
    Calculate accuracy
    
    Args:
        y_true: True labels
        y_pred: Predicted labels
        
    Returns:
        Accuracy score
    """
    if isinstance(y_true, torch.Tensor):
        y_true = y_true.cpu().numpy()
    if isinstance(y_pred, torch.Tensor):
        y_pred = y_pred.cpu().numpy()
    
    return accuracy_score(y_true, y_pred)


def calculate_confusion_matrix(y_true, y_pred, class_names=None):
    """
    Calculate and return confusion matrix
    
    Args:
        y_true: True labels
        y_pred: Predicted labels
        class_names: Optional class names for visualization
        
    Returns:
        Confusion matrix
    """
    if isinstance(y_true, torch.Tensor):
        y_true = y_true.cpu().numpy()
    if isinstance(y_pred, torch.Tensor):
        y_pred = y_pred.cpu().numpy()
    
    cm = confusion_matrix(y_true, y_pred)
    return cm


def plot_confusion_matrix(y_true, y_pred, class_names=None, save_path=None):
    """
    Plot confusion matrix
    
    Args:
        y_true: True labels
        y_pred: Predicted labels
        class_names: Optional class names
        save_path: Path to save the plot
    """
    cm = calculate_confusion_matrix(y_true, y_pred, class_names)
    
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=class_names, yticklabels=class_names)
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.title('Confusion Matrix')
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Confusion matrix saved to {save_path}")
    
    plt.show()


def classification_report_metrics(y_true, y_pred, class_names=None):
    """
    Print classification report
    
    Args:
        y_true: True labels
        y_pred: Predicted labels
        class_names: Optional class names
    """
    if isinstance(y_true, torch.Tensor):
        y_true = y_true.cpu().numpy()
    if isinstance(y_pred, torch.Tensor):
        y_pred = y_pred.cpu().numpy()
    
    report = classification_report(y_true, y_pred, target_names=class_names)
    print(report)
    return report


def evaluate_model(model, data_loader, device, criterion=None):
    """
    Evaluate model on a dataset
    
    Args:
        model: Model to evaluate
        data_loader: DataLoader for the dataset
        device: Device to run on
        criterion: Optional loss function
        
    Returns:
        Dictionary with evaluation metrics
    """
    model.eval()
    all_preds = []
    all_labels = []
    total_loss = 0.0
    
    with torch.no_grad():
        for inputs, labels in data_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            
            outputs = model(inputs)
            
            if criterion:
                loss = criterion(outputs, labels)
                total_loss += loss.item()
            
            _, predicted = torch.max(outputs.data, 1)
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    results = {
        'predictions': np.array(all_preds),
        'labels': np.array(all_labels),
        'accuracy': calculate_accuracy(all_labels, all_preds)
    }
    
    if criterion:
        results['loss'] = total_loss / len(data_loader)
    
    return results