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

class LinearFiltration(Strategy):
    def __init__(self, target_class_index):
        super().__init__(target_class_index=target_class_index)

    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
        model.eval()
        # Freeze internal params
        for param in model.parameters():
            param.requires_grad = False

        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 _sums_and_counts(self, model, num_classes, loader, device, forget_index, h_dim):
        model.eval()

        sums = torch.zeros(num_classes, h_dim, device=device)
        counts = torch.zeros(num_classes, device=device)

        # Generate values for retain
        with torch.no_grad():
            for inputs, targets in loader:
                inputs = inputs.to(device)
                targets = targets.to(device)
                # predictions
                outputs = model(inputs)

                for j in range(num_classes):
                    if j == forget_index:
                        continue
                    mask = (targets == j)

                    if mask.any():
                        sums[j] += outputs[mask].sum(dim=0)
                        counts[j] += mask.sum()
        
        return sums, counts
    

    # 
    def _get_means(self,model, num_classes, loader, device, forget_index):
        h_dim = model.fc.out_features

        # all predictions
        sums, counts = self._sums_and_counts(
            model=model,
            num_classes=num_classes,
            loader=loader,
            device=device,
            forget_index=forget_index,
            h_dim=h_dim
        )

        #A = []

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

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

        K = tensor.shape[0] 
        
        # pi_a_forget should match the feature space dimensions (h_dim)
        pi_a_f = torch.zeros(tensor.shape[1], device=tensor.device)

        t_1 = pi_a_f
        # Extracting the 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


    # Normalisation filtration
    def normalise(self, model, retain_loader, forget_loader, device, forget_index):
        W = model.fc.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
        A = self._get_means(
            model=model,
            num_classes=num_classes,
            loader=full_loader,
            device=device,
            forget_index=forget_index
        )
        # 9
        Z = self._compute_z(tensor=A, forget_index=forget_index)
        B_Z_rows = []

        for i in range(num_classes):
            if i == forget_index:
                B_Z_rows.append(Z)
            else:
                # Retained classes maintain their original ideal feature directions
                B_Z_rows.append(A[i])
        
        # 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(A)
        # 11 
        W_Z = B_Z @ A_inv @ W

        # 12
        model.fc.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
        )