import torch
import torch.nn as nn
from .Strategy import Strategy

class LinearFiltration(Strategy):
    def __init__(self, target_class_idx: int):
        super().__init__()  # Automatically configures 'NormalizingLinearFiltration_metrics.txt'
        self.target_class_idx = target_class_idx

    def _run(self, model: nn.Module) -> nn.Module:
        model.eval()
        for param in model.parameters():
            param.requires_grad = False
            
        with torch.no_grad():
            W = model.fc.weight.data.clone()
            num_classes = W.shape[0]
            
            A = self._calculate_filtration_matrix(num_classes, self.target_class_idx, W.device)
            sanitized_W = torch.mm(A, W)
            model.fc.weight.copy_(sanitized_W)
            
        return model

    '''@staticmethod
    def _calculate_filtration_matrix(num_classes: int, forget_class: int, device: torch.device) -> torch.Tensor:
        A = torch.eye(num_classes, device=device)
        num_remaining_classes = num_classes - 1
        for j in range(num_classes):
            if j == forget_class:
                A[forget_class, j] = 0.0
            else:
                A[forget_class, j] = 1.0 / num_remaining_classes
        return A'''
    @staticmethod
    def _calculate_filtration_matrix(num_classes: int, forget_class: int, device: torch.device) -> torch.Tensor:
        A = torch.eye(num_classes, device=device)
        num_remaining = num_classes - 1
        
        # The row of the forgotten class should average the inputs of all other classes
        for j in range(num_classes):
            if j == forget_class:
                A[forget_class, j] = 0.0
            else:
                A[forget_class, j] = 1.0 / num_remaining
                
        return A