import torch
import numpy as np
from abc import ABC, abstractmethod
from torchvision import transforms, datasets
from torch.utils.data import Dataset, DataLoader, Subset

class Data(ABC):
    """
    Handles image pipelines, identity filtering, indexing, and unlearning splits.
    """
    def __init__(self, res: int = 224, sample_size = 30, class_size = 20):
        self.res = res
        self.sample_size = sample_size
        self.class_size = class_size
        self.target = None # will have to be set in get_set()

    def train_transform(self):
        return transforms.Compose([
            # ResNet expects 224 x 224 res 
            # Inception expects 299 x 299
            transforms.Resize((self.res, self.res)),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.ColorJitter(
                brightness=0.2, 
                contrast=0.2, 
                saturation=0.1
            ),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406], 
                std=[0.229, 0.224, 0.225]
            )
        ])


    def test_transform(self):
        return transforms.Compose([
            transforms.Resize((self.res, self.res)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406], 
                std=[0.229, 0.224, 0.225]
            )
        ])

    @abstractmethod
    def get_set(self)-> datasets.TorchDataset:
        """Loads and returns the raw underlying PyTorch Dataset instance."""
        pass

    def get_targets(self) -> torch.Tensor:
        return self.target

    def get_ids_and_counts(self) -> tuple[torch.Tensor, torch.Tensor]:
        if self.target is None:
            raise ValueError ("This should be called after the 'target' variable has been set.")
        return torch.unique(
            self.target, 
            return_counts=True
        )

    def select_ids(self) -> np.ndarray:
        ids, counts = self.get_ids_and_counts()
        eligible_mask = counts >= self.sample_size
        eligible_ids = ids[eligible_mask].numpy()

        if len(eligible_ids) < self.class_size:
            raise ValueError(
                f"Only found {len(eligible_ids)} identities with {self.sample_size}+ images."
            )

        return np.random.choice(eligible_ids, self.class_size, replace=False)

    # Function to get max amount of samples per class
    def select_top_ids(self) -> np.ndarray:

        ids, counts = self.get_ids_and_counts()
        # sort by number of images (descending)
        sorted_indices = torch.argsort(counts, descending=True)
        top_ids = ids[sorted_indices][:self.class_size].numpy()
        return np.array(top_ids, dtype=int)

    def get_indices(self, identities: np.ndarray, split_at: int, max_size: int = None) -> tuple[list, list]:
        '''train_indices = []
        test_indices = []
        max_size = self.sample_size if max_size is None else max_size
        
        # Pull raw target tensor array using concrete implementation rules
        all_targets = np.array(self.get_targets().cpu())
        np.random.seed(42)

        for person_id in identities:
            indices = np.where(all_targets == person_id)[0]
            np.random.shuffle(indices)
            
            # Constrain total sample tracking size if requested (e.g. CelebA ceiling)
            current_pool = indices[:max_size] if max_size else indices
            
            if split_at >= len(current_pool):
                raise ValueError(f"Split point ({split_at}) exceeds slice size ({len(current_pool)}) for class {person_id}.")

            train_indices.extend(current_pool[:split_at])
            test_indices.extend(current_pool[split_at:])

        return train_indices, test_indices'''
        if split_at >= self.sample_size: # debug safety
            raise ValueError(f"Split point ({split_at}) must be less than total size ({self.sample_size}).")
    
        train_indices = []
        test_indices = []

        #training_sample = int(sample_size * training_ratio)
        np.random.seed(42)
        target = self.get_targets()
        for person_id in identities:
            # Get all indices for this specific person
            indices = torch.where(target == person_id)[0].numpy()

            # Shuffle the indices for this person
            np.random.shuffle(indices)
            
            # split data to testing and training
            train_indices.extend(indices[:split_at])
            test_indices.extend(indices[split_at:self.sample_size])

        return train_indices, test_indices

    @staticmethod
    def get_unlearn_loaders(
        dataset: Dataset, 
        forget_class_idx: int, 
        batch_size: int = 32
    ) -> tuple[DataLoader, DataLoader]:

        """Splits an IdentitySubset into forget/retain parts based on local class index."""
        if hasattr(dataset, 'targets'):
            targets = dataset.targets
        elif hasattr(dataset, 'identity'): 
            targets = dataset.identity
        else:
            targets = [dataset[i][1] for i in range(len(dataset))]
            
        if not isinstance(targets, torch.Tensor):
            targets = torch.tensor(targets)
            
        forget_indices = torch.where(targets == forget_class_idx)[0].tolist()
        retain_indices = torch.where(targets != forget_class_idx)[0].tolist()
        
        forget_subset = Subset(dataset, forget_indices)
        retain_subset = Subset(dataset, retain_indices)
        
        forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=False)
        retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True)
        
        print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.")
        
        return forget_loader, retain_loader
    

    @staticmethod
    def getDataSet(set:SetType, sample_size):
        # some test
        if set == SetType.CASIA:
            from sets.CasiaFace import CasiaFace
            return CasiaFace(sample_size = sample_size)
        if set == SetType.CELEBA:
            from sets.CelebA import CelebA
            return CelebA(sample_size=sample_size)


from enum import Enum, auto
class SetType(Enum):
    CASIA = auto()
    CELEBA = auto()