import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import torchvision.models as models

class ZeroRetrainForgettingEvaluator:
    def __init__(self, unlearned_model: nn.Module, num_classes: int):
        """
        Initializes the ZRF Evaluator.
        
        Args:
            unlearned_model (nn.Module): Your fine-tuned & unlearned ResNet-50.
            num_classes (int): Number of classes used in your CelebA task.
        """
        # select device
        if torch.cuda.is_available():
            self.device = torch.device("cuda")
        elif hasattr(torch, "xpu") and torch.xpu.is_available():
            self.device = torch.device("xpu") # For Intel GPUs using IPEX
        else:
            self.device = torch.device("cpu")
            
        print(f"[INFO] Using device: {self.device}")

        # prepare the unlearned model
        self.unlearned_model = unlearned_model.to(self.device)
        self.unlearned_model.eval()

        # Instantiate a structurally matching, completely random model
        print(f"[INFO] Initializing random baseline ResNet-50 with {num_classes} classes...")
        self.random_model = self.get_random_model(num_classes)
        self.random_model = self.random_model.to(self.device)
        self.random_model.eval()

   # gets randomly initialised model 
   # for comparison with unlearned model
    def get_random_model(num_classes):
        print(f"[INFO] Initializing random baseline ResNet-50 with {num_classes} classes...")
        model = models.resnet50(weights=None)
        model.fc = nn.Linear(model.fc.in_features, num_classes)
        return model


    # compute divergence
    def _compute_js_divergence(self, p: torch.Tensor, q: torch.Tensor) -> float:
        """
        Computes the Jensen-Shannon (JS) Divergence between two probability distributions.
        
        Args:
            p, q (Tensor): Tensors of shape (batch_size, num_classes) containing probabilities.
        """
        # Avoid log(0) issues by adding a tiny epsilon
        eps = 1e-12
        p = torch.clamp(p, eps, 1.0)
        q = torch.clamp(q, eps, 1.0)
        
        # Calculate the midpoint distribution
        m = 0.5 * (p + q)
        
        # Compute KL Divergence natively: KL(P || M) and KL(Q || M)
        kl_pm = torch.sum(p * (torch.log(p) - torch.log(m)), dim=1)
        kl_qm = torch.sum(q * (torch.log(q) - torch.log(m)), dim=1)
        
        # JS Divergence is the average of both KL divergences
        js_div = 0.5 * (kl_pm + kl_qm)
        
        # Return the mean divergence across the entire batch
        return js_div.mean().item()

    def evaluate_forget_class(self, dataset, batch_size: int = 32) -> float:
        """
        Evaluates the unlearned model against the random model using images 
        from the forgotten class/identity.
        
        Args:
            dataset (Dataset): A PyTorch Dataset containing images of the forget set.
            batch_size (int): Batch size for evaluation.
            
        Returns:
            float: The ZRF score (JS Divergence). A lower divergence means 
                   the unlearned model is behaving exactly like a random model.
        """
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
        
        total_js_div = 0.0
        total_samples = 0

        # No gradients needed for evaluation
        with torch.no_grad():
            for images, _ in dataloader:
                images = images.to(self.device)
                batch_len = images.size(0)

                # Get raw outputs (logits)
                unlearned_logits = self.unlearned_model(images)
                random_logits = self.random_model(images)

                # Convert logits to probability distributions via Softmax
                unlearned_probs = F.softmax(unlearned_logits, dim=1)
                random_probs = F.softmax(random_logits, dim=1)

                # Calculate JS divergence for this batch
                batch_js = self._compute_js_divergence(unlearned_probs, random_probs)
                
                # Weighted average based on batch size (handles final smaller batches perfectly)
                total_js_div += batch_js * batch_len
                total_samples += batch_len

        final_zrf_score = total_js_div / total_samples
        return final_zrf_score