import torch.nn as nn
import time
import os
from pathlib import Path
from torch.utils.data import DataLoader
import Util

class Strategy:
    """Abstract base class for unlearning algorithms with automated, strategy-specific logging."""
    
    def __init__(self, target_class_index):
        # Dynamically set file name based on the class name (e.g., 'NormalizingLinearFiltration.txt')
        self.strategy_name = self.__class__.__name__
        self.target_class_index = target_class_index

    def set_target_class(self, target_class_index: int):
        """Dynamically switch the unlearning target without retraining."""
        self.target_class_index = target_class_index
        

    def apply(self, model: nn.Module, dataset) -> nn.Module:
        log_file = Path(f"reports/{self.strategy_name}/{model.__class__.__name__}/time_metrics.txt")
        Util._initialize_log_file(log_file=log_file)
        """
        Wraps the unlearning execution with automated timing and strategy-specific logging.
        DO NOT override this method in subclasses. Override _run instead.
        """
    


        retain_loader, forget_loader = self._split_data(dataset)
        
        # record start time to evaluate efficiency
        start_time = time.perf_counter()
        # Execute core unlearning logic
        processed_model = self._run(model, forget_loader, retain_loader)
        
        end_time = time.perf_counter()
        execution_time = end_time - start_time
        
        # Log to the strategy's specific file
        Util.log_metric(
            log_file=log_file, 
            execution_time=execution_time
        )

        print(f"[{self.strategy_name}] Completed in {execution_time:.6f} seconds. Saved to {log_file}")
        
        return processed_model

    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
        """Subclasses implement their core unlearning logic here."""
        raise NotImplementedError
    
    '''
    different strategies split data in to different partitions differently. 
    So a strategy will implement its own and since this part is startegy specific. 
    not all should compute it the same.
    '''
    def _split_data(self,dataset):
        pass