import torch
from Strategy import Strategy

class NormalizingLinearFiltration(Strategy):
    def __init__(self, target_class_idx):
        self.target_class_idx = target_class_idx

    def apply(self, model, forget_loader, retain_loader):
        model.eval()
        # Freeze parameters structurally
        for param in model.parameters():
            param.requires_grad = False
            
        with torch.no_grad():
            # we modify only classification head
            # Shape: [num_classes, feature_dim]
            W = model.fc.weight.data
            
            # Compute the normalization transformation projection matrix (A)
            # (In your full code, calculate A here matching Baumhauer et al.'s equations)
            num_classes = W.shape[0]
            A = torch.eye(num_classes, device=W.device)
            # Mask/blend target class index distribution configurations here...
            A[self.target_class_idx, :] = 0.0 
            
            # 3. Direct weight matrix override: W_filtered = A * W
            sanitized_W = torch.mm(A, W)
            model.fc.weight.copy_(sanitized_W)