import unittest

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')

        feature = 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]
        model = tfkit.classifier.MtClassifier({"taskA": ["a", "b"]}, tokenizer, pretrained)
        print(model(feature))
        self.assertTrue(isinstance(model(feature), Tensor))
        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)
        print(model.predict(task="taskA", input=input))
        print(model.predict(task="taskA", input=input, topk=2))

    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')

        feature = tfkit.qa.get_feature_from_data(tokenizer, input, target, maxlen=512)
        for k, v in feature.items():
            feature[k] = [v, v]
        model = tfkit.qa.QA(tokenizer, pretrained, maxlen=512)
        print(model(feature))
        self.assertTrue(isinstance(model(feature), Tensor))
        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)
        result, model_dict = model.predict(input=input)
        print("model_dict", model_dict, input, result)
        self.assertTrue('label_prob_all' in model_dict)
        self.assertTrue('label_map' in model_dict)
        self.assertTrue(len(result) == 1)
        result, model_dict = model.predict(input=input, topk=2)
        self.assertTrue('label_prob_all' in model_dict)
        self.assertTrue('label_map' in model_dict)
        self.assertTrue(len(result) == 2)

    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_tiny')
        pretrained = AutoModel.from_pretrained('voidful/albert_chinese_tiny')
        feature = 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]
        model = tfkit.tag.Tagger(["O", "A"],
                                 tokenizer, pretrained)

        self.assertTrue(isinstance(model(feature), Tensor))
        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)
        result, model_dict = model.predict(input=input)
        self.assertTrue('label_prob_all' in model_dict)
        self.assertTrue('label_map' in model_dict)
        print(result, len(result))
        self.assertTrue(isinstance(result, list))
        self.assertTrue(isinstance(result[0][0], str))

    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')

        feature = 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)

        print(model(feature))
        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)
        self.assertTrue('label_map' in model_dict)
        print(result, len(result))
        self.assertTrue(isinstance(result, list))
        self.assertTrue(isinstance(result[0][0], str))

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

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

        feature = 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]
        model = tfkit.gen_onebyone.OneByOne(tokenizer, pretrained)

        print(model(feature))
        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)
        self.assertTrue('label_map' in model_dict)
        print(result, len(result))
        self.assertTrue(isinstance(result, list))
        self.assertTrue(isinstance(result[0][0], str))
        result, model_dict = model.predict(input=input, beamsearch=True, beamsize=3)
        print("beamsaerch", result, len(result), model_dict)
