# Copyright 2022 MosaicML Composer authors
# SPDX-License-Identifier: Apache-2.0

import logging
from typing import List, cast

from torch.utils.data import DataLoader, Dataset

from composer.utils import MissingConditionalImportError, dist

log = logging.getLogger(__name__)


def build_lm_dataloader(
    datadir: List[str],
    tokenizer_name: str,
    global_batch_size: int,
    *,
    split: str = 'train',
    shuffle: bool = True,
    drop_last: bool = True,
    use_masked_lm: bool = False,
    num_tokens: int = 0,
    mlm_probability: float = 0.15,
    subsample_ratio: float = 1.0,
    **dataloader_kwargs,
):
    """Builds a dataloader for a generic language modeling dataset.

    Args:
        datadir (list): List containing the string of the path to the HuggingFace
            Datasets directory.
        dataloader_hparams (DataLoaderHparams): DataLoaderHparams object.
        tokenizer_name (str): The name of the HuggingFace tokenizer to
            preprocess text with. See `HuggingFace documentation
            <https://huggingface.co/models>`_.
        global_batch_size (int): Global batch size.
        split (str): the dataset split to use either 'train', 'val', or 'test'. Default: ``'train```. Default: ``'train'``.
        shuffle (bool): whether to shuffle the dataset. Default: ``True``.
        drop_last (bool): whether to drop last samples. Default: ``True``.
        use_masked_lm (bool): Whether the dataset should be encoded with masked
            language modeling or not.
        num_tokens (int, optional): Number of tokens to train on. ``0``
            will train on all tokens in the dataset. Default: ``0``.
        mlm_probability (float, optional): If using masked language modeling, the
            probability with which tokens will be masked. Default: ``0.15``.
        subsample_ratio (float, optional): Proportion of the dataset to use. Default:
            ``1.0``.
        **dataloader_kwargs (Dict[str, Any]): Additional settings for the dataloader (e.g. num_workers, etc.)
    """
    if global_batch_size % dist.get_world_size() != 0:
        raise ValueError(
            f'global_batch_size ({global_batch_size}) must be divisible by world_size ({dist.get_world_size()}).')
    batch_size = global_batch_size // dist.get_world_size()
    try:
        import transformers
    except ImportError as e:
        raise MissingConditionalImportError(extra_deps_group='nlp', conda_package='transformers') from e

    try:
        import datasets
    except ImportError as e:
        raise MissingConditionalImportError(extra_deps_group='nlp', conda_package='datasets') from e

    assert tokenizer_name is not None

    tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_name)
    # loads a dataset that is assumed to be pre-tokenized
    lm_datasets = [datasets.load_from_disk(i) for i in datadir]  #type: ignore (thirdparty)

    # merge the dataset to re-sample from
    if split is None:
        raise ValueError('A dataset split is required')
    merged_dataset = [[d[split]] for d in lm_datasets]
    # flatten merged_dataset
    merged_dataset = [item for sublist in merged_dataset for item in sublist]
    lm_datasets = datasets.concatenate_datasets(merged_dataset)  #type: ignore (thirdparty)

    total_num_samples = len(lm_datasets)  # type: ignore
    tokens_per_sample = len(lm_datasets[0]['input_ids'])  #type: ignore (thirdparty)
    total_num_tokens = total_num_samples * tokens_per_sample

    # truncate the dataset to a specified size
    num_samples = total_num_samples
    if num_tokens > 0:
        assert num_tokens <= total_num_tokens, f'Requested {num_tokens} tokens must be <= total_num_tokens={total_num_tokens}'
        assert num_tokens % tokens_per_sample == 0, f'Requested {num_tokens} tokens is not divisible by tokens_per_sample={tokens_per_sample}'
        num_samples = num_tokens // tokens_per_sample
        subsample_ratio = num_samples / total_num_samples
    elif subsample_ratio < 1.0:
        num_samples = round(total_num_samples * subsample_ratio)
        num_tokens = num_samples * tokens_per_sample
    elif subsample_ratio == 1.0 and num_tokens == 0:
        num_tokens = total_num_tokens
    else:
        log.warning('No subsampling going on!')

    lm_datasets = lm_datasets.select(range(num_samples))  # type: ignore (thirdparty)
    log.info(f'LM datasets: {lm_datasets}')
    log.info(f'Subsample ratio: {subsample_ratio}')
    log.info(f'Total number of samples: {num_samples:e}')
    log.info(f'Total number of tokens: {num_tokens:e}')
    dataset = lm_datasets

    # for some tokenizers, e.g. GPT-2, they don't have padding tokens. Hence, we cannot use the LM collator.
    if tokenizer.pad_token_id is None:
        data_collator = transformers.default_data_collator
    else:
        data_collator = transformers.DataCollatorForLanguageModeling(tokenizer=tokenizer,
                                                                     mlm=use_masked_lm,
                                                                     mlm_probability=mlm_probability)

    sampler = dist.get_sampler(
        cast(Dataset, dataset),  # HF datasets do not subclass torch datasets, so this cast is needed
        drop_last=drop_last,
        shuffle=shuffle)

    return DataLoader(
        dataset=dataset,  # type: ignore
        batch_size=batch_size,
        sampler=sampler,
        drop_last=drop_last,
        collate_fn=data_collator,
        **dataloader_kwargs)
