# https://huggingface.co/datasets/openai/gsm8k

import json
import re

import pandas as pd

df = pd.read_parquet("data_00/train-00000-of-00001.parquet")

data = []

for row in df.itertuples():
    cur = {}
    question, explanation = str(row.question), str(row.answer)

    temp = re.search(r"\n#### (.+)$", explanation)
    if temp:
        explanation = explanation[: temp.start()]
        temp = re.sub(",", "", temp.group(1))
        try:
            answer = float(temp)
        except ValueError:
            continue
    else:
        continue

    equations = re.findall(r"<<(.+)=.+>>", explanation)

    validation = {
        "input": {},
        "process": [{"id": "", "expr": i} for i in equations],
        "output": "",
    }

    cur["question"] = question
    cur["explanation"] = explanation
    cur["answer"] = answer
    cur["validation"] = validation

    data.append(cur)

with open('data_01/data.json', 'w') as f:
    json.dump(data, f, indent=4)