import json

dataset = ['2wikimqa', 'gov_report', 'hotpotqa', 'lcc', 'multi_news', 'multifieldqa_en', 'passage_count', 'passage_retrieval_en', 'qasper', 'repobench-p', 'samsum', 'trec', 'triviaqa']

dataset2type = {
    "2wikimqa" : "mqa",
    'gov_report' : "sum",
    "hotpotqa" : "mqa",
    "lcc" : "code", 
    "multi_news": "sum", 
    "multifieldqa_en" : "qa",
    "passage_count" : "syn", 
    "passage_retrieval_en": "syn",
    'qasper': 'qa', 
    'repobench-p': 'code',
    'samsum': 'few',
    'trec': 'few',
    'triviaqa': 'few'

}

type2dataset = {
    "qa": [ 'multifieldqa_en', 'qasper'], 
    'mqa': ['hotpotqa','2wikimqa' ],
    'sum': ['gov_report', 'multi_news'],
    'few': ['trec', 'triviaqa', 'samsum',], 
    'syn': ['passage_count', 'passage_retrieval_en'], 
    'code': ['lcc', 'repobench-p']
}

def get_results(kv_method):
    base_path = f"./results_longbench/longchat-7b-v1.5-32k_20000_{kv_method}/result.json"
    with open(f"{base_path}", "r", encoding="utf-8") as f:
        data = json.load(f)
    return data

def get_results_task(kv_method):
    data = get_results(kv_method)
        
    res = {aa:0 for aa in type2dataset.keys()}
    # res = {}

    for dataset in dataset2type.keys():
        score = data[dataset]
        type_ = dataset2type[dataset]
        res[type_] = res[type_] + score
    return res


print(f"     ", end='')
for type_ in type2dataset.keys():
    print(f"| {type_:>7} ", end='')
print("|")
for kv_method in ['exact', 'h2o', 'sink', 'clustergen']
    res = get_results_task(kv_method)
    print(f"{kv_method:<8} ", end='')
    for type_ in type2dataset.keys():
        print(f"& {res[type_]:<8.2f}", end='')
    print("\\\\")
