import torch
import torch.nn as nn
from .Strategy import Strategy
from torch.utils.data import DataLoader
from sets.Data import get_unlearning_loaders, _combine_set, vertical_split

class LinearFiltration(Strategy):
    def __init__(self, target_class_index, num_classes = 20):
        super().__init__(target_class_index=target_class_index)
        self.A = None
        self.num_classes = num_classes

    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
        model.eval()
        device = next(model.parameters()).device
            
        return self.normalise(
            model=model, 
            retain_loader=retain_loader, 
            forget_loader=forget_loader,
            device=device,
            forget_index=self.target_class_index
        )
    

    def _get_classifier(self, model: nn.Module) -> nn.Linear:

        inner_model = getattr(model, "model", model)
        
        # looking for standard naming conventions in named modules
        for name, module in inner_model.named_modules():
            # Check if it's our target linear layer
            if (name == "fc" or name == "classifier") and isinstance(module, nn.Linear):
                return module
                
            # Handle models (like EfficientNet) where the classifier is a Sequential block
            if name == "classifier" and isinstance(module, nn.Sequential):
                for sub_module in reversed(list(module.children())):
                    if isinstance(sub_module, nn.Linear):
                        return sub_module

        # scan backwards for the last Linear layer
        for module in reversed(list(inner_model.modules())):
            if isinstance(module, nn.Linear):
                return module
                
        raise RuntimeError(f"Could not locate a linear classification head for {model.__class__.__name__}")
    
  

    def _compute_A(self, model, loader, device):
        model.eval()
        
        
        # Initialize tracking tensors
        sums = torch.zeros(self.num_classes, self.num_classes, device=device)
        counts = torch.zeros(self.num_classes, device=device)

        with torch.no_grad():
            for inputs, targets in loader:
                inputs, targets = inputs.to(device), targets.to(device)
                
                # the logit predictions
                outputs = model(inputs) 
                
                # One-hot encode targets to act as a routing mask
                one_hot = torch.nn.functional.one_hot(targets, num_classes=self.num_classes).float()
                
                # add 
                sums += torch.t(one_hot) @ outputs
                
                # Sum columns of one-hot to get counts per class in this batch
                counts += one_hot.sum(dim=0)

        # means
        counts_safe = counts.unsqueeze(1)
        self.A = torch.where(
            counts_safe > 0, 
            sums / counts_safe, 
            torch.zeros_like(sums)
        )

    # 9
    def _compute_z(self, tensor, forget_index):

        K = tensor.shape[0] 
        
        
        pi_a_f = torch.zeros(tensor.shape[1], device=tensor.device)

        t_1 = pi_a_f
        # row vector for the forgotten class
        a_f = tensor[forget_index, :] 

        mask_a_f = torch.ones(
            a_f.shape[0],
            dtype=torch.bool, 
            device=tensor.device
        )
        # We compute the target shift over features
        t_2 = -(1.0 / (K - 1)) * a_f[mask_a_f].sum()

        mask_rows = torch.ones(K, dtype=torch.bool, device=tensor.device)
        mask_rows[forget_index] = False

        r_A = tensor[mask_rows, :]
        t_3 = (1.0 / ((K - 1)) ** 2) * r_A.sum()

        return t_1 + t_2 + t_3
    
    def _pi(self, a_tensor):
        """
        Affine transformation to normalize the logit distribution.
        This maps the logit mean to 0 and scales based on variance 
        to prevent logit shrinkage/expansion 'scars'.
        """
        # Calculate mean and std across the feature dimension
        mean = a_tensor.mean(dim=0, keepdim=True)
        std = a_tensor.std(dim=0, keepdim=True)
        
        # Avoid division by zero
        std = torch.where(std > 1e-6, std, torch.ones_like(std))
        
        return (a_tensor - mean) / std
    

    # Normalisation filtration
    def normalise(self, model, retain_loader, forget_loader, device, forget_index):
        clf = self._get_classifier(model)
        W = clf.weight.data.clone()
        #num_classes = W.shape[0]

        # we combine the data so we can calculate the mean of prdictions
        #full_loader = _combine_set(retain_loader, forget_loader)
        # 8
        # Computing A is the most resource intensive part of this algorithm
        # and to optimise the process, we computr it only once and re-use it
        # because mean of all prdictions is the same for all
        if self.A is None:
            self._compute_A(
                model = model,
                #num_classes = num_classes,
                loader = forget_loader,
                device = device
            )

        # 9
        Z = self._compute_z(tensor=self.A, forget_index=forget_index)
        B_Z_rows = []

        for i in range(self.num_classes):
            if i == forget_index:
                # Normalise the 'erased' target vector Z
                B_Z_rows.append(self._pi(Z.unsqueeze(0)).squeeze(0))
            else:
                # Normalise the retained class vectors
                B_Z_rows.append(self._pi(self.A[i].unsqueeze(0)).squeeze(0))
        
        # 10
        # Stack back along dim=0 to match (num_classes, h_dim)
        # to get mean
        B_Z = torch.stack(B_Z_rows, dim=0)

        A_inv = torch.linalg.pinv(self.A)
        # 11 
        W_Z = B_Z @ A_inv @ W

        # 12
        clf = self._get_classifier(model)
        with torch.no_grad():
            clf.weight.copy_(W_Z)

        return model
    
    # overriden function
    def _split_data(self, dataset):
        '''return get_unlearning_loaders(
            dataset=dataset,
            forget_class_idx=self.target_class_index,
            batch_size = 32
        )'''
        return vertical_split(
            dataset= dataset,
            batch_size=32,
            num_classes=self.num_classes,
            ratio=0.1
        )