import torch
import torch.nn as nn
from torch.utils.data import DataLoader, RandomSampler
from torch.autograd import grad
from unlearning.Strategy import Strategy

class CertifiedRemoval(Strategy):
    """
    Implements Certified Unlearning for non-convex DNNs (Zhang et al.).
    Uses a modified, stabilized stochastic Newton step using Taylor-expansion 
    HVP estimation across the entire parameter space, capped with calibrated noise.
    """
    def __init__(self, target_class_index: int, l2_reg: float = 0.0005, 
                 gamma: float = 0.01, scale: float = 1000.0, 
                 s1: int = 10, s2: int = 1000, std: float = 0.001, unlearn_bs: int = 2):
        super().__init__(target_class_index)
        self.l2_reg = l2_reg        
        self.gamma = gamma          
        self.scale = scale          
        self.s1 = s1                
        self.s2 = s2                
        self.std = std              
        self.unlearn_bs = unlearn_bs

    '''
    def _compute_loss_gradient(self, model, loader, device: torch.device):
        model.eval()
        criterion = nn.CrossEntropyLoss(reduction='sum')
        params = [p for p in model.parameters() if p.requires_grad]
        grad_accumulator = [torch.zeros_like(p).cpu() for p in params]
        total_samples = 0
        
        for data, targets in loader:
            total_samples += targets.shape[0]
            data, targets = data.to(device), targets.to(device)
            outputs = model(data)
            
            mini_grads = list(grad(criterion(outputs, targets), params))
            for i in range(len(grad_accumulator)):
                grad_accumulator[i] += mini_grads[i].cpu().detach()

        for i in range(len(grad_accumulator)):
            grad_accumulator[i] /= total_samples

        l2_reg_term = 0.0
        for param in model.parameters():
            l2_reg_term += torch.norm(param, p=2)
        
        reg_grads = list(grad(self.l2_reg * l2_reg_term, params))
        for i in range(len(grad_accumulator)):
            grad_accumulator[i] += reg_grads[i].cpu().detach()
            
        return [p.to(device) for p in grad_accumulator]'''
    def _compute_loss_gradient(self, model, loader, device: torch.device):
        model.eval()
        # Use reduction='sum' matching the original framework
        criterion = nn.CrossEntropyLoss(reduction='sum')
        params = [p for p in model.parameters() if p.requires_grad]
        grad_accumulator = [torch.zeros_like(p).cpu() for p in params]
        total_samples = 0
        
        for data, targets in loader:
            total_samples += targets.shape[0]
            data, targets = data.to(device), targets.to(device)
            outputs = model(data)
            
            loss = criterion(outputs, targets)
            
            # Incorporate L2 weight regularization directly inside the backprop graph
            # to keep scaling bounded and aligned with the data volume
            l2_reg_term = 0.0
            for param in model.parameters():
                if param.requires_grad:
                    l2_reg_term += torch.norm(param, p=2)
            
            total_loss = loss + (self.l2_reg * l2_reg_term)
            
            mini_grads = list(grad(total_loss, params, retain_graph=False))
            for i in range(len(grad_accumulator)):
                grad_accumulator[i] += mini_grads[i].cpu().detach()

        for i in range(len(grad_accumulator)):
            grad_accumulator[i] /= total_samples
            
        return [p.to(device) for p in grad_accumulator]

    
    def grad_batch(batch_loader, lam, model, device):
        model.eval()
        criterion = nn.CrossEntropyLoss(reduction='sum')
        params = [p for p in model.parameters() if p.requires_grad]
        grad_batch = [torch.zeros_like(p).cpu() for p in params]
        num = 0
        for batch_idx, (data, targets) in enumerate(batch_loader):
            num += targets.shape[0]
            data, targets = data.to(device), targets.to(device)
            outputs = model(data)

            grad_mini = list(grad(criterion(outputs, targets), params))
            for i in range(len(grad_batch)):
                grad_batch[i] += grad_mini[i].cpu().detach()

        for i in range(len(grad_batch)):
            grad_batch[i] /= num

        l2_reg = 0
        for param in model.parameters():
            l2_reg += torch.norm(param, p=2)
        grad_reg = list(grad(lam * l2_reg, params))
        for i in range(len(grad_batch)):
            grad_batch[i] += grad_reg[i].cpu().detach()
        return [p.to(device) for p in grad_batch]

    def _hvp(self, loss, params, v):
        first_grads = grad(loss, params, retain_graph=True, create_graph=True)
        elemwise_products = 0
        for grad_elem, v_elem in zip(first_grads, v):
            elemwise_products += torch.sum(grad_elem * v_elem)
        # FIX 1: Set create_graph to False to prevent massive nested graph accumulation
        return grad(elemwise_products, params, create_graph=False)

    def _stochastic_newton_update(self, g, retain_dataset, model, device):
        model.eval()
        criterion = nn.CrossEntropyLoss()
        params = [p for p in model.parameters() if p.requires_grad]
        h_res = [torch.zeros_like(p) for p in g]
        
        for _ in range(self.s1):
            h_estimate = [p.clone() for p in g]
            sampler = RandomSampler(retain_dataset, replacement=True, num_samples=self.unlearn_bs * self.s2)
            res_loader = DataLoader(retain_dataset, batch_size=self.unlearn_bs, sampler=sampler)
            res_iter = iter(res_loader)
            
            for j in range(self.s2):
                try:
                    data, target = next(res_iter)
                except StopIteration:
                    res_iter = iter(res_loader)
                    data, target = next(res_iter)
                    
                data, target = data.to(device), target.to(device)
                outputs = model(data)
                
                loss = criterion(outputs, target)
                l2_reg_term = 0.0
                for param in model.parameters():
                    l2_reg_term += torch.norm(param, p=2)
                loss += (self.l2_reg + self.gamma) * l2_reg_term
                
                h_s = self._hvp(loss, params, h_estimate)
                
                with torch.no_grad():
                    for k in range(len(params)):
                        # FIX 2: Added .detach() to decouple history strings across iterative update blocks
                        #h_estimate[k] = (h_estimate[k] + g[k] - h_s[k] / self.scale).detach()
                        next_estimate = h_estimate[k].data + g[k].data - (h_s[k].data / self.scale)
                        h_estimate[k] = next_estimate.clone()
                del h_s, loss, outputs
                        
            for k in range(len(params)):
                h_res[k] = h_res[k] + h_estimate[k] / self.scale
                
        return [p / self.s1 for p in h_res]

    '''def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
        device = next(model.parameters()).device

        num_forget = len(forget_loader.dataset)
        num_retain = len(retain_loader.dataset)
        scaling_ratio = num_forget / num_retain
        
        print(">> Calculating base gradients over target FORGET set...")
        # FIX 3: Base gradients MUST be evaluated from forget_loader to drop target class distributions
        g = self._compute_loss_gradient(model, forget_loader, device)
        
        print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...")
        retain_dataset = retain_loader.dataset 
        delta = self._stochastic_newton_update(g, retain_dataset, model, device)
        
        print(">> Applying stabilized parameter adjustments and randomized certification noise...")
        with torch.no_grad():
            for i, param in enumerate(model.parameters()):
                if param.requires_grad:
                    noise = self.std * torch.randn(param.data.size(), device=device)
                    #param.data.add_(-delta[i] + noise)
                    param.data.add_(scaling_ratio * delta[i] + noise)
                
        print(">> Certified Unlearning process completed successfully across the complete landscape.")
        return model'''
    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
        device = next(model.parameters()).device

        print(">> Calculating stable base gradients over the RETAIN set...")
        # To match the author's snippet perfectly, g MUST be computed on the retain data.
        # If this loader is too large for your VRAM, use a smaller batch size (e.g. 16 or 32) 
        # in your main training script when creating retain_loader.
        g = self._compute_loss_gradient(model, retain_loader, device)
        
        print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...")
        retain_dataset = retain_loader.dataset 
        delta = self._stochastic_newton_update(g, retain_dataset, model, device)
        
        print(">> Applying parameter removal adjustments (-delta)...")
        with torch.no_grad():
            for i, param in enumerate(model.parameters()):
                if param.requires_grad:
                    noise = self.std * torch.randn(param.data.size(), device=device)
                    
                    # MATCHING THE SNIPPET: Subtract delta exactly as the authors do
                    # This removes the influence trace of the omitted data.
                    param.data.add_(-delta[i] + noise)
                
        print(">> Certified Unlearning process completed successfully.")
        return model