from torch.utils.data import DataLoader
from data_loaders.tensors import collate as all_collate
from data_loaders.tensors import t2m_collate, t2m_prefix_collate

def get_dataset_class(name):
    if name == "amass":
        from .amass import AMASS
        return AMASS
    elif name == "uestc":
        from .a2m.uestc import UESTC
        return UESTC
    elif name == "humanact12":
        from .a2m.humanact12poses import HumanAct12Poses
        return HumanAct12Poses
    elif name == "humanml":
        from data_loaders.humanml.data.dataset import HumanML3D
        return HumanML3D
    elif name == "kit":
        from data_loaders.humanml.data.dataset import KIT
        return KIT
    elif name == "preprocessed_posterior":
        from data_loaders.preprocessed_posterior_loader import PreprocessedPosteriorDataset
        return PreprocessedPosteriorDataset
    else:
        raise ValueError(f'Unsupported dataset name [{name}]')

def get_collate_fn(name, hml_mode='train', pred_len=0, batch_size=1):
    if hml_mode == 'gt':
        from data_loaders.humanml.data.dataset import collate_fn as t2m_eval_collate
        return t2m_eval_collate
    if name in ["humanml", "kit"]:
        if pred_len > 0:
            return lambda x: t2m_prefix_collate(x, pred_len=pred_len)
        return lambda x: t2m_collate(x, batch_size)
    elif name == "preprocessed_posterior":
        from data_loaders.preprocessed_posterior_loader import collate_preprocessed_posterior
        return collate_preprocessed_posterior
    else:
        return all_collate


def get_dataset(name, num_frames, split='train', hml_mode='train', abs_path='.', fixed_len=0, 
                device=None, autoregressive=False, cache_path=None, posterior_dir=None, 
                max_samples=None, resample=False): 
    DATA = get_dataset_class(name)
    if name in ["humanml", "kit"]:
        dataset = DATA(split=split, num_frames=num_frames, mode=hml_mode, abs_path=abs_path, fixed_len=fixed_len, 
                       device=device, autoregressive=autoregressive)
    elif name == "preprocessed_posterior":
        dataset = DATA(posterior_dir=posterior_dir, max_samples=max_samples, resample=resample)
    else:
        dataset = DATA(split=split, num_frames=num_frames)
    return dataset


def get_dataset_loader(name, batch_size, num_frames, split='train', hml_mode='train', fixed_len=0, pred_len=0, 
                       device=None, autoregressive=False, posterior_dir=None, max_samples=None, resample=False,
                       num_workers=8, shuffle=True):
    """
    Get a data loader.
    
    Args:
        name: dataset name; supports 'humanml', 'kit', 'preprocessed_posterior', etc.
        batch_size: batch size
        num_frames: number of frames
        split: dataset split (train/val/test)
        hml_mode: HumanML mode
        fixed_len: fixed length
        pred_len: prediction length
        device: device
        autoregressive: whether to use autoregressive mode
        posterior_dir: directory path for the preprocessed_posterior dataset
        max_samples: maximum number of samples (for preprocessed_posterior)
        resample: whether to resample (for preprocessed_posterior)
        num_workers: number of data loading workers
        shuffle: whether to shuffle data
    
    Returns:
        DataLoader object
    """
    dataset = get_dataset(name, num_frames, split=split, hml_mode=hml_mode, fixed_len=fixed_len, 
                device=device, autoregressive=autoregressive, posterior_dir=posterior_dir,
                max_samples=max_samples, resample=resample)
    
    collate = get_collate_fn(name, hml_mode, pred_len, batch_size)

    loader = DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle,
        num_workers=num_workers, drop_last=True, collate_fn=collate
    )

    return loader


def get_preprocessed_posterior_loader(posterior_dir, batch_size=32, shuffle=True,
                                       num_workers=8, max_samples=None, resample=False,
                                       mean=None, std=None, drop_last=True):
    """
    Convenience function for creating a preprocessed posterior data loader.
    
    Args:
        posterior_dir: directory containing preprocessed posterior files
        batch_size: batch size
        shuffle: whether to shuffle
        num_workers: number of data loading workers
        max_samples: maximum number of samples
        resample: whether to resample the latent while loading
        mean: mean used for normalization
        std: standard deviation used for normalization
        drop_last: whether to drop the last incomplete batch
    
    Returns:
        DataLoader object with a dataset attribute
    """
    from data_loaders.preprocessed_posterior_loader import get_preprocessed_posterior_loader as _get_loader
    return _get_loader(
        posterior_dir=posterior_dir,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        max_samples=max_samples,
        resample=resample,
        mean=mean,
        std=std,
        drop_last=drop_last
    )