# Copyright (c) OpenMMLab. All rights reserved.
import argparse
from functools import partial

import numpy as np
from datasets import DatasetDict
from mmengine.config import Config

from xtuner.dataset.utils import Packer, encode_fn
from xtuner.registry import BUILDER


def parse_args():
    parser = argparse.ArgumentParser(
        description='Verify the correctness of the config file for the '
        'custom dataset.')
    parser.add_argument('config', help='config file name or path.')
    args = parser.parse_args()
    return args


def is_standard_format(dataset):
    example = next(iter(dataset))
    if 'conversation' not in example:
        return False
    conversation = example['conversation']
    if not isinstance(conversation, list):
        return False
    for item in conversation:
        if (not isinstance(item, dict)) or ('input'
                                            not in item) or ('output'
                                                             not in item):
            return False
        input, output = item['input'], item['output']
        if (not isinstance(input, str)) or (not isinstance(output, str)):
            return False
    return True


def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)

    tokenizer = BUILDER.build(cfg.tokenizer)
    if cfg.get('framework', 'mmengine').lower() == 'huggingface':
        train_dataset = cfg.train_dataset
    else:
        train_dataset = cfg.train_dataloader.dataset

    dataset = train_dataset.dataset
    max_length = train_dataset.max_length
    dataset_map_fn = train_dataset.get('dataset_map_fn', None)
    template_map_fn = train_dataset.get('template_map_fn', None)
    max_dataset_length = train_dataset.get('max_dataset_length', 10)
    split = train_dataset.get('split', 'train')
    remove_unused_columns = train_dataset.get('remove_unused_columns', False)
    rename_maps = train_dataset.get('rename_maps', [])
    shuffle_before_pack = train_dataset.get('shuffle_before_pack', True)
    pack_to_max_length = train_dataset.get('pack_to_max_length', True)
    input_ids_with_output = train_dataset.get('input_ids_with_output', True)

    if dataset.get('path', '') != 'json':
        raise ValueError(
            'You are using custom datasets for SFT. '
            'The custom datasets should be in json format. To load your JSON '
            'file, you can use the following code snippet: \n'
            '"""\nfrom datasets import load_dataset \n'
            'dataset = dict(type=load_dataset, path=\'json\', '
            'data_files=\'your_json_file.json\')\n"""\n'
            'For more details, please refer to Step 5 in the '
            '`Using Custom Datasets` section of the documentation found at'
            ' docs/zh_cn/user_guides/single_turn_conversation.md.')

    try:
        dataset = BUILDER.build(dataset)
    except RuntimeError:
        raise RuntimeError(
            'Unable to load the custom JSON file using '
            '`datasets.load_dataset`. Your data-related config is '
            f'{train_dataset}. Please refer to the official documentation on'
            ' `load_dataset` (https://huggingface.co/docs/datasets/loading) '
            'for more details.')

    if isinstance(dataset, DatasetDict):
        dataset = dataset[split]

    if not is_standard_format(dataset) and dataset_map_fn is None:
        raise ValueError(
            'If the custom dataset is not in the XTuner-defined '
            'format, please utilize `dataset_map_fn` to map the original data'
            ' to the standard format. For more details, please refer to '
            'Step 1 and Step 5 in the `Using Custom Datasets` section of the '
            'documentation found at '
            '`docs/zh_cn/user_guides/single_turn_conversation.md`.')

    if is_standard_format(dataset) and dataset_map_fn is not None:
        raise ValueError(
            'If the custom dataset is already in the XTuner-defined format, '
            'please set `dataset_map_fn` to None.'
            'For more details, please refer to Step 1 and Step 5 in the '
            '`Using Custom Datasets` section of the documentation found at'
            ' docs/zh_cn/user_guides/single_turn_conversation.md.')

    max_dataset_length = min(max_dataset_length, len(dataset))
    indices = np.random.choice(len(dataset), max_dataset_length, replace=False)
    dataset = dataset.select(indices)

    if dataset_map_fn is not None:
        dataset = dataset.map(dataset_map_fn)

    print('#' * 20 + '   dataset after `dataset_map_fn`   ' + '#' * 20)
    print(dataset[0]['conversation'])

    if template_map_fn is not None:
        template_map_fn = BUILDER.build(template_map_fn)
        dataset = dataset.map(template_map_fn)

    print('#' * 20 + '   dataset after adding templates   ' + '#' * 20)
    print(dataset[0]['conversation'])

    for old, new in rename_maps:
        dataset = dataset.rename_column(old, new)

    if pack_to_max_length and (not remove_unused_columns):
        raise ValueError('We have to remove unused columns if '
                         '`pack_to_max_length` is set to True.')

    dataset = dataset.map(
        partial(
            encode_fn,
            tokenizer=tokenizer,
            max_length=max_length,
            input_ids_with_output=input_ids_with_output),
        remove_columns=list(dataset.column_names)
        if remove_unused_columns else None)

    print('#' * 20 + '   encoded input_ids   ' + '#' * 20)
    print(dataset[0]['input_ids'])
    print('#' * 20 + '   encoded labels    ' + '#' * 20)
    print(dataset[0]['labels'])

    if pack_to_max_length and split == 'train':
        if shuffle_before_pack:
            dataset = dataset.shuffle()
            dataset = dataset.flatten_indices()
        dataset = dataset.map(Packer(max_length), batched=True)

        print('#' * 20 + '   input_ids after packed to max_length   ' +
              '#' * 20)
        print(dataset[0]['input_ids'])
        print('#' * 20 + '   labels after packed to max_length    ' + '#' * 20)
        print(dataset[0]['labels'])


if __name__ == '__main__':
    main()