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

from enum import Enum, auto
class Set_Name(Enum):
    CELEBA = auto()
    CASIAFACES = auto()

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

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

# Load data with 'identity' as target and transform it
def get_set(set_name:Set_Name):
    return fetch_celeb_a() if set_name == Set_Name.CELEBA else fetch_casia_faces()

def fetch_celeb_a():
    return datasets.CelebA(
        root='./data', 
        split='all', 
        target_type='identity', 
        download=True, 
        transform=None
    )

def fetch_casia_faces():
    # location of the data (path relative to project root)
    final_path = os.path.abspath("./data/casia-set")
    
    if not os.path.exists(final_path):
        raise FileNotFoundError(
            f"Unpacked dataset not found at {final_path}. "
            "Please run Extractor.py first!"
        )
    
    print(f"Loading unpacked CASIA dataset from: {final_path}")
    
    return datasets.ImageFolder(
        root=final_path, 
        transform=None
    )



def get_ids_and_counts(dataset):

    target = get_target(dataset=dataset)
    return  torch.unique(
        input = target,
        return_counts = True
    )



# filter selected identities from dataset
# How many classes, how many images per class
def select_ids( dataset, sample_size, class_size):
    ids, counts = get_ids_and_counts(dataset = dataset)
    eligible_mask = counts >= sample_size
    eligible_ids = ids[eligible_mask].numpy()

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

    # Randomly select identities
    return np.random.choice(eligible_ids, class_size, replace=False)

# optional function to get max amount of samples per class
def select_top_ids(dataset, class_size):
    ids, counts = get_ids_and_counts(dataset = dataset)

    # sort by number of images (descending)
    sorted_indices = torch.argsort(counts, descending = True)

    top_ids = ids[sorted_indices][:class_size].numpy()

    return np.array(top_ids, dtype=int)


def get_target(dataset):
    """
    Unified target extractor. 
    Instantly reads raw dataset arrays or safely scales down to unpack wrapped Subsets.
    """
    if hasattr(dataset, 'identity'):
        # celebA
        targets = dataset.identity
    elif hasattr(dataset, 'targets'):
        # others
        targets = dataset.targets
    else:
        # If it's an IdentitySubset or standard Subset, extract mapped targets sequentially
        # This guarantees we get the 0 -> (n-1) remapped labels
        targets = [dataset[i][1] for i in range(len(dataset))]

    if not isinstance(targets, torch.Tensor):
        targets = torch.tensor(targets)
        
    return targets


# split class images to train and test set.
def get_indices(dataset, identities, split_at, size = 30):

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

    target = get_target(dataset=dataset)

    #training_sample = int(sample_size * training_ratio)
    np.random.seed(42)
    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:size])

    return train_indices, test_indices



def get_unlearning_loaders(dataset: Dataset, forget_class_idx: int, batch_size: int = 32) -> tuple[DataLoader, DataLoader]:
    """
    Splits an IdentitySubset or standard Dataset into forget and retain sets 
    based on a remapped target class index.
    """
    # extract targets
    targets = get_target(dataset=dataset)
        
    # mask indices local to this subset
    forget_indices = torch.where(targets == forget_class_idx)[0].tolist()
    retain_indices = torch.where(targets != forget_class_idx)[0].tolist()
    
    # PyTorch Subsets
    forget_subset = Subset(dataset, forget_indices)
    retain_subset = Subset(dataset, retain_indices)
    
    # DataLoaders
    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