import string

from datasets import load_dataset
from tokenizers import ByteLevelBPETokenizer
from transformers import PreTrainedTokenizerFast


# dataset_0 = (
#     load_dataset('wikimedia/wikisource', lang, split='train')
#     for lang in ['20231201.ar', '20231201.as', '20231201.az', '20231201.ban', '20231201.be', '20231201.bg', '20231201.bn', '20231201.br', '20231201.bs', '20231201.ca', '20231201.cs', '20231201.cy', '20231201.da', '20231201.de', '20231201.el', '20231201.en', '20231201.eo', '20231201.es', '20231201.et', '20231201.eu', '20231201.fa', '20231201.fi', '20231201.fo', '20231201.fr', '20231201.gl', '20231201.gu', '20231201.he', '20231201.hi', '20231201.hr', '20231201.hu', '20231201.hy', '20231201.id', '20231201.is', '20231201.it', '20231201.ja', '20231201.jv', '20231201.kn', '20231201.ko', '20231201.la', '20231201.li', '20231201.lij', '20231201.lt', '20231201.mk', '20231201.ml', '20231201.mr', '20231201.nap', '20231201.nl', '20231201.no', '20231201.or', '20231201.pa', '20231201.pl', '20231201.pms', '20231201.pt', '20231201.ro', '20231201.ru', '20231201.sa', '20231201.sah', '20231201.sk', '20231201.sl', '20231201.sr', '20231201.su', '20231201.sv', '20231201.ta', '20231201.te', '20231201.th', '20231201.tr', '20231201.uk', '20231201.vec', '20231201.vi', '20231201.wa', '20231201.yi', '20231201.zh', '20231201.zh-min-nan']
# )

dataset_1 = (
    load_dataset('xu-song/cc100-samples', lang, split='train')
    for lang in ['am', 'ar', 'as', 'az', 'be', 'bg', 'bn', 'bn_rom', 'br', 'bs', 'ca', 'cs', 'cy', 'da', 'de', 'el', 'en', 'eo', 'es', 'et', 'eu', 'fa', 'ff', 'fi', 'fr', 'fy', 'ga', 'gd', 'gl', 'gn', 'gu', 'ha', 'he', 'hi', 'hi_rom', 'hr', 'ht', 'hu', 'hy', 'id', 'ig', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km', 'kn', 'ko', 'ku', 'ky', 'la', 'lg', 'li', 'ln', 'lo', 'lt', 'lv', 'mg', 'mk', 'ml', 'mn', 'mr', 'ms', 'my', 'my_zaw', 'ne', 'nl', 'no', 'ns', 'om', 'or', 'pa', 'pl', 'ps', 'pt', 'qu', 'rm', 'ro', 'ru', 'sa', 'si', 'sc', 'sd', 'sk', 'sl', 'so', 'sq', 'sr', 'ss', 'su', 'sv', 'sw', 'ta', 'ta_rom', 'te', 'te_rom', 'th', 'tl', 'tn', 'tr', 'ug', 'uk', 'ur', 'ur_rom', 'uz', 'vi', 'wo', 'xh', 'yi', 'yo', 'zh-Hans', 'zh-Hant', 'zu']
)

dataset_2 = (
    load_dataset('csebuetnlp/xlsum', lang, split='train')
    for lang in ['amharic', 'arabic', 'azerbaijani', 'bengali', 'burmese', 'chinese_simplified', 'chinese_traditional', 'english', 'french', 'gujarati', 'hausa', 'hindi', 'igbo', 'indonesian', 'japanese', 'kirundi', 'korean', 'kyrgyz', 'marathi', 'nepali', 'oromo', 'pashto', 'persian', 'pidgin', 'portuguese', 'punjabi', 'russian', 'scottish_gaelic', 'serbian_cyrillic', 'serbian_latin', 'sinhala', 'somali', 'spanish', 'swahili', 'tamil', 'telugu', 'thai', 'tigrinya', 'turkish', 'ukrainian', 'urdu', 'uzbek', 'vietnamese', 'welsh', 'yoruba']
)

# dataset_3 = load_dataset('recursal/SuperWikiNEXT-32B', split='train')
dataset_4 = load_dataset('m-a-p/CodeFeedback-Filtered-Instruction', split='train')
dataset_5 = load_dataset('nampdn-ai/tiny-codes', split='train')
# dataset_6 = load_dataset('ajibawa-2023/Maths-College', split='train')
dataset_7 = load_dataset('microsoft/orca-math-word-problems-200k', split='train')
dataset_8 = load_dataset('mlabonne/FineTome-100k', split='train')
dataset_9 = load_dataset('arcee-ai/agent-data', split='train')
dataset_10 = [
    load_dataset('cognitivecomputations/SystemChat-2.0', data_files='SystemChat_filtered.jsonl', split='train'),
    load_dataset('cognitivecomputations/SystemChat-2.0', data_files='SystemChat_multilingual.jsonl', split='train'),
]
dataset_11 = load_dataset('badrex/llm-emoji-dataset', split='train')


def batch_iterator():
    # for d in dataset_0:
    #     for row in d['text']:
    #         yield row
    #         break
    #
    #     break

    for d in dataset_1:
        for row in d['text']:
            yield row
            # break

        # break

    for d in dataset_2:
        for row in d['text']:
            yield row
            # break

        # break

    # for row in dataset_3['text']:
    #     yield row
    #     break

    for row in dataset_4:
        yield row['query'] + '\n' + row['answer']
        # break

    for row in dataset_5:
        yield row['prompt'] + '\n' + row['response']
        # break

    # for row in dataset_6:
    #     yield row['instruction'] + '\n' + row['output']
    #     break

    for row in dataset_7:
        yield row['question'] + '\n' + row['answer']
        # break

    for row in dataset_8['conversations']:
        yield '\n'.join(n['value'] for n in row)
        # break

    for row in dataset_9['conversations']:
        yield '\n'.join(n['value'] for n in row)
        # break

    for d in dataset_10:
        for row in d['messages']:
            yield '\n'.join(n['content'] for n in row)
            # break

    for row in dataset_11:
        yield f'{row["character"]}\n{row["unicode"]}\n{row["short description"]}\n{row["tags"]}\n{row["LLM description"]}'
        # break


# for row in batch_iterator():
#     print(f'{row = }')


special_tokens = [
    '<s>',
    '</s>',
    '<pad>',
    '<unk>',
    '<mask>',
    '<|im_start|>',
    '<|im_end|>',
    '<tools>',
    '</tools>',
    '<tool_call>',
    '</tool_call>',
    '<tool_response>',
    '</tool_response>',
    'system',
    'user',
    'assistant',
    *list(string.printable),
]

for i in range(64 - len(special_tokens)):
    special_tokens.append(f'<|reserved_{i}|>')

ascii_chars = string.ascii_letters + string.ascii_lowercase + string.ascii_uppercase + string.digits + string.punctuation

tokenizer = ByteLevelBPETokenizer()

tokenizer.train_from_iterator(
    [ascii_chars],
    vocab_size=len(ascii_chars),
    min_frequency=1,
    special_tokens=[],
)

tokenizer.train_from_iterator(
    batch_iterator(),
    vocab_size=32064,
    min_frequency=2,
    special_tokens=special_tokens,
)

tokenizer.save_model('..')

CHATML_CHAT_TEMPLATE = (
    "{% for message in messages %}"
        "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}"
    "{% endfor %}"
    "{% if add_generation_prompt %}"
        "{{ '<|im_start|>assistant\n' }}"
    "{% endif %}"
)

fast_tokenizer = PreTrainedTokenizerFast(
    tokenizer_object=tokenizer,
    chat_template=CHATML_CHAT_TEMPLATE,
    bos_token='<s>',
    eos_token='</s>',
    unk_token='<unk>',
    pad_token='<pad>',
    mask_token='<mask>',
)

fast_tokenizer.save_pretrained('..')