import torch
import numpy as np

@torch.inference_mode() # More memory-efficient than no_grad()
def get_loss_per_sample(model, data_loader, device):
    """
    Returns a list of individual losses for every sample in the loader.
    Useful for MIA to see how 'certain' the model is about specific images.
    """
    model.eval()
    criterion = torch.nn.CrossEntropyLoss(reduction='none') # Crucial: returns loss per image
    all_losses = []

    for inputs, labels in data_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        
        outputs = model(inputs)
        
        # Calculate loss for each image in the batch individually
        loss = criterion(outputs, labels)
        
        all_losses.extend(loss.cpu().numpy())

    return all_losses


@torch.inference_mode()
def get_losses_by_class(model, data_loader, device):
    """
    Returns a dictionary: { class_id: [list_of_losses_for_this_class] }
    """
    model.eval()
    criterion = torch.nn.CrossEntropyLoss(reduction='none')
    
    class_losses = {}

    for inputs, labels in data_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        
        # Get individual losses
        losses = criterion(outputs, labels).cpu().numpy()
        labels_np = labels.cpu().numpy()

        for i, class_id in enumerate(labels_np):
            if class_id not in class_losses:
                class_losses[class_id] = []
            class_losses[class_id].append(losses[i])

    return class_losses


# evaluate MIA
def eval_MIA(forgotten_losses, never_seen_losses):
    avg_f_loss = np.mean(forgotten_losses)
    avg_ns_loss = np.mean(never_seen_losses)

    print(f"Average Loss on Forgotten Identity: {avg_f_loss:.4f}")
    print(f"Average Loss on Unknown Identities: {avg_ns_loss:.4f}")

    if avg_f_loss < avg_ns_loss * 0.8:
        print("MIA Warning: Model still shows high certainty on forgotten data.")
    else:
        print("MIA Success: Model treats forgotten data like unknown data.")