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

# train set transform 
def train_transform(res):
    return transforms.Compose([
        transforms.Resize((res, 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]
        )
    ])

# 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 using ImageFolder for CASIA-WebFace
'''
def get_set():
    # This will check local cache first, then download if missing
    print("Checking for CASIA-WebFace dataset...")
    path = kagglehub.dataset_download("debarghamitraroy/casia-webface")
    
    # Kagglehub often downloads a nested structure (e.g., path/casia-webface/casia-webface)
    # We need the folder that directly contains the identity subfolders
    # We'll check if there's a 'casia-webface' subfolder inside the downloaded path
    sub_path = os.path.join(path, "casia-webface")
    final_path = sub_path if os.path.exists(sub_path) else path
    
    print(f"Loading dataset from: {final_path}")
    
    return datasets.ImageFolder(
        root=final_path, 
        transform=None
    )'''
# Load data using ImageFolder for your UNPACKED images
def get_set():
    # This must point to the folder created by Extractor.py
    # NOT the kagglehub cache path
    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):
    # ImageFolder stores labels in .targets
    targets = torch.tensor(dataset.targets)
    return torch.unique(
        input = targets,
        return_counts=True
    )

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."
        )

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

def select_balanced_ids(dataset, class_size):
    ids, counts = get_ids_and_counts(dataset=dataset)
    sorted_indices = torch.argsort(counts, descending=True)
    top_ids = ids[sorted_indices][:class_size].numpy()
    return np.array(top_ids, dtype=int)

def get_indices(dataset, identities, split_at):
    train_indices = []
    test_indices = []
    
    # We convert to numpy for faster searching with np.where
    all_targets = np.array(dataset.targets)

    for person_id in identities:
        # Get all indices for this specific person
        indices = np.where(all_targets == person_id)[0]

        # Shuffle the indices for this person
        np.random.shuffle(indices)
        
        # Split data based on your split_at value
        train_indices.extend(indices[:split_at])
        test_indices.extend(indices[split_at:])

    return train_indices, test_indices



# 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_forget_retain_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.
    """
    # 1. Safely extract targets whether it's a standard dataset or a Subset wrapper
    if hasattr(dataset, 'targets'):
        targets = dataset.targets
    elif hasattr(dataset, 'identity'):  # Raw CelebA support
        targets = dataset.identity
    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)
        
    # 2. Generate 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()
    
    # 3. Create PyTorch Subsets
    forget_subset = Subset(dataset, forget_indices)
    retain_subset = Subset(dataset, retain_indices)
    
    # 4. Wrap into clean 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