import os
import re
import json
import logging
import pandas as pd


def load_a_file(a_file, content=None):
    if not content:
        content = []
    with open(a_file, 'r', encoding="utf-8", errors="ignore") as f:
        lines = f.readlines()
        content.extend(lines)
    return content


def save_a_file(a_file_content, file_name):
    with open(file_name, "w", encoding='utf-8') as f:
        f.writelines(a_file_content)


def get_file_list(folder, file_type_list):
    filelist = []
    for dirpath, _, filenames in os.walk(folder):
        for file in filenames:
            file_type = file.split('.')[-1]
            if file_type in file_type_list:
                file_fullname = os.path.join(dirpath, file)
                filelist.append(file_fullname)

    return filelist


def process_gen_files(gen_path_dir, len_ori):
    txt_files_lst = get_file_list(gen_path_dir, file_type_list)
    try:
        txt_files_lst = sorted(txt_files_lst, key=lambda x: int(x.split('_')[-1].split('.')[0]), reverse=True)
    except Exception as e:
        print(e)

    content = []
    for i in range(len(txt_files_lst)):
        a_file = txt_files_lst[i]
        content = load_a_file(a_file, content)

    diff_len = len(content) - len_ori

    if diff_len > len(txt_files_lst) or diff_len < 0:
        return content

    content_all = []
    for i in range(diff_len):
        a_file = txt_files_lst[i]
        content = load_a_file(a_file)
        content_all.extend(content[:-1])
    for i in range(diff_len, len(txt_files_lst)):
        a_file = txt_files_lst[i]
        content = load_a_file(a_file)
        content_all.extend(content)

    return content_all


def get_ans_gen_letter(qa_text):
    if not qa_text:
        return qa_text
    a_text = qa_text.strip().split('[SEP]')[1]
    q_text = qa_text.split('[SEP]')[0].replace('．', '.')

    flag_q = re.findall(r'(?:错误)|(?:incorrect)|(?:not true)|(?:False)|(?:无法推断)', q_text)
    if flag_q:
        flag_a = False
    else:
        flag_a = True
    a_text = a_text.replace('$', '').replace('Ｃ', 'C').replace('Ａ', 'A').replace('Ｂ', 'B').replace('Ｄ', 'D')
    a_text = re.sub(r'\\rm\s*\{([A-Z]+)\}', '\\1', a_text)
    a_text = re.sub(r'\\mathrm\s*\{([A-Z]+)\}', '\\1', a_text)
    a_text = re.sub(r'\\boxed\s*\{\s*([A-Z]+)\s*\}\s*', '\\1', a_text)
    a_text = re.sub(r'\\text\s*\{[\(]?([A-Z]+)[\)]?\}', '\\1', a_text)
    a_text = re.sub(r'<font[^>]*?>([A-K]+)</font>', '\\1', a_text)
    a_text = re.sub(r'\\textbf\s*\{\(([A-Z]+)\)\}', '\\1.', a_text)
    a_text = re.sub(r'\\mathbf\s*\{([^{]*?)\}', '\\1', a_text)

    ans_text = re.findall(
        r'(?:correct|most appropriate|most accurate|plausible|reasonable|best|closest) (?:answer|option|Option|statement) (?:is|would be)[option\s<n>:*\(]*?[A-K]\s*[\),:\.]', a_text)

    if not ans_text:
        ans_text = re.findall(
            r'(?:correct|most appropriate|most accurate|plausible|reasonable|best|closest) answer [\sa-z]*? "[^"]*?"[isoption\s<n>:*\(]*?[A-K]\s*[\),:\.]', a_text)
    else:
        ans_text = re.sub(
            r'(?:correct|most appropriate|most accurate|plausible|reasonable|best|closest) (?:answer|option|Option|statement) (?:is|would be)[option\s<n>:*\(]*?([A-K])\s*[\),:\.]', '\\1', ans_text[-1])

    if not ans_text:
        ans_text = re.findall(
            r'(?:option|Option) [A-K]\s*[\),:\.][\sa-z]*? "[^"]*?"[a-z\s]*?the correct answer', a_text)
    else:
        ans_text = re.sub(
            r'(?:correct|most appropriate|most accurate|plausible|reasonable|best) answer [\sa-z]*? "[^"]*?"[isoption\s<n>:*\(]*?([A-K])\s*[\),:\.]', '\\1', ans_text[-1])
    if not ans_text:
        ans_text = re.findall(
            r'(?:option|Option) [A-K](?:\s*[\.\)][\s:a-zA-Z]{0,10}|) is the correct answer', a_text)
    else:
        ans_text = re.sub(
            r'(?:option|Option) ([A-K])\s*[\),:\.][\sa-z]*? "[^"]*?"[a-z\s]*?the correct answer', '\\1', ans_text[-1])
    if not ans_text:
        ans_text = re.findall(
            r' answer (?:is|would be) .*?[A-K]\s*[\)\.]', a_text)
    else:
        ans_text = re.sub(
            r'(?:option|Option) ([A-K])(?:\s*[\.\)][\s:a-zA-Z]{0,10}|) is the correct answer', '\\1', ans_text[-1])
    if not ans_text:
        if not flag_a:
            ans_text = re.findall(
                r'incorrect statement (?:is|would be)[option\s<n>:*\(]*?[A-K]\s*[\),:\.]', a_text)

            if ans_text:
                ans_true = []
                for al in range(len(ans_text)):
                    choice_sub = re.sub(
                        r'incorrect statement (?:is|would be)[option\s<n>:*\(]*?([A-K])\s*[\),:\.]', '\\1', ans_text[al])
                    ans_true.append(choice_sub)
                if ans_true:
                    ans_true = ''.join(sorted(list(set(ans_true))))
                    ans_text = [ans_true]
    else:
        ans_text = [
            re.sub(r' answer (?:is|would be) [^\.]*?([A-K])\s*[\)\.]', '\\1', ans_text[-1])]
    if not ans_text:
        split_n = a_text.split('<n>')
        ans_text = re.findall(
            r'(?:option|answer|Option|answer|)[a-z\s]*?[A-K][\s\.,:\)]', split_n[-1])
        if not ans_text and len(split_n) >= 2:
            ans_text = re.findall(
                r'(?:option|answer|Option|answer|)[a-z\s]*?[A-K][\s\.,:\)]', split_n[-2])
        if ans_text:
            ans_text = [re.sub(
                r'(?:option|answer|Option|answer|)[a-z\s]*?([A-K])[\s\.,:\)]', '\\1', ans_text[-1])]
    if re.findall(r'none of the [a-z\s*] options are correct.', a_text):
        ans_text = None
    if ans_text:
        try:
            if re.findall(r'[A-K]+', ans_text[-1]) != ans_text:
                ans = re.findall(r'[A-K]', ans_text[-1])[-1]
            else:
                ans = ans_text
            ans = ''.join(sorted(list(set(ans))))
            return ans
        except Exception as e:
            print(e)
            return None
    else:
        return None


def eval_gen_res(qa_content, accuracy_df, origin_file_path=None, txt_eval_res_dir=None):
    qa_content_new = qa_content[:]
    try:
        file_name = os.path.basename(origin_file_path).split('.')[0]
        content_yuan_all_dict = {}
        content_qa = []
        with open(origin_file_path, 'r', encoding='utf-8') as f:
            for id, line in enumerate(f):
                line = line.strip()
                if not line:
                    continue
                g = re.sub(r'(<n>)+', '<n>', line).strip()
                g = g.split('[SEP]')

                if len(g) == 2:
                    content_yuan_all_dict[g[0].strip() + '[SEP]' + str(id)] = g[1].replace(' ', '')
                    content_qa.append(line)

    except Exception as e:
        print('Error: ', e)

    text_answer_true, text_answer_false = [], []
    # print(f'------------Eval--{file_name}.--------------------')
    # print(f'len(content_yuan_all_dict):{len(content_yuan_all_dict)}')
    content_no_ans = []
    content_q = []
    for i in range(len(qa_content)):
        qa_content[i] = qa_content[i].strip()
        if not qa_content[i]:
            continue
        qa_content[i] = re.sub('(\s*<n>\s*)+?$', '', qa_content[i])

        q_text = qa_content[i].split('[SEP]')[0].strip()
        q_text = re.sub(r'(<n>)+', '<n>', q_text).strip()
        ans_true = ''
        for key in list(content_yuan_all_dict.keys()):
            if q_text in key.split('[SEP]')[0] or key.split('[SEP]')[0] in q_text:
                ans_true = content_yuan_all_dict[key]
                content_q.append(key)
                del content_yuan_all_dict[key]
                break
        if not ans_true:
            continue
        if '[SEP]' not in qa_content[i]:
            qa_content[i] = qa_content[i].replace(
                key.split('[SEP]')[0], key.split('[SEP]')[0] + '[SEP]')

        ans_true = ans_true.replace('Answer:', '').replace('Answer：', '').replace('.', '').replace('Answers:', '').replace('Answers：', '').strip()
        if not ans_true:
            print(f'No correct answer found for question: "{q_text}".')
            continue
        qa_content[i] = qa_content[i].split('[EOD]')[0]
        try:
            ans_gen = get_ans_gen_letter(qa_content[i])
        except Exception as e:
            ans_gen = None
        qa_content_new[i] = ''
        if not ans_gen:
            content_no_ans.append(qa_content[i] + '\n')
        else:
            if ans_true == ans_gen:
                text_answer_true.append(qa_content[i] + ' [EOD] ' + ans_gen + ' [ANS] ' + ans_true + '\n')
                continue
        if ans_gen:
            text_answer_false.append(qa_content[i] + ' [EOD] ' + ans_gen + ' [ANS] ' + ans_true + '\n')
        else:
            text_answer_false.append(qa_content[i] + ' [EOD] ' + ' None ' + ' [ANS] ' + ans_true + '\n')

    try:
        accuracy = len(text_answer_true) / (len(text_answer_true) + len(text_answer_false))
    except:
        accuracy = 0
    # print(f'Number of correct answers:{len(text_answer_true)}')
    # print(f'Number of incorrect answers:{len(text_answer_false)}')
    # print('accuracy:', accuracy)

    accuracy_dct = {"file_name": file_name, "accuracy": accuracy, 'true_num': len(text_answer_true), 'false_num': len(
        text_answer_false), f'Raw total lines of {file_name}': len(content_qa)}
    if accuracy_df.empty:
        accuracy_df = pd.DataFrame(accuracy_dct, index=[0])
    else:
        accuracy_df.loc[len(accuracy_df)] = accuracy_dct.values()

    try:
        if txt_eval_res_dir:
            if not os.path.exists(txt_eval_res_dir):
                os.makedirs(txt_eval_res_dir)
            save_a_file(text_answer_true, os.path.join(txt_eval_res_dir, f"{file_name}_true.txt"))
            save_a_file(text_answer_false, os.path.join(txt_eval_res_dir, f"{file_name}_false.txt"))
            with open(os.path.join(txt_eval_res_dir, 'accuracy_split.jsonl'), 'a', encoding='utf-8') as f:
                json.dump(accuracy_dct, f, ensure_ascii=False)
                f.write('\n')

    except Exception as e:
        print('save Error: ', e)
    qa_content_new = [item for item in qa_content_new if item]
    return qa_content_new, accuracy_df


def main():
    global file_type_list
    file_type_list = ['jsonl']
    len_ori = 14042

    origin_file_path = "<Specify path>"
    eval_file_path = "<Specify path>"
    txt_eval_res_dir = "<Specify path>"
    
    origin_file_lst = get_file_list(origin_file_path, ['txt'])
    if os.path.isfile(eval_file_path):
        qa_content = load_a_file(eval_file_path)
    else:
        qa_content = process_gen_files(eval_file_path, len_ori)
    accuracy_df = pd.DataFrame()
    for origin_file_path in origin_file_lst:
        qa_content, accuracy_df = eval_gen_res(qa_content, accuracy_df, origin_file_path, txt_eval_res_dir)

    accuracy_df.to_csv(f'{txt_eval_res_dir}/accuracy_split.csv', index=False)

    true_num_all = accuracy_df['true_num'].sum()
    false_num_all = accuracy_df['false_num'].sum()

    print(f'Number of correct answers:{true_num_all}')
    print(f'Number of incorrect answers:{false_num_all}')
    print(f'accuracy:{true_num_all/(true_num_all + false_num_all)}')


if __name__ == "__main__":

    main()
