from pydantic import BaseModel
import os


class AdditionalSpecialTokens(BaseModel):
    # word embedding
    space_token = "[Space]"

    # task token
    tag_task_token = "[TASK:TAG]"

    # segment embedding
    comment_token = "[Comment]"
    tag_token = '[Tag]'
    question_token = '[Question]'
    answer_token = '[Answer]'
    content_token = '[Content]'
    img_token = "[IMG]"
    task_token = '[TASK]'
    speaker1_token = '[Speaker1]'
    speaker2_token = '[Speaker2]'
    comment1_token = "[Comment1]"
    comment2_token = "[Comment2]"
    comment3_token = "[Comment3]"
    comment4_token = "[Comment4]"
    comment5_token = "[Comment5]"
    comment6_token = "[Comment6]"
    comment7_token = "[Comment7]"
    comment8_token = "[Comment8]"
    comment9_token = "[Comment9]"
    comment10_token = "[Comment10]"
    article_start_token = "[ArticleStart]"
    article_end_token = "[ArticleEnd]"
    article_split_token = "[ArticleSplit]"
    chat_history_start_token = "[ChatHistoryStart]"
    chat_history_end_token = "[ChatHistoryEnd]"
    sex_male_token = "[SEX_MALE]"
    sex_female_token = "[SEX_FEMALE]"
    sex_unk_token = "[SEX_UNK]"


class TokenizerParams:
    def __init__(self):
        self.article_start = "[ArticleStart]"
        self.article_end = "[ArticleEnd]"
        self.article_split = "[ArticleSplit]"

        from transformers import BertTokenizer
        self.vocab_path = os.getcwd() + '/nlp_utils/vocab.txt'

        self.tokenizer = BertTokenizer(vocab_file=self.vocab_path, do_lower_case=True)

        self.special_tokens_obj = AdditionalSpecialTokens()
        self.special_tokens = self.special_tokens_obj.__dict__.values()
        self.tokenizer.add_special_tokens({'additional_special_tokens': [x for x in self.special_tokens]})

        # Special tokens

        self.space_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.space_token)
        self.tag_task_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.tag_task_token)
        self.unk_id = self.tokenizer.convert_tokens_to_ids("[UNK]")
        self.sep_id = self.tokenizer.convert_tokens_to_ids("[SEP]")
        self.article_start_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.article_start_token)
        self.article_split_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.article_split_token)
        self.article_end_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.article_end_token)
        self.comment_seg_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.comment_token)
        self.tag_seg_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.tag_token)
        self.question_seg_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.question_token)
        self.answer_seg_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.answer_token)
        self.content_seg_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.content_token)
        self.img_seg_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.img_token)
        self.task_seg_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.task_token)
        self.speaker1_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.speaker1_token)
        self.speaker2_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.speaker2_token)
        self.chat_history_start_id = self.tokenizer.convert_tokens_to_ids(
            self.special_tokens_obj.chat_history_start_token)
        self.chat_history_end_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.chat_history_end_token)
        self.sex_male_token_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.sex_male_token)
        self.sex_female_token_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.sex_female_token)
        self.sex_unk_token_id = self.tokenizer.convert_tokens_to_ids(self.special_tokens_obj.sex_unk_token)
        self.sex_chrs = {"1": "[SEX_MALE]", "0": "[SEX_FEMALE]", "None": "[SEX_UNK]"}
