import pandas as pd
from tabulate import tabulate

# 创建一个示例 DataFrame
data = {
    'Model': [
        'GPT-4', 'GLM-4', 'GPT-3.5-turbo', 'Qwen-72B-Chat', 'ERNIE-Bot-4.0', 'LLaMA-2-70B',
        'DevOps-Model-14B-Chat', 'GLM-3-turbo', 'Qwen-14B-Chat', 'LLaMA-2-13B', 'InternLM2-Chat-20B',
        'LLaMA-2-7B', 'Qwen-7B-Chat', 'Baichuan2-13B-Chat', 'InternLM2-Chat-7B', 'Mistral-7B', 'ChatGLM3-6B'
    ],
    'Naive': [
        '/', '64.77', '68.30', '70.32', '60.00', '55.00', '63.85', '59.53', '62.60', '53.30', '60.48',
        '48.20', '52.10', '51.90', '48.20', '47.22', '42.10'
    ],
    'SC': [
        '/', '64.77', '68.30', '70.32', '60.00', '56.20', '61.96', '59.53', '59.70', '53.00', '60.48',
        '46.80', '51.00', '51.60', '48.20', '47.22', '42.10'
    ],
    'CoT': [
        '88.70', '77.06', '70.90', '70.13', '70.00', '66.80', '41.15', '63.65', '50.58', '56.80', '45.10',
        '52.00', '48.30', '44.50', '49.74', '45.58', '43.47'
    ],
    'CoT+SC': [
        '/', '77.06', '72.50', '70.22', '70.00', '67.20', '44.01', '63.65', '55.88', '61.00', '45.10',
        '55.20', '49.80', '47.45', '49.74', '45.58', '43.47'
    ]
}

df = pd.DataFrame(data)

# 使用tabulate生成LaTeX表格
latex_table = tabulate(df, headers='keys', tablefmt='latex', showindex=False,
                       colalign='left')

def gen_latex_table(caption, label, dataframe):
    table = tabulate(dataframe, headers='keys', tablefmt='latex', showindex=False,
                     colalign='left')
    table = (
        "\\begin{table}[]\n"
        f"\\caption{{{caption}}}\n"
        f"\\label{{{label}}}\n"
        "\\footnotesize\n"
        f"{table}\n"
        "\\end{table}"
    )
    # 确认生成的\hline只有三个
    assert table.count("\\hline") == 3
    # 将table中的第一个\hline改为\toprule
    table = table.replace("\\hline", "\\toprule", 1)
    # 将table中的第二个\hline改为\midrule
    table = table.replace("\\hline", "\\midrule", 1)
    # 将table中的最后一个\hline改为\bottomrule
    table = table.replace("\\hline", "\\bottom", 1)
    return table

# # 添加表格环境
# latex_table = (
#     "\\begin{table}[]\n"
#     "\\caption{LLMs' overall performance (Accuracy\\%) on Wired Network Operations English test set (3-shot). "
#     "\\normalfont Models are ranked based on their best performance (marked as bold) among different settings.}\n"
#     "\\label{tab:network_eng_3shot}\n"
#     "\\footnotesize\n"
#     f"{latex_table}\n"
#     "\\end{table}"
# )
# latex_table = gen_latex_table(
#     caption="LLMs' overall performance (Accuracy\%) on Wired Network Operations English test set (3-shot). "
#             "Models are ranked based on their best performance (marked as bold) among different settings.",
#     label="tab:network_eng_3shot",
#     table=latex_table
# )

# print(latex_table)