import time
from pathlib import Path
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from unlearning.Strategy import Strategy

class Retrain(Strategy):
    """
    Implements the Exact Unlearning Baseline by retraining the model architecture
    completely from scratch using only the retained dataset partition.
    """
    def __init__(self, target_class_index: int, lr: float = 0.01, 
                 weight_decay: float = 0.0005, epochs: int = 5):
        super().__init__(target_class_index)
        self.lr = lr
        self.weight_decay = weight_decay
        self.epochs = epochs

    def _reset_weights(self, model: nn.Module):
        """
        Re-initializes all learnable parameters of the model to clear pre-trained memories.
        """
        inner_model = getattr(model, "model", model)
        for layer in inner_model.modules():
            if hasattr(layer, 'reset_parameters'):
                layer.reset_parameters()

    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
        device = next(model.parameters()).device
        
        print(">> Triggering Exact Unlearning Baseline (Retraining from scratch)...")
        
        # 1. Clear the pre-trained state completely
        self._reset_weights(model) # model should be loaded here or weights reset to ImageNet (pretrained default)
        model.train()
        
        # fresh optimizer for this clean environment
        optimizer = optim.SGD(
            model.parameters(), 
            lr=self.lr, 
            momentum=0.9, 
            weight_decay=self.weight_decay
        )
        criterion = nn.CrossEntropyLoss()

        # 3. Standard training loop over the Retain set
        for epoch in range(self.epochs):
            running_loss = 0.0
            total_samples = 0
            
            for data, targets in retain_loader:
                data, targets = data.to(device), targets.to(device)
                
                optimizer.zero_grad()
                outputs = model(data)
                loss = criterion(outputs, targets)
                
                loss.backward()
                optimizer.step()
                
                running_loss += loss.item() * targets.size(0)
                total_samples += targets.size(0)
                
            epoch_loss = running_loss / total_samples
            print(f"   [Retrain] Epoch {epoch+1}/{self.epochs} completed. Loss: {epoch_loss:.4f}")

        print(">> Retraining pipeline finished. Baseline weights established.")
        return model

    def _split_data(self, dataset):
        # Dynamically pulls loaders from your Data.py script
        from sets.Data import get_unlearning_loaders
        return get_unlearning_loaders(
            dataset=dataset,
            forget_class_idx=self.target_class_index,
            batch_size=32
        )