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

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
        )
    
    # FIX: Added staticmethod decorator
    @staticmethod
    def get_features(model, inputs):
        # For ResNet, pass through everything up to the fc layer
        x = model.conv1(inputs)
        x = model.bn1(x)
        x = model.relu(x)
        x = model.maxpool(x)

        x = model.layer1(x)
        x = model.layer2(x)
        x = model.layer3(x)
        x = model.layer4(x)

        x = model.avgpool(x)
        x = torch.flatten(x, 1)
        return x

    @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
        
        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
    

    @staticmethod
    def _sums_and_counts(model, num_classes, retain_loader, forget_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 retain_loader:
                inputs = inputs.to(device)
                targets = targets.to(device)
                # FIX: Call get_features instead of model() directly
                outputs = LinearFiltration.get_features(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()
        
        # Values for forget
        with torch.no_grad():
            for inputs, targets in forget_loader:
                inputs = inputs.to(device)
                targets = targets.to(device)
                # FIX: Call get_features instead of model() directly
                outputs = LinearFiltration.get_features(model, inputs)

                mask = (targets == forget_index)

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

        return sums, counts
    
    @staticmethod
    def _get_means(model, num_classes, retain_loader, forget_loader, device, forget_index):
        h_dim = model.fc.in_features

        sums, counts = LinearFiltration._sums_and_counts(
            model=model,
            num_classes=num_classes,
            retain_loader=retain_loader,
            forget_loader=forget_loader,
            device=device,
            forget_index=forget_index,
            h_dim=h_dim
        )
        A = []

        for i in range(num_classes):
            if counts[i] > 0:
                A.append(sums[i] / counts[i])
            else:
                A.append(torch.zeros(h_dim, device=device))
        
        # CORRECT: Stack along dim=0 to make it (num_classes, h_dim)
        return torch.stack(A, dim=0)
    

    @staticmethod
    def _compute_z(tensor, forget_index):
        # Now tensor has shape (num_classes, h_dim) -> tensor.shape[0] is num_classes
        K = tensor.shape[0] 
        
        # pi_a0 should match the feature space dimensions (h_dim)
        pi_a0 = torch.zeros(tensor.shape[1], device=tensor.device)

        t_1 = pi_a0
        a0 = tensor[forget_index, :] # Extracting the row vector for the forgotten class

        mask_a0 = torch.ones(
            a0.shape[0],
            dtype=torch.bool, 
            device=tensor.device
        )
        # We compute the target shift over features
        t_2 = -(1.0 / (K - 1)) * a0[mask_a0].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


    @staticmethod
    def normalise(model, retain_loader, forget_loader, device, forget_index):
        W = model.fc.weight.data.clone()
        num_classes = W.shape[0]
        
        A = LinearFiltration._get_means(
            model=model,
            num_classes=num_classes,
            retain_loader=retain_loader,
            forget_loader=forget_loader,
            device=device,
            forget_index=forget_index
        )

        Z = LinearFiltration._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])
        
        # Stack back along dim=0 to match (num_classes, h_dim)
        B_Z = torch.stack(B_Z_rows, dim=0)

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

        model.fc.weight.copy_(W_Z)

        return model