import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, ConcatDataset, Subset
from unlearning.Strategy import Strategy
import numpy as np
from sklearn.metrics import classification_report
from architectures.WFNet import WF_Net_Model

from sets.Data import vertical_split

class WeightFiltration(Strategy):
    def __init__(self, 
        target_class_index: int, 
        arch,
        num_classes: int = 20, 
        epochs: int = 6, 
        lr: float = 100.0, 
        gamma: float = 0.01,
        lambda_1 = 25           
    ):
        super().__init__(target_class_index=target_class_index)
        self.epochs = epochs
        self.lr = lr
        self.gamma = gamma 
        self.num_classes = num_classes
        self.wf_model = None 
        self.lambda_1 = lambda_1
        self.arch = arch
    

    def _optimise_filter(self, model: nn.Module, retain_loader: DataLoader, forget_loader: DataLoader, device) -> nn.Module:

        # new WF_Model instance
        wf_model = WF_Net_Model(
            device=device, 
            arch=self.arch,
            size=self.num_classes, 
            original_model=model, 
            target_class_index=self.target_class_index
        )

        wf_net = wf_model.get()
        optimizer = optim.SGD([wf_net.alpha], lr=self.lr)
        # Use reduction='none' so we can manipulate individual item losses
        criterion_none = nn.CrossEntropyLoss(reduction='none')
        criterion_mean = nn.CrossEntropyLoss()

        for epoch in range(self.epochs):
            t_loss_r, t_loss_f = 0.0, 0.0
            steps = 0
            # forget and retain
            for (r_inputs, r_labels), (f_inputs, f_labels) in zip(retain_loader, forget_loader):
                r_inputs, r_labels = r_inputs.to(device), r_labels.to(device)
                f_inputs, f_labels = f_inputs.to(device), f_labels.to(device)
                
                optimizer.zero_grad()
                
                # retain data paired with randomly selected rows of alpha to compute the retaining loss 
                random_offset = torch.randint(0, self.num_classes - 1, size=r_labels.shape, device=device)
                gate_signals_r = torch.where(random_offset >= r_labels, random_offset + 1, random_offset)
                
                outputs_r = wf_net(r_inputs, target_class_indices=gate_signals_r)
                loss_r = criterion_mean(outputs_r, r_labels)

                # Forget set is paired with corresponding labels as row selectors for alpha 
                # and used to compute unlearning loss
                outputs_f = wf_net(f_inputs, target_class_indices=f_labels)
                
                # Calculate loss for every single item in the batch at once
                per_item_forget_loss = criterion_none(outputs_f, f_labels)
                
                # Use a scatter/sum approach to get class-wise losses without a Python loop
                # Create a mask of unique classes present in this batch
                unique_classes, inverse_indices = torch.unique(f_labels, return_inverse=True)
                classes_in_batch = unique_classes.size(0)
                
                if classes_in_batch > 0:
                    # average CE loss per class
                    class_loss_sums = torch.zeros(classes_in_batch, device=device)
                    class_loss_sums.scatter_add_(0, inverse_indices, per_item_forget_loss)
                    
                    class_counts = torch.zeros(classes_in_batch, device=device)
                    class_counts.scatter_add_(0, inverse_indices, torch.ones_like(per_item_forget_loss))
                    
                    mean_class_ce_loss = class_loss_sums / class_counts
                    
                    # Poppi et al. suggest employing reciprocal of the forget loss 
                    # to avoid shortcomings of negative gradient approach
                    loss_f = torch.mean(1.0 / (mean_class_ce_loss + 1e-6))
                else:
                    loss_f = torch.tensor(0.0, device=device)

                # Regularisation penalty
                loss_reg = torch.sum(1.0 - torch.sigmoid(wf_net.alpha))
                                
                # Backpropagation
                total_loss = loss_r + (self.lambda_1 * loss_f) + (self.gamma * loss_reg)
                total_loss.backward()
                optimizer.step()
                                
                # Keep tracking stats
                t_loss_r += loss_r.item()
                t_loss_f += loss_f.item()
                steps += 1
                            
            print(f"  Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {t_loss_f/steps:.4f}")
        
        return wf_model
    

    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
        device = next(model.parameters()).device
        model.eval()
        
        if self.wf_model is None:
            print("Initializing and compiling global WF-Net matrix (Run Once for all classes)...")
            
            self.wf_model = self._optimise_filter(
                model, 
                retain_loader=retain_loader, 
                forget_loader=forget_loader,
                device=device
            )
        else:
            print(f"Gating matrix loaded. Switching layout to target class index: {self.target_class_index}")
            self.wf_model.target_class_index = self.target_class_index
            
        return self.wf_model
    
    def _split_data(self, dataset):
        return vertical_split(
            dataset= dataset,
            batch_size=32,
            num_classes=self.num_classes
        )