# 可视化的模型对比测试
import re
import time
from pprint import pprint

from paddlenlp import Taskflow
from paddlenlp.utils.doc_parser import DocParser


def visual_model_test(task_path, test_img, schema):
    img = re.split(r'[\\/]', test_img)[-1]
    img_name = ""
    img_type = "jpg"
    last_dot_index = img.rfind(".")
    if last_dot_index != -1:
        img_name = img[:last_dot_index]
        img_type = img[last_dot_index + 1:]

    # 默认模型
    ie = Taskflow("information_extraction", schema=schema, model="uie-x-base")
    results = ie({"doc": test_img})
    pprint(results[0])
    DocParser.write_image_with_results(
        test_img,
        result=results[0],
        save_path="./img_result/" + img_name + "_default." + img_type)
    # 自己训练的模型
    my_ie = Taskflow("information_extraction", schema=schema, model="uie-x-base", task_path=task_path)
    my_results = my_ie({"doc": test_img})
    pprint(my_results[0])
    DocParser.write_image_with_results(
        test_img,
        result=my_results[0],
        save_path="./img_result/" + img_name + "_my." + img_type)


def main(model_type):
    # 开始时间
    start_time = time.time()

    # 结算清单
    if model_type == "settlement":
        task_path = "../../config/model/settlement_list_model"
        test_img_path = "img/PH20240511000638_1_094306_1.jpg"
        schema = ["姓名", "入院日期", "出院日期", "费用总额", "个人现金支付", "个人账户支付", "自费", "医保类型"]
    elif model_type == "discharge":
        task_path = "../../config/model/discharge_record_model"
        test_img_path = "img/PH20240401000003_3_001938_2.jpg"
        schema = ["医院", "科别", "姓名", "入院日期", "出院日期", "主治医生"]
    elif model_type == "cost":
        task_path = "../../config/model/cost_list_model"
        test_img_path = "img/PH20240511000648_4_094542_2.jpg"
        schema = ["姓名", "入院日期", "出院日期", "费用总额"]
    else:
        print("请输入正确的类型！")
        return
    visual_model_test(task_path, test_img_path, schema)

    # 结束时间
    end_time = time.time()
    pprint(f"处理时长：{end_time - start_time}秒")


if __name__ == '__main__':
    main("settlement")
    # main("discharge")
    # main("cost")