#! /usr/bin/env python3

import sys, gzip
import jsonlines
from typing import Any


"""
(define (problem c01)
	(:domain countdown)
	(:objects
		n1 n2 n3 n4  - num
	)
  (:init
		(= (value n1) 3)
		(= (value n2) 4)
		(= (value n3) 5)
		(= (value n4) 6)

		(= (targetvalue) 24)
		(= (numactive) 4)
		(active n1)
		(active n2)
		(active n3)
		(active n4)
	)
	(:goal
		(and
		(goalreached)
		)
	)
)

"""

def create_pddl(example: dict[str, Any]) -> str:
    """Create a PDDL problem for a given example

    Args:
        example: The example

    Returns:
        A PDDL problem string
    """
    state = example["state"]
    goal = example["goal"]
    name = "_".join([str(x) for x in state]) + "_" + str(goal)
    objs = " ".join([f"n{i+1}" for i in range(len(state))])
    problem_str_components = [f"(define (problem c{name})"]
    problem_str_components.append("  (:domain countdown)")
    problem_str_components.append("  (:objects")
    problem_str_components.append(f"    {objs}  - num")
    problem_str_components.append("  )")
    problem_str_components.append("  (:init")
    for i, x in enumerate(state):
        problem_str_components.append(f"    (= (value n{i+1}) {x})")
        problem_str_components.append(f"    (active n{i+1})")

    problem_str_components.append(f"    (= (targetvalue) {goal})")
    problem_str_components.append(f"    (= (numactive) {len(state)})")
    problem_str_components.append("  )")
    problem_str_components.append("  (:goal")
    problem_str_components.append("    (and (goalreached))")
    problem_str_components.append("  )")
    problem_str_components.append(")")
    return "\n".join(problem_str_components)


if __name__ == "__main__":

    if len(sys.argv) < 3:
        print("Usage: generate_pddl.py <source.jsonl.gz> <folder>")
        exit(1)
    file = sys.argv[1]
    folder = sys.argv[2]
    with gzip.open(file) as f:
        for i, line in enumerate(jsonlines.Reader(f)):
            with open(f"{folder}/p{i+1:04}.pddl", "w") as wf:
                wf.write(create_pddl(line))
