import json
import os
import time

# load each json in data/train
data_dir = "./data/train"
data_list = os.listdir(data_dir)
for data in data_list:
    with open(os.path.join(data_dir, data)) as file:
        thy_series = json.load(file)
    Q = thy_series['problem']
    A = thy_series['solution']
    answer = thy_series['answer']
    #### check the formal problem
    check = thy_series.get("check", "")
    thy = thy_series.get(f'generation_0', {})
    # if thy == {}:
        # thy['formal problem'] = thy_series.get("formal problem", "")
        # thy['answer'] = thy_series.get("check", "")
    # thy_series[f'generation_0'] = thy
    fp = thy_series.get("formal problem", "")
    thy_series['formal problem'] = thy["formal problem"]
    thy_series['check'] = thy["answer"]
    with open(os.path.join(data_dir, data), 'w') as file:
        json.dump(thy_series, file, indent=4)