# 导入必备工具包
import torch
import numpy as np
from rich import print
from datasets import load_dataset
from transformers import AutoTokenizer
# import sys
# sys.path.append('..')
from bert_P_Tuning.ptune_config import *
from functools import partial


def convert_example(
        examples: dict,
        tokenizer,
        max_seq_len: int,
        max_label_len: int,
        p_embedding_num=6,
        train_mode=True,
        return_tensor=False
) -> dict:
    """
    将样本数据转换为模型接收的输入数据。

    Args:
        examples (dict): 训练数据样本, e.g. -> {
                                                "text": [
                                                            '娱乐	嗨放派怎么停播了',
                                                            '体育	世界杯为何迟迟不见宣传',
                                                            ...
                                                ]
                                            }
        max_label_len (int): 最大label长度，若没有达到最大长度，则padding为最大长度
        p_embedding_num (int): p-tuning token 的个数
        train_mode (bool): 训练阶段 or 推理阶段。
        return_tensor (bool): 是否返回tensor类型，如不是，则返回numpy类型。

    Returns:
        dict (str: np.array) -> tokenized_output = {
                            'input_ids': [[101, 3928, ...], ],
                            'token_type_ids': [[0, 0, ...], ],
                            'mask_positions': [[5, 6, ...], ],
                            'mask_labels': [[183, 234], [298, 322], ...]
                        }
    """
    tokenized_output = {
        'input_ids': [],
        'attention_mask': [],
        'mask_positions': [],  # 记录label的位置（即MASK Token的位置）
        'mask_labels': []  # 记录MASK Token的原始值（即Label值）
    }
    #print(f'examples--》{examples}')
    start_mask_position = 1  # 将 prompt token(s) 插在 [CLS] 之后
    for i, example in enumerate(examples['text']):
        try:
            if train_mode:
                label, content = example.strip().split('\t', 1)
            else:
                content = example.strip()

            encoded_inputs = tokenizer(
                text=content,
                truncation=True,
                max_length=max_seq_len,
                padding='max_length')
        except:
            continue

        input_ids = encoded_inputs['input_ids']
        #print(f'encoded_inputs-->{encoded_inputs}')
        #print(f'input_ids-->{input_ids}')

        mask_tokens = ['[MASK]'] * max_label_len  # 1.生成 MASK Tokens, 和label长度一致
        #print(f'mask_tokens-->{mask_tokens}')
        mask_ids = tokenizer.convert_tokens_to_ids(mask_tokens)  # token 转 id
        #print(f'mask_ids-->{mask_ids}')
        p_tokens = ["[unused{}]".format(i + 1) for i in range(p_embedding_num)]  # 2.构建 软模版 prompt token(s)
        #print(f'p_tokens-->{p_tokens}')

        p_tokens_ids = tokenizer.convert_tokens_to_ids(p_tokens)  # token 转 id
        #print(f'p_tokens_ids-->{p_tokens_ids}')

        input_ids = input_ids[
                        :max_seq_len - len(mask_ids) - len(p_tokens_ids)]  # 根据最大长度-p_token长度-label长度，裁剪content的长度
        #print(f'input_ids1-->{input_ids}')
        #print(len(input_ids))
        input_ids = input_ids[:start_mask_position] + p_tokens_ids+mask_ids + input_ids[
                                                                         # 3.插入 MASK -> [CLS][unused1][unused2]...[MASK][MASK]世界杯...[SEP]
                                                                         start_mask_position:]
        #print(f'input_ids2--{input_ids}')



        mask_positions = [len(p_tokens_ids) + start_mask_position + i for  # 将 Mask Tokens 的位置记录下来
                          i in range(max_label_len)]
        #print(f'mask_positions--{mask_positions}')
        tokenized_output['input_ids'].append(input_ids)
        # 如果输入需要token_type_ids，可以进行添加，

        if 'token_type_ids' in encoded_inputs:  # 兼容不需要 token_type_id 的模型, e.g. Roberta-Base
            tmp = encoded_inputs['token_type_ids']
            if 'token_type_ids' not in tokenized_output:
                tokenized_output['token_type_ids'] = [tmp]
            else:
                tokenized_output['token_type_ids'].append(tmp)

        #处理attention_mask
        attention_mask = np.where(np.array(input_ids) > 0, 1, 0)
        #print(f'attention_mask--{attention_mask}')
        tokenized_output['attention_mask'].append(attention_mask)
        tokenized_output['mask_positions'].append(mask_positions)

        if train_mode:
            mask_labels = tokenizer(text=label)  # label token 转 id
            mask_labels = mask_labels['input_ids'][1:-1]  # 丢掉[CLS]和[SEP]
            mask_labels = mask_labels[:max_label_len]
            mask_labels += [tokenizer.pad_token_id] * (max_label_len - len(mask_labels))  # 将 label 补到最长
            tokenized_output['mask_labels'].append(mask_labels)

    for k, v in tokenized_output.items():
        if return_tensor:
            tokenized_output[k] = torch.LongTensor(v)
        else:
            tokenized_output[k] = np.array(v)

    return tokenized_output





if __name__ == '__main__':
    pc = ProjectConfig()
    train_dataset = load_dataset('text', data_files={'train': pc.train_path})
    #print(f'train_dataset==>{train_dataset}')
    # print(train_dataset['train']['text'])
    tokenizer = AutoTokenizer.from_pretrained(pc.pre_model)

    examples = {"text": ['手机	这个手机也太卡了。']}
    tokenized_output = convert_example(examples=examples,
                                       tokenizer=tokenizer,
                                       max_seq_len=30,
                                       max_label_len=2,
                                       train_mode=True, return_tensor=False)

    print(f'tokenized_output-->{tokenized_output}')