import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from unlearning.Strategy import Strategy

class CertifiedRemoval(Strategy):
    """
    Implements Certified Removal (Guo et al.) adapted for deep architectures
    like ResNet50 by isolating and updating the final classification layer.
    """
    def __init__(self, removal_bound: float, epsilon: float, l2_reg: float = 0.1):
        super().__init__()
        self.removal_bound = removal_bound  # gamma in the paper
        self.epsilon = epsilon              # Privacy budget
        self.l2_reg = l2_reg                # Lambda regularization term

    def _get_features(self, backbone: nn.Module, loader: DataLoader, device: torch.device):
        """Passes data through the frozen ResNet backbone to extract embedding features."""
        backbone.eval()
        all_features = []
        all_labels = []
        
        with torch.no_grad():
            for inputs, labels in loader:
                inputs = inputs.to(device)
                # Pass through backbone to get the 2048-dimensional feature vector
                features = backbone(inputs)
                all_features.append(features.cpu())
                all_labels.append(labels.cpu())
                
        return torch.cat(all_features, dim=0), torch.cat(all_labels, dim=0)

    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
        """
        Entry point expected by your Model.unlearn() architecture interface.
        Applies Certified Removal strictly to the final linear layer (model.fc).
        """
        device = next(model.parameters()).device
        
        # Isolate the final NN (Fully connected) layer from the model
        linear_head = model.fc
        # Temporarily turn the fc layer into a identity pass-through
        model.fc = nn.Identity() 
        
        print(">> Extracting deep features from model backbone...")
        retain_features, retain_labels = self._get_features(model, retain_loader, device)
        forget_features, forget_labels = self._get_features(model, forget_loader, device)
        
        # Restore the linear head back
        model.fc = linear_head
        
        # Extract weights from the classification layer
        # w shape: [num_classes, 2048]
        w = model.fc.weight.data.clone().cpu()
        
        # Compute the Exact Hessian Matrix over the remaining (retained) features
        # Formula: H = (X^T * X) / N + lambda * I
        N_retain = retain_features.size(0)
        hessian = self._compute_hessian(retain_features=retain_features, retain_features_size = N_retain)
        
        grad_forget = self._compute_loss_gradient(
            forget_labels=forget_labels,
            forget_features=forget_features, 
            model_weights=w)
        #torch.matmul(error.t(), forget_features) / forget_features.size(0)

        # Compute the Newton step update via solving: H * Delta_W^T = Grad_forget^T
        delta_w = self._compute_newton_step(
            tensor = hessian,
            gradient= grad_forget
        )
        # Apply the Certified Removal update rule: W_new = W + Delta_W
        new_w = w + delta_w
        # Calibrate noise based on your epsilon budget 
        # (Guo et al. use a perturbation based on the regularization lambda and epsilon)
        sigma = 2.0 / (self.l2_reg * self.epsilon)
        noise = torch.randn_like(new_w) * (sigma / N_retain)
        new_w = new_w + noise

        # Theoretical Guarantee verification
        norm_delta = torch.norm(delta_w).item()
        if norm_delta > self.removal_bound:
            print(f"!! Warning: Removal budget exceeded! Norm: {norm_delta:.4f} > Bound: {self.removal_bound}")
        else:
            print(f">> Certificate valid. Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}")

        # Push updated parameters back into the model instance in-place
        model.fc.weight.data = new_w.to(device)
        
        print(">> Certified Removal process completed successfully.")
        return model
    

    # computing the hessian matrix
    def _compute_hessian(self, retain_features, retain_features_size):
        print(">> Computing exact Hessian matrix...")
        # N_retain = retain_features.size(0)
        X_T_X = torch.matmul(retain_features.t(), retain_features)
        reg_matrix = self.l2_reg * torch.eye(retain_features.size(1))
        return (X_T_X / retain_features_size) + reg_matrix
    

    def _compute_loss_gradient(self, forget_features, forget_labels, model_weights):
        print(">> Calculating forget set gradients...")
        num_classes = model_weights.size(0)
        # Pass features through linear layer weights to get logits
        logits_forget = torch.matmul(forget_features, model_weights.t()) 
        # Apply softmax to get true class probabilities
        preds_softmax = torch.softmax(logits_forget, dim=1)
        
        forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float()
        

        error = preds_softmax - forget_labels_one_hot
        # grad_forget shape: [num_classes, 2048]
        return torch.matmul(error.t(), forget_features) / forget_features.size(0)
    

    def _compute_newton_step(self,tensor, gradient):
        print(">> Solving Newton step via system optimization...")
        try:
            delta_w_t = torch.linalg.solve(tensor, gradient.t())
            delta_w = delta_w_t.t()
        except RuntimeError:
            print(">> Warning: Hessian matrix is singular. Falling back to pseudo-inverse.")
            delta_w = torch.matmul(gradient, torch.linalg.pinv(tensor).t())
        return delta_w