import ast
import json
import random
from pathlib import Path

random.seed(42)

with open("data_02/data_processed.json", encoding="utf-8") as f:
    data = json.load(f)

def extract_operations(expr):
    tree = ast.parse(expr, mode="eval")
    operations = []

    names = {
        ast.Add: "ADD",
        ast.Sub: "SUB",
        ast.Mult: "MUL",
        ast.Div: "DIV"
    }

    def visit(node):
        if isinstance(node, ast.BinOp):
            visit(node.left)
            visit(node.right)
            if type(node.op) in names:
                operations.append(names[type(node.op)])

    visit(tree.body)
    return operations

def convert(item):
    validation = item["validation"]
    steps = validation["process"]

    operation_sequence = []
    for step in steps:
        operation_sequence.extend(
            extract_operations(step["expr"])
        )

    prompt = (
        "Generate one math word problem with the "
        "following calculation structure.\n"
        f"Calculation steps: {len(steps)}\n"
        f"Operations: {' -> '.join(operation_sequence)}\n"
        "Return valid JSON containing question, "
        "explanation, answer, and validation. "
        "The validation program must reference "
        "previous calculation results correctly."
    )

    completion = json.dumps(
        {
            "question": item["question"],
            "explanation": item["explanation"],
            "answer": item["answer"],
            "validation": validation,
        },
        ensure_ascii=False,
        separators=(",", ":")
    )

    return {"prompt": prompt, "completion": completion}

examples = [convert(item) for item in data]

random.shuffle(examples)

n = len(examples)
n_train = int(n * 0.9)
n_valid = int(n * 0.05)

splits = {
    "train": examples[:n_train],
    "valid": examples[n_train:n_train + n_valid],
    "test": examples[n_train + n_valid:],
}

Path("data_03").mkdir(exist_ok=True)

for name, rows in splits.items():
    with open(
        f"data_03/{name}.jsonl",
        "w",
        encoding="utf-8"
    ) as f:
        for row in rows:
            f.write(
                json.dumps(row, ensure_ascii=False) + "\n"
            )

    print(name, len(rows))