import json


baselines = ['GPT_score', 'BERT_score', 'ROUGE_score', 'BLEU_score', 'BART_score', 'BLEURT_score', 'UNIEVAL_score']
modes = ['human_a', 'human_b', 'model_a', 'model_b']

FILES = ['../benchmark/english/geoen_llm.json', '../benchmark/english/polen_llm.json', '../benchmark/english/hisen_llm.json', '../benchmark/english/psyen_llm.json', '../benchmark/english/meden_llm.json', '../benchmark/english/lawen_llm.json']

# compress float in 2 decimal places
# TypeError: type dict doesn't define __round__ method

for file in FILES:
    with open(file, 'r') as f:
        data = json.load(f)
    if 'geo' in file or 'his' in file or 'pol' in file:
        for d in data:
            for baseline in baselines:
                for mode in modes:
                    if baseline == 'ROUGE_score':
                        key = f"{baseline}_{mode}"
                        d[f"{key}_rounded"] = round(d[key]['rougeL'][2], 2)
                    else:
                        key = f"{baseline}_{mode}"
                    # AttributeError: 'float' object has no attribute 'round'
                        d[f"{key}_rounded"] = round(d[key], 2)
        with open(file
                    , 'w') as f:
                json.dump(data, f, indent=4)
    else:
        for d in data:
            for baseline in baselines:
                for key in ['model_a', 'model_b']:
                    if baseline == 'ROUGE_score':
                        key = f"{baseline}_{key}"
                        d[f"{key}_rounded"] = round(d[key]['rougeL'][2], 2)
                    else:
                        key = f"{baseline}_{key}"
                        d[f"{key}_rounded"] = round(d[key], 2)
    with open(file, 'w') as f:
        json.dump(data, f, indent=4)
        
                    