import json
import pprint
import ipdb
import os

domains = [
    'translate',
    'qa',
    'chat',
    'summary',
    'harmlessness',
    'math_cot',
    'math_pot',
    'code_exec',
    'code_not_exec'
]

folders = [
    'sub_dev_data',
    'sub_test_data'
]

counter = {
    q: {
        'subtle': 0,
        'complex': 0,
        'obvious': 0
    } for q in ['low', 'med', 'high']
}

for folder in folders:
    for domain in domains:
        path = f'{folder}/{domain}_feedback_correction.json'
        with open(path) as f:
            data = json.load(f)
        for sample in data:
            try:
                score = float(sample['feedback_score'])
            except:
                continue
            quality = sample['metadata']['quality']
            if score <= 1:
                counter[quality]['obvious'] += 1
            elif 1 <  score <= 4:
                counter[quality]['complex'] += 1
            else:
                counter[quality]['subtle'] += 1


pprint.pprint(counter)

