import os
import unittest

import torch
from torch import Tensor
from transformers import BertTokenizer, AutoModel

import tfkit


class TestModel(unittest.TestCase):

    def testClassifier(self):
        input = "One hundred thirty-four patients suspected of having pancreas cancer successfully underwent gray scale ultrasound examination of the pancreas ."
        target = "a"
        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
        pretrained = AutoModel.from_pretrained('voidful/albert_chinese_tiny')

        model = tfkit.classifier.MtClassifier({"taskA": ["a", "b"]}, tokenizer, pretrained)
        for feature in tfkit.classifier.get_feature_from_data(tokenizer, task_lables={"taskA": ["a", "b"]},
                                                              task="taskA",
                                                              input=input, target=target, maxlen=512):
            for k, v in feature.items():
                feature[k] = [v, v]

            print(feature)
            # test train
            print(model(feature))
            self.assertTrue(isinstance(model(feature), Tensor))
            # test eval
            print(model(feature, eval=True))
            model_dict = model(feature, eval=True)
            self.assertTrue('label_prob_all' in model_dict)
            self.assertTrue('label_map' in model_dict)
            # test predict
            tok_label = model.predict(task="taskA", input=input)
            self.assertTrue(len(tok_label) == 2)
            # test predict with top k 2
            top_k_label, top_k_dict = model.predict(task="taskA", input=input, topk=2)
            print("test predict with top k 2, ", top_k_label, top_k_dict)
            self.assertTrue(len(top_k_label) == 2)

        # test exceed 512
        for merge_strategy in ['minentropy', 'maxcount', 'maxprob']:
            result, model_dict = model.predict(task="taskA", input=" ".join([str(i) for i in range(2000)]),
                                               merge_strategy=merge_strategy)
            print(result, len(model_dict), model_dict)
            self.assertTrue(isinstance(result, list))
            self.assertTrue(len(result) == 1)

    def testQA(self):
        input = "梵 語 在 社 交 中 口 頭 使 用 , 並 且 在 早 期 古 典 梵 語 文 獻 的 發 展 中 維 持 口 頭 傳 統 。 在 印 度 , 書 寫 形 式 是 當 梵 語 發 展 成 俗 語 之 後 才 出 現 的 ; 在 書 寫 梵 語 的 時 候 , 書 寫 系 統 的 選 擇 受 抄 寫 者 所 處 地 域 的 影 響 。 同 樣 的 , 所 有 南 亞 的 主 要 書 寫 系 統 事 實 上 都 用 於 梵 語 文 稿 的 抄 寫 。 自 1 9 世 紀 晚 期 , 天 城 文 被 定 為 梵 語 的 標 準 書 寫 系 統 , 十 分 可 能 的 原 因 是 歐 洲 人 有 用 這 種 文 字 印 刷 梵 語 文 本 的 習 慣 。 最 早 的 已 知 梵 語 碑 刻 可 確 定 為 公 元 前 一 世 紀 。 它 們 採 用 了 最 初 用 於 俗 語 而 非 梵 語 的 婆 羅 米 文 。 第 一 個 書 寫 梵 語 的 證 據 , 出 現 在 晚 於 它 的 俗 語 的 書 寫 證 據 之 後 的 幾 個 世 紀 , 這 被 描 述 為 一 種 悖 論 。 在 梵 語 被 書 寫 下 來 的 時 候 , 它 首 先 用 於 行 政 、 文 學 或 科 學 類 的 文 本 。 宗 教 文 本 口 頭 傳 承 , 在 相 當 晚 的 時 候 才 「 不 情 願 」 地 被 書 寫 下 來 。 [Question] 最 初 梵 語 以 什 麼 書 寫 系 統 被 記 錄 下 來 ?"
        target = [201, 205]
        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
        pretrained = AutoModel.from_pretrained('voidful/albert_chinese_tiny')
        model = tfkit.qa.QA(tokenizer, pretrained, maxlen=512)

        for feature in tfkit.qa.get_feature_from_data(tokenizer, input, target, maxlen=512):
            for k, v in feature.items():
                feature[k] = [v, v]

            # test train
            print(model(feature))
            self.assertTrue(isinstance(model(feature), Tensor))
            # test eval
            print(model(feature, eval=True))
            model_dict = model(feature, eval=True)
            self.assertTrue('label_prob_all' in model_dict)
            self.assertTrue('label_map' in model_dict)

        # test predict
        result, model_dict = model.predict(input=input)
        print("model_dict", model_dict, input, result)
        self.assertTrue('label_prob_all' in model_dict[0])
        self.assertTrue('label_map' in model_dict[0])
        self.assertTrue(len(result) == 1)

        # test eval top k = 2
        top_k_label, top_k_dict = model.predict(input=input, topk=10)
        print("top_k_label", top_k_label)
        self.assertTrue(len(top_k_label) == 10)

        # test exceed 512
        for merge_strategy in ['minentropy', 'maxcount', 'maxprob']:
            result, model_dict = model.predict(input=" ".join([str(i) for i in range(550)]),
                                               merge_strategy=merge_strategy)
            print(result, len(model_dict))
            self.assertTrue(isinstance(result, list))
            self.assertTrue(len(result) == 1)

    def testTag(self):
        input = "在 歐 洲 , 梵 語 的 學 術 研 究 , 由 德 國 學 者 陸 特 和 漢 斯 雷 頓 開 創 。 後 來 威 廉 · 瓊 斯 發 現 印 歐 語 系 , 也 要 歸 功 於 對 梵 語 的 研 究 。 此 外 , 梵 語 研 究 , 也 對 西 方 文 字 學 及 歷 史 語 言 學 的 發 展 , 貢 獻 不 少 。 1 7 8 6 年 2 月 2 日 , 亞 洲 協 會 在 加 爾 各 答 舉 行 。 [SEP] 陸 特 和 漢 斯 雷 頓 開 創 了 哪 一 地 區 對 梵 語 的 學 術 研 究 ?"
        target = "O A A O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O"
        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_small')
        pretrained = AutoModel.from_pretrained('voidful/albert_chinese_small')

        model = tfkit.tag.Tagger(['O', 'A'],
                                 tokenizer, pretrained)

        for feature in tfkit.tag.get_feature_from_data(tokenizer, labels=["O", "A"], input=input, target=target,
                                                       maxlen=512):
            for k, v in feature.items():
                feature[k] = [v, v]
            print(feature)

            # test train
            print(model(feature))
            self.assertTrue(isinstance(model(feature), Tensor))
            # test eval
            model_dict = model(feature, eval=True)
            self.assertTrue('label_prob_all' in model_dict)
            self.assertTrue('label_map' in model_dict)

        # test predict
        result, model_dict = model.predict(input=input, start_contain="A", end_contain="A")
        print("model_dict", model_dict)
        self.assertTrue('label_prob_all' in model_dict[0])
        self.assertTrue('label_map' in model_dict[0])
        print("result", result, len(result))
        self.assertTrue(isinstance(result, list))

        # test exceed 512
        for merge_strategy in ['minentropy', 'maxcount', 'maxprob']:
            result, model_dict = model.predict(input=" ".join([str(i) for i in range(1000)]),
                                               merge_strategy=merge_strategy, start_contain="A", end_contain="A")
            print(result)
            self.assertTrue(isinstance(result, list))

    def testMask(self):
        input = "今 天 [MASK] 情 [MASK] 好"
        target = "心 很"

        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_small')
        pretrained = AutoModel.from_pretrained('voidful/albert_chinese_small')
        model = tfkit.gen_mask.Mask(tokenizer, pretrained)

        for feature in tfkit.gen_mask.get_feature_from_data(tokenizer, input=input, target=target, maxlen=512):
            for k, v in feature.items():
                feature[k] = [v]

            print(feature)
            self.assertTrue(isinstance(model(feature), Tensor))

            model_dict = model(feature, eval=True)
            print(model_dict)
            self.assertTrue('label_map' in model_dict)
            self.assertTrue('label_prob' in model_dict)

        result, model_dict = model.predict(input=input)
        self.assertTrue('label_prob' in model_dict[0])
        self.assertTrue('label_map' in model_dict[0])
        print("predict", result, len(result))

        self.assertTrue(isinstance(result, list))
        self.assertTrue(isinstance(result[0][0], str))

        # test exceed 512
        result, model_dict = model.predict(input="T " * 512)
        self.assertTrue(isinstance(result, list))
        self.assertTrue(len(result[0][0]) == 0)

    def testMCQ(self):
        input = "你 是 誰 [SEP] [MASK] 我 [MASK] 你 [MASK] 他"
        target = 1

        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_small')
        pretrained = AutoModel.from_pretrained('voidful/albert_chinese_small')
        model = tfkit.mcq.MCQ(tokenizer, pretrained)

        for feature in tfkit.mcq.get_feature_from_data(tokenizer, input=input, target=target, maxlen=512):
            for k, v in feature.items():
                feature[k] = [v]

            print(feature)
            self.assertTrue(isinstance(model(feature), Tensor))

            model_dict = model(feature, eval=True)
            print(model_dict)
            self.assertTrue('label_map' in model_dict)
            self.assertTrue('label_max' in model_dict)

        result, model_dict = model.predict(input=input)
        self.assertTrue('label_max' in model_dict[0])
        self.assertTrue('label_map' in model_dict[0])
        print("predict", result, len(result))

        self.assertTrue(isinstance(result, list))
        print(result)
        self.assertTrue(isinstance(result[0], str))

        # test exceed 512
        result, model_dict = model.predict(input="T " * 300 + "[MASK]" + "T " * 300)
        self.assertTrue(isinstance(result, list))

    def testOnce(self):
        input = "See you next time"
        target = "下 次 見"
        ntarget = "不 見 不 散"

        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
        pretrained = AutoModel.from_pretrained('voidful/albert_chinese_tiny')

        for feature in tfkit.gen_once.get_feature_from_data(tokenizer, input=input, target=target, maxlen=512):
            for k, v in feature.items():
                feature[k] = [v, v]
            model = tfkit.gen_once.Once(tokenizer, pretrained)
            self.assertTrue(isinstance(model(feature), Tensor))
            model_dict = model(feature, eval=True)
            self.assertTrue('label_prob_all' in model_dict)
            self.assertTrue('label_map' in model_dict)

        result, model_dict = model.predict(input=input)
        self.assertTrue('label_prob_all' in model_dict[0])
        self.assertTrue('label_map' in model_dict[0])
        print(result, len(result))
        self.assertTrue(isinstance(result, list))
        self.assertTrue(isinstance(result[0][0], str))
        # test exceed 512
        result, model_dict = model.predict(input="T " * 512)
        self.assertTrue(isinstance(result, list))

    def testOnebyone(self):
        input = "See you next time"
        previous = "下 次"
        target = "下 次 見"
        ntarget = "不 見 不 散"

        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_small')
        pretrained = AutoModel.from_pretrained('voidful/albert_chinese_small')

        model = tfkit.gen_onebyone.OneByOne(tokenizer, pretrained)
        # package = torch.load('./cache/model.pt', map_location='cpu')
        # for model_tag, state_dict in zip(package['tags'], package['models']):
        #     model.load_state_dict(state_dict)

        # test filter sim
        dummy_result = [[['表', '示', '事', '業'], 4.438073633101437],
                        [['表', '示', '事', '情'], 9.86092332722302],
                        [['表', '示', '事', '情'], 9.86092332722302]]
        model._filterSimilar(dummy_result, 3)
        self.assertTrue(len(dummy_result), 3)

        dummy_result = [[['表', '示', '事', '業'], 4.438073633101437],
                        [['表', '示', '事', '情'], 9.86092332722302]]
        model._filterSimilar(dummy_result, 2)
        self.assertTrue(len(dummy_result), 2)

        for feature in tfkit.gen_onebyone.get_feature_from_data(tokenizer, input=input,
                                                                tokenized_previous=tokenizer.tokenize(
                                                                    " ".join(previous)),
                                                                tokenized_target=tokenizer.tokenize(" ".join(target)),
                                                                maxlen=512):
            for k, v in feature.items():
                feature[k] = [v, v]

            print(model(feature))
            self.assertTrue(isinstance(model(feature), Tensor))
            model_dict = model(feature, eval=True)
            self.assertTrue('label_map' in model_dict)

        # greedy
        result, model_dict = model.predict(input=input)
        print(result, model_dict)
        self.assertTrue('label_map' in model_dict[0])
        self.assertTrue(len(result) == 1)
        self.assertTrue(isinstance(result, list))
        self.assertTrue(isinstance(result[0][0], str))

        # TopK
        result, model_dict = model.predict(input=input, decodenum=3, mode='topK', topK=3, filtersim=False)
        print("TopK no filter sim", result, len(result), model_dict)
        self.assertTrue('label_map' in model_dict[0])
        self.assertTrue(len(result) == 3)
        self.assertTrue(isinstance(result, list))
        self.assertTrue(isinstance(result[0][0], str))

        # beamsearch
        result, model_dict = model.predict(input=input, decodenum=3)
        print("beamsaerch", result, len(result), model_dict)
        self.assertTrue('label_map' in model_dict[0])
        self.assertTrue(len(result) == 3)
        self.assertTrue(isinstance(result, list))
        self.assertTrue(isinstance(result[0][0], str))

        # TopK
        result, model_dict = model.predict(input=input, decodenum=3, mode='topK', topK=20)
        print("TopK", result, len(result), model_dict)
        self.assertTrue('label_map' in model_dict[0])
        self.assertTrue(len(result) == 3)
        self.assertTrue(isinstance(result, list))
        self.assertTrue(isinstance(result[0][0], str))

        # TopP
        result, model_dict = model.predict(input=input, decodenum=3, mode='topP', topP=0.8)
        print("TopP", result, len(result), model_dict)
        self.assertTrue('label_map' in model_dict[0])
        self.assertTrue(len(result) == 3)
        self.assertTrue(isinstance(result, list))
        self.assertTrue(isinstance(result[0][0], str))

        # test exceed 512
        result, model_dict = model.predict(input="T " * 512)
        self.assertTrue(isinstance(result, list))
        print("exceed max len", result)
        self.assertTrue(len(result[0]) == 0)
        result, model_dict = model.predict(input="T " * 550, reserved_len=10)
        self.assertTrue(isinstance(result, list))
        print("exceed max len with reserved len:", result)
        self.assertTrue(result)

    def testOnebyoneWithReservedLen(self):
        ROOT_DIR = os.path.dirname(os.path.abspath(__file__ + "/../"))
        DATASET_DIR = os.path.join(ROOT_DIR, 'demo_data')

        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')

        for i in tfkit.gen_onebyone.get_data_from_file(os.path.join(DATASET_DIR, 'generate.csv')):
            tasks, task, input, target, negative_text = i
            input = input.strip()
            tokenized_target = tokenizer.tokenize(" ".join(target))
            for j in range(1, len(tokenized_target) + 1):
                feature = tfkit.gen_onebyone.get_feature_from_data(tokenizer, input=input,
                                                                   tokenized_previous=tokenized_target[:j - 1],
                                                                   tokenized_target=tokenized_target[:j],
                                                                   maxlen=20, reserved_len=0)[-1]
                target_start = feature['start']
                print(f"input: {len(feature['input'])}, {tokenizer.decode(feature['input'][:target_start])} ")
                print(f"type: {len(feature['type'])}, {feature['type'][:target_start]} ")
                print(f"mask: {len(feature['mask'])}, {feature['mask'][:target_start]} ")
                if tokenized_target is not None:
                    print(
                        f"target: {len(feature['target'])}, {tokenizer.convert_ids_to_tokens(feature['target'][target_start])} ")
