k"""
Final model evaluation script.

This script loads the trained model and computes advanced metrics
on the CIFAR-100 test set as per the project feedback.

Calculates:
- Top-1 Accuracy
- Top-5 Accuracy
- Macro F1-Score
- Macro AUC (One-vs-Rest)
- Classification Report
"""

import torch
import torch.nn.functional as F
import numpy as np
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score, classification_report
from tqdm import tqdm

from src.model import create_resnet50_model
from src.dataset import get_cifar100_data_loaders

MODEL_WEIGHTS_PATH = "cifar100_resnet50_final.pth"
BATCH_SIZE = 128
NUM_WORKERS = 4
NUM_CLASSES = 100

def calculate_top5_accuracy(labels, top5_preds):
    assert labels.ndim == 1, "Labels should be 1D"
    assert top5_preds.ndim == 2 and top5_preds.shape[1] == 5, "Top-5 preds should be (N, 5)"
    
    correct = 0
    for i in range(len(labels)):
        if labels[i] in top5_preds[i]:
            correct += 1
    
    return correct / len(labels)

def evaluate_model():
    
    print("--- Starting Model Evaluation ---")
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    _, test_loader = get_cifar100_data_loaders(
        batch_size=BATCH_SIZE, 
        num_workers=NUM_WORKERS
    )

    print(f"Loading model from {MODEL_WEIGHTS_PATH}...")
    try:
        model = create_resnet50_model(pretrained=False)
        
        model.load_state_dict(torch.load(MODEL_WEIGHTS_PATH, map_location=device))
        model.to(device)
        model.eval()
        print("Model loaded successfully.")
    except FileNotFoundError:
        print(f"Error: Model weights file not found at '{MODEL_WEIGHTS_PATH}'")
        print("Please run train.py first to generate the model file.")
        return
    except Exception as e:
        print(f"An error occurred loading the model: {e}")
        return

    print("Running inference on test set...")
    
    all_labels = []
    all_preds_top1 = []
    all_preds_top5 = []
    all_probabilities = []

    with torch.no_grad():
        for images, labels in tqdm(test_loader, desc="Evaluating"):
            images = images.to(device)
            labels = labels.to(device)
            
            outputs = model(images)
            
            probs = F.softmax(outputs, dim=1)
            all_probabilities.append(probs.cpu().numpy())
            
            _, preds_top1 = torch.max(outputs, 1)
            all_preds_top1.append(preds_top1.cpu().numpy())
            
            _, preds_top5 = torch.topk(outputs, 5, dim=1)
            all_preds_top5.append(preds_top5.cpu().numpy())
            
            all_labels.append(labels.cpu().numpy())

    all_labels = np.concatenate(all_labels)
    all_preds_top1 = np.concatenate(all_preds_top1)
    all_preds_top5 = np.concatenate(all_preds_top5)
    all_probabilities = np.concatenate(all_probabilities)

    print("Inference complete. Calculating metrics...")
    
    top1_acc = accuracy_score(all_labels, all_preds_top1)
    
    top5_acc = calculate_top5_accuracy(all_labels, all_preds_top5)
    
    f1 = f1_score(all_labels, all_preds_top1, average='macro')
    
    auc = roc_auc_score(
        all_labels, 
        all_probabilities, 
        multi_class='ovr', 
        average='macro'
    )
    
    report = classification_report(all_labels, all_preds_top1)
    
    print("\n--- Evaluation Results ---")
    print(f"Top-1 Accuracy: {top1_acc:.4f}")
    print(f"Top-5 Accuracy: {top5_acc:.4f}")
    print(f"Macro F1-Score: {f1:.4f}")
    print(f"Macro AUC (ovr): {auc:.4f}")
    print("\n--- Classification Report ---")
    print(report)
    print("--------------------------")

if __name__ == "__main__":
    evaluate_model()