# Modified from: https://github.com/Yuliang-Liu/MultimodalOCR/blob/main/OCRBench/scripts/LLaVA1_5.py

from copy import deepcopy
import os
import csv
import json
import pickle
import argparse
import numpy as np
from tqdm import tqdm
from datasets import load_dataset
import logging
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO,)


def save_json(json_list,save_path):
    with open(save_path, 'w') as file:
        json.dump(json_list, file,indent=4)

OCRBench_score = {"Regular Text Recognition":0,"Irregular Text Recognition":0,"Artistic Text Recognition":0,"Handwriting Recognition":0,
"Digit String Recognition":0,"Non-Semantic Text Recognition":0,"Scene Text-centric VQA":0,"Doc-oriented VQA":0,"Doc-oriented VQA":0,
"Key Information Extraction":0,"Handwritten Mathematical Expression Recognition":0}
AllDataset_score = {"IIIT5K":0,"svt":0,"IC13_857":0,"IC15_1811":0,"svtp":0,"ct80":0,"cocotext":0,"ctw":0,"totaltext":0,"HOST":0,"WOST":0,"WordArt":0,"IAM":0,"ReCTS":0,"ORAND":0,"NonSemanticText":0,"SemanticText":0,
"STVQA":0,"textVQA":0,"ocrVQA":0,"ESTVQA":0,"ESTVQA_cn":0,"docVQA":0,"infographicVQA":0,"ChartQA":0,"ChartQA_Human":0,"FUNSD":0,"SROIE":0,"POIE":0,"HME100k":0}
num_all = {"IIIT5K":0,"svt":0,"IC13_857":0,"IC15_1811":0,"svtp":0,"ct80":0,"cocotext":0,"ctw":0,"totaltext":0,"HOST":0,"WOST":0,"WordArt":0,"IAM":0,"ReCTS":0,"ORAND":0,"NonSemanticText":0,"SemanticText":0,
"STVQA":0,"textVQA":0,"ocrVQA":0,"ESTVQA":0,"ESTVQA_cn":0,"docVQA":0,"infographicVQA":0,"ChartQA":0,"ChartQA_Human":0,"FUNSD":0,"SROIE":0,"POIE":0,"HME100k":0}

def ocrbench_eval(eval_file, save_file):
    dataset = load_dataset("echo840/OCRBench")["test"]
    answers = [json.loads(q) for q in open(eval_file)]
    predictions = {pred['question_id']: pred["text"] for pred in answers}

    for i, data in enumerate(dataset):
        predict = predictions.get(i, '')
        data_type = data["question_type"]
        dataset_name = data["dataset"]
        answers = data["answer"]
        data['result'] = 0
        if dataset_name == "HME100k":
            if type(answers)==list:
                for j in range(len(answers)):
                    answer = answers[j].strip().replace("\n"," ").replace(" ","")
                    predict = predict.strip().replace("\n"," ").replace(" ","")
                    if answer in predict:
                        data['result'] = 1
            else:
                answers = answers.strip().replace("\n"," ").replace(" ","")
                predict = predict.strip().replace("\n"," ").replace(" ","")
                if answers in predict:
                    data['result'] = 1
        else:
            if type(answers)==list:
                for j in range(len(answers)):
                    answer = answers[j].lower().strip().replace("\n"," ")
                    predict = predict.lower().strip().replace("\n"," ")
                    if answer in predict:
                        data['result'] = 1
            else:
                answers = answers.lower().strip().replace("\n"," ")
                predict = predict.lower().strip().replace("\n"," ")
                if answers in predict:
                    data['result'] = 1
        OCRBench_score[data_type] += data['result']

    recognition_score = OCRBench_score['Regular Text Recognition']+OCRBench_score['Irregular Text Recognition']+OCRBench_score['Artistic Text Recognition']+OCRBench_score['Handwriting Recognition']+OCRBench_score['Digit String Recognition']+OCRBench_score['Non-Semantic Text Recognition']
    Final_score = recognition_score+OCRBench_score['Scene Text-centric VQA']+OCRBench_score['Doc-oriented VQA']+OCRBench_score['Key Information Extraction']+OCRBench_score['Handwritten Mathematical Expression Recognition']
    print("###########################OCRBench##############################")
    print(f"Text Recognition(Total 300):{recognition_score/3}")
    print("------------------Details of Recognition Score-------------------")
    print(f"Regular Text Recognition(Total 50): {OCRBench_score['Regular Text Recognition']}")
    print(f"Irregular Text Recognition(Total 50): {OCRBench_score['Irregular Text Recognition']}")
    print(f"Artistic Text Recognition(Total 50): {OCRBench_score['Artistic Text Recognition']}")
    print(f"Handwriting Recognition(Total 50): {OCRBench_score['Handwriting Recognition']}")
    print(f"Digit String Recognition(Total 50): {OCRBench_score['Digit String Recognition']}")
    print(f"Non-Semantic Text Recognition(Total 50): {OCRBench_score['Non-Semantic Text Recognition']}")
    print("----------------------------------------------------------------")
    print(f"Scene Text-centric VQA(Total 200): {OCRBench_score['Scene Text-centric VQA']}")
    print("----------------------------------------------------------------")
    print(f"Doc-oriented VQA(Total 200): {OCRBench_score['Doc-oriented VQA']}")
    print("----------------------------------------------------------------")
    print(f"Key Information Extraction(Total 200): {OCRBench_score['Key Information Extraction']}")
    print("----------------------------------------------------------------")
    print(f"Handwritten Mathematical Expression Recognition(Total 100): {OCRBench_score['Handwritten Mathematical Expression Recognition']}")
    print("----------------------Final Score-------------------------------")
    print(f"Final Score(Total 1000): {Final_score}")

    results = {
        "Final Score": Final_score,
        "Text Recognition": recognition_score,
        "Regular Text Recognition": OCRBench_score['Regular Text Recognition'],
        "Irregular Text Recognition": OCRBench_score['Irregular Text Recognition'],
        "Artistic Text Recognition": OCRBench_score['Artistic Text Recognition'],
        "Handwriting Recognition": OCRBench_score['Handwriting Recognition'],
        "Digit String Recognition": OCRBench_score['Digit String Recognition'],
        "Non-Semantic Text Recognition": OCRBench_score['Non-Semantic Text Recognition'],
        "Scene Text-centric VQA": OCRBench_score['Scene Text-centric VQA'],
        "Doc-oriented VQA": OCRBench_score['Doc-oriented VQA'],
        "Key Information Extraction": OCRBench_score['Key Information Extraction'],
        "Handwritten Mathematical Expression Recognition": OCRBench_score['Handwritten Mathematical Expression Recognition'],
    }
    os.makedirs(os.path.dirname(save_file), exist_ok=True)
    save_json(results, save_file)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--result-file", type=str)
    parser.add_argument("--result-save", type=str)
    args = parser.parse_args()

    ocrbench_eval(args.result_file, args.result_save)