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

class WeightFiltration(Strategy):
    """
    Implements Poppi et al.'s Weight Filtering framework for linear layers.
    Uses a standard functional hook to guarantee native PyTorch autograd tracking.
    """
    def __init__(self, target_class_index,num_classes: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
        super().__init__(target_class_index = target_class_index)
        self.num_classes = num_classes
        self.epochs = epochs
        self.lr = lr
        self.gamma = gamma 
        self.alpha = None
        self.hook_handle = None

    
    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
        device = next(model.parameters()).device
        model.eval() 
        
        # Locate layer4 for dynamic optimization
        target_layer = model.layer4 if hasattr(model, 'layer4') else None
        fc_layer = model.fc if hasattr(model, 'fc') and isinstance(model.fc, nn.Linear) else None
        
        if target_layer is None or fc_layer is None:
            raise AttributeError("Model does not have the required layers.")
        
        # Match alpha dimensions to the channels outputted by layer4
        num_features = fc_layer.weight.shape[1] 
        self.alpha = nn.Parameter(torch.ones(self.num_classes, num_features, device=device) * 1.5)

        # Freeze everything except our channel mask
        for p in model.parameters():
            p.requires_grad = False
        self.alpha.requires_grad = True

        # Hook into layer4 dynamically to run the untraining optimization
        self.hook_handle = target_layer.register_forward_hook(self._get_hook())
        # optimise the filter to maintain accuracy on retain set
        # and decrease accuracy on forget set
        self._optimise_filter(model, forget_loader, retain_loader, device)

        # Remove the runtime hook
        self.hook_handle.remove()
        
        # Transfer the channel suppression permanently into model.fc
        with torch.no_grad():
            mask = torch.sigmoid(self.alpha[self.target_class_index]) # Shape: (num_features,)
            
            # Suppress the channels ONLY for the target class row in fc
            fc_layer.weight[self.target_class_index].copy_(
                fc_layer.weight[self.target_class_index] * mask
            )
            print(f">> Baked deep channel filter into Class {self.target_class_index} weights.")
            
        return model

    def _get_hook(self):
        """
        Filters the internal feature map channels of layer4.
        The mask scales the channels across the batch.
        """
        def functional_hook(module, layer_input, layer_output):
            # layer_output shape: (batch, channels, height, width) -> e.g., (16, 2048, 7, 7)
            # self.alpha shape: (num_classes, channels) -> e.g., (20, 2048)
            
            # Extract 1D mask for the target class: (channels,)
            mask = torch.sigmoid(self.alpha[self.target_class_index])
            
            # Reshape mask to (1, channels, 1, 1) so it broadcasts over batch, height, and width
            mask = mask.view(1, -1, 1, 1)
            
            # Scale the internal feature maps before they move to the next layer
            return layer_output * mask
            
        return functional_hook
    
    
    def _optimise_filter(self, model, forget_loader, retain_loader, device):
        optimizer = optim.Adam([self.alpha], lr=self.lr)
        criterion = nn.CrossEntropyLoss()

        print(f"[{self.__class__.__name__}] Unlearning Class {self.target_class_index} with gamma={self.gamma}...")

        # To optimise this loop we will watch improvements after each optimisation
        temp_forget_loss = None
        # this can be adjusted to optimise the best escape point
        # it is the value we set to evaluate performance improvement after each itteration.
        # if improvement is less than this, then we break itteration.
        threshold = 0.05

        for epoch in range(self.epochs):
            forget_iter = iter(forget_loader)
            t_loss_r, t_loss_f = 0.0, 0.0
            steps = 0
            
            for r_inputs, r_labels in retain_loader:
                r_inputs, r_labels = r_inputs.to(device), r_labels.to(device)
                
                try:
                    f_inputs, _ = next(forget_iter)
                except StopIteration:
                    forget_iter = iter(forget_loader)
                    f_inputs, _ = next(forget_iter)
                f_inputs = f_inputs.to(device)
                
                optimizer.zero_grad()
                
                # Compute Losses
                # The hook handles the weight filtering smoothly behind the scenes
                loss_r = criterion(model(r_inputs), r_labels)
                loss_f = -torch.sum((torch.ones_like(model(f_inputs)) / self.num_classes) * torch.log_softmax(model(f_inputs), dim=-1))
                
                total_loss = loss_r + (self.gamma * loss_f)
                total_loss.backward()
                optimizer.step()
                
                t_loss_r += loss_r.item()
                t_loss_f += loss_f.item()

                steps += 1
            forget_loss = t_loss_f / steps
            print(f"  Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {forget_loss:.4f}")

            if temp_forget_loss is not None:

                improvement = temp_forget_loss - forget_loss
                # if optimisation reaches a point of diminishing returns (improvements is less than threshold)
                # we break the loop
                if improvement < threshold:
                    break
            # else we update the lasst recorded loss.
            temp_forget_loss = forget_loss