from torchvision import datasets, transforms, models
import torch
import numpy as np

# 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([
        # Just standard resize to 224x224
        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():
    return datasets.CelebA(
        root='./data', 
        split='all', 
        target_type='identity', 
        download=True, 
        transform=None
    )


def get_ids_and_counts(dataset):
    return  torch.unique(
        dataset.identity, 
        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 50 identities
    return np.random.choice(eligible_ids, class_size, replace=False)

# optional function to get max amount of samples per class
def select_balanced_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)


# split class images to train and test set.
def get_indices(dataset, identities, split_at):
    train_indices = []
    test_indices = []

    #training_sample = int(sample_size * training_ratio)
    for person_id in identities:
        # Get all indices for this specific person
        indices = torch.where(dataset.identity == 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:])

    return train_indices, test_indices