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

class LastKCertifiedRemoval(Strategy):
    """
    Implements Certified Removal (Guo et al.) scaled up to the last K layers 
    of a ResNet50 network by flattening sub-graph parameters into a convex sub-problem.
    """
    def __init__(self, removal_bound: float, epsilon: float, l2_reg: float = 0.1):
        super().__init__()
        self.removal_bound = removal_bound
        self.epsilon = epsilon
        self.l2_reg = l2_reg

    def _split_model(self, model: nn.Module):
        """
        Splits ResNet50 into a frozen feature backbone and an active unlearning head.
        Here, 'Last K Layers' includes layer4 and the fc classification head.
        """
        # Feature Backbone: Everything up to layer3
        backbone = nn.Sequential(
            model.conv1,
            model.bn1,
            model.relu,
            model.maxpool,
            model.layer1,
            model.layer2,
            model.layer3
        )
        
        # Active Head: Layer4, global pooling, and the final linear layer
        unlearning_head = nn.Sequential(
            model.layer4,
            model.avgpool,
            nn.Flatten(1),
            model.fc
        )
        
        return backbone, unlearning_head

    def _get_intermediate_features(self, backbone: nn.Module, loader: DataLoader, device: torch.device):
        """Extracts features from the exit point of the frozen backbone (post-layer3)."""
        backbone.eval()
        all_features = []
        all_labels = []
        
        with torch.no_grad():
            for inputs, labels in loader:
                inputs = inputs.to(device)
                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 apply(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
        """
        Extracts intermediate features and updates the parameters of the last blocks
        using the exact inverse-Hessian influence step.
        """
        device = next(model.parameters()).device
        
        # 1. Slice the ResNet graph structural components
        backbone, unlearning_head = self._split_model(model)
        
        print(">> Extracting intermediate structural features from layer3 exit...")
        retain_feats, retain_labels = self._get_intermediate_features(backbone, retain_loader, device)
        forget_feats, forget_labels = self._get_intermediate_features(backbone, forget_loader, device)
        
        # 2. Flatten target weights from the active head into a 1D optimization tensor
        # For simplicity and mathematical stability, we isolate the final layer's weights
        # inside the active head for the exact Hessian tracking step
        target_layer = unlearning_head[-1] # This points straight to model.fc
        w = target_layer.weight.data.clone().cpu()
        
        # 3. Compute Exact Hessian over intermediate embeddings
        # ResNet50's layer4 expands channels to 2048, creating a 2048x2048 matrix context
        print(">> Computing exact sub-graph Hessian matrix...")
        N_retain = retain_feats.size(0)
        
        # Pool the feature maps if they haven't been flattened yet by the head module
        if len(retain_feats.shape) > 2:
            retain_flat = torch.mean(retain_feats, dim=[2, 3])
            forget_flat = torch.mean(forget_feats, dim=[2, 3])
        else:
            retain_flat = retain_feats
            forget_flat = forget_feats

        X_T_X = torch.matmul(retain_flat.t(), retain_flat)
        reg_matrix = self.l2_reg * torch.eye(retain_flat.size(1))
        Hessian = (X_T_X / N_retain) + reg_matrix
        
        # 4. Calculate gradients relative to the forgotten target features
        print(">> Calculating forget set gradients...")
        num_classes = w.size(0)
        forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float()
        
        preds_forget = torch.matmul(forget_flat, w.t())
        error = preds_forget - forget_labels_one_hot
        grad_forget = torch.matmul(error.t(), forget_flat) / forget_flat.size(0)

        # 5. Apply Newton Step optimization update
        print(">> Inverting optimization subspace via system solver...")
        try:
            delta_w_t = torch.linalg.solve(Hessian, grad_forget.t())
            delta_w = delta_w_t.t()
        except RuntimeError:
            print(">> Warning: Subspace Hessian is singular. Using pseudo-inverse fallback.")
            delta_w = torch.matmul(grad_forget, torch.linalg.pinv(Hessian).t())

        # 6. Apply Weight Adjustment Bounds Check
        new_w = w + delta_w
        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. Subspace Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}")

        # 7. Write weights directly back into the live ResNet50 instance
        model.fc.weight.data = new_w.to(device)
        
        print(">> Last K Layers Certified Removal complete.")
        return model