import torch
import torch.distributions
from torchvision import datasets
from torch.utils.data import DataLoader
from .augmentations.imagenet_augmentation import get_imageNet_augmentation
from .paths import get_imagenet_path

DEFAULT_TRAIN_BATCHSIZE = 128
DEFAULT_TEST_BATCHSIZE = 128


def get_imagenet_labels():
    path = get_imagenet_path()
    dataset = datasets.ImageNet(path, split='val', transform='none')
    classes_extended = dataset.classes
    labels = []
    for a in classes_extended:
        labels.append(a[0])
    return labels

def get_imagenet_label_wid_pairs():
    path = get_imagenet_path()
    dataset = datasets.ImageNet(path, split='val', transform='none')
    classes_extended = dataset.classes
    wids = dataset.wnids

    label_wid_pairs = []
    for a, b in zip(classes_extended, wids) :
        label_wid_pairs.append((a[0], b))
    return label_wid_pairs

def get_ImageNet(train=True, batch_size=None, shuffle=None, augm_type='test', num_workers=8, size=224, config_dict=None, path=None):
    if batch_size == None:
        if train:
            batch_size = DEFAULT_TRAIN_BATCHSIZE
        else:
            batch_size = DEFAULT_TEST_BATCHSIZE

    augm_config = {}
    transform = get_imageNet_augmentation(type=augm_type, out_size=size, config_dict=augm_config)
    if not train and augm_type != 'none':
        print('Warning: ImageNet test set with ref_data augmentation')

    if shuffle is None:
        shuffle = train
    
    if train == True:
        dataset = datasets.ImageNet(path, split='train', transform=transform)
    else:
        dataset = datasets.ImageNet(path, split='val', transform=transform)

    loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
                                         shuffle=shuffle, num_workers=num_workers, drop_last=True)

    if config_dict is not None:
        config_dict['Dataset'] = 'ImageNet'
        config_dict['Batch out_size'] = batch_size
        config_dict['Augmentation'] = augm_config

    return loader


def get_OpenImageO(batch_size=None, shuffle=True, augm_type='test', num_workers=8, size=224, config_dict=None, path=None, max_samples=None):
    """Load OpenImageO dataset using ImageFolder"""
    if batch_size is None:
        batch_size = DEFAULT_TEST_BATCHSIZE

    augm_config = {}
    transform = get_imageNet_augmentation(type=augm_type, out_size=size, config_dict=augm_config)
    
    dataset = datasets.ImageFolder(path, transform=transform)
    original_size = len(dataset)
    
    # Create random subset if max_samples is specified
    if max_samples is not None and max_samples < len(dataset):
        # Use local random generator with fixed seed for reproducible subset selection
        generator = torch.Generator()
        generator.manual_seed(42)
        indices = torch.randperm(len(dataset), generator=generator)[:max_samples].tolist()
        dataset = torch.utils.data.Subset(dataset, indices)
        print(f"OpenImageO dataset: Using {len(dataset)} samples out of {original_size} total samples")
    else:
        print(f"OpenImageO dataset: Using all {original_size} samples")
    
    loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
                                         shuffle=shuffle, num_workers=num_workers, drop_last=True)

    if config_dict is not None:
        config_dict['Dataset'] = 'OpenImageO'
        config_dict['Batch size'] = batch_size
        config_dict['Augmentation'] = augm_config
        config_dict['Original dataset size'] = original_size
        config_dict['Used samples'] = len(dataset)
        if max_samples is not None:
            config_dict['Max samples'] = max_samples

    return loader




