import os
from pathlib import Path
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from unlearning.Strategy import Strategy
from architectures.Model import Model, Architecture

class Retrain(Strategy):
    """
    Implements the Exact Unlearning Baseline by re-instantiating a fresh, 
    pre-trained instance of the specific architecture and training it from scratch 
    on the retain set using the Model's internal train function.
    """
    def __init__(self, target_class_index: int, arch: Architecture, size: int, 
                 lr: float = 0.001, epochs: int = 5):
        super().__init__(target_class_index)
        self.arch = arch
        self.size = size
        self.lr = lr
        self.epochs = epochs

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

        device = next(model.parameters()).device

        # we need to check if a retrained copy exists on disk
        checkpoint_path = f"trained_models/class_{self.target_class_index}_retrained.pth"
        if os.path.exists(checkpoint_path):
            print(f"Found existing retrained model checkpoint at '{checkpoint_path}'. Loading parameters directly...")
            
            # Load the state dict using safe configuration flags
            state_dict = torch.load(checkpoint_path, map_location=device, weights_only=True)
            
            # Safely apply the parameter weights to the model in-place
            model.load_state_dict(state_dict)
            print("Retrained parameter loading complete (Retraining bypassed).")
            return model

        # Cache Miss: Execute the standard retraining pipeline
        print(f"No naive model found for class {self.target_class_index} retraining a new one")
        
        
        print(f"Retraining {self.arch.name} from pristine state)...")
        inner_model = getattr(model, "model", model)
        if hasattr(inner_model, "fc"):
            total_classes = inner_model.fc.out_features
        elif hasattr(inner_model, "classifier"):
            # Fallback for alternative architecture layout types
            total_classes = inner_model.classifier[-1].out_features
        else:
            total_classes = self.size
        
        # a new model with default params is created
        fresh = Model.create(self.arch, device, total_classes)
        
        # we train it with retain set
        fresh.train(
            epochs=self.epochs, 
            loader=retain_loader, 
            rate=self.lr, 
            mode="retrain"
        )
        
        # Extract module parameter state dict and copy in place
        model.load_state_dict(fresh.model.state_dict())
        
        print("Retraining pipeline complete")
        return model

    def _split_data(self, dataset):
        from sets.Data import get_unlearning_loaders
        return get_unlearning_loaders(
            dataset=dataset,
            forget_class_idx=self.target_class_index,
            batch_size=16
        )