from Routes.pokemon_game_desc import PokemonGameDesc


def format_name(name):
    return name.replace(' ', '_').replace('.', '').replace("'", '')

def generate_pddl_domain():
    # Define the types, predicates, and action as shown above
    domain_pddl = """
    (define (domain pokemon)
        (:requirements :strips :typing :equality :quantified-preconditions :conditional-effects :negative-preconditions)
        (:types
            location
            condition
        )
        (:predicates
            (at ?loc - location)
            (connected ?from ?to - location)
            (has ?cond - condition)
            (grants ?loc - location ?cond - condition)
            (requires ?from ?to - location ?cond - condition)
            (requires_grant ?loc - location ?cond - condition ?req - condition)
            (visited ?loc - location)
            (is_town_or_city ?loc - location)
        )
        (:action move
            :parameters (?from ?to - location)
            :precondition (and
                (at ?from)
                (connected ?from ?to)
                ; Remove (not (visited ?to)) if re-visiting is allowed
                ; Universal preconditions for required conditions
                (forall (?cond - condition)
                    (imply (requires ?from ?to ?cond) (has ?cond))
                )
            )
            :effect (and
                (at ?to)
                (not (at ?from))
                (visited ?to)
            )
        )
        (:action acquire
            :parameters (?loc - location ?cond - condition)
            :precondition (and
                (at ?loc)
                (grants ?loc ?cond)
                (not (has ?cond))
                ; Universal preconditions for required grants
                (forall (?req - condition)
                    (imply (requires_grant ?loc ?cond ?req) (has ?req))
                )
            )
            :effect (and
                (has ?cond)
            )
        )
        (:action fly
            :parameters (?from ?to - location)
            :precondition (and
                (at ?from)
                (has Fly)
                (has Fly_out_of_battle)
                (visited ?to)
                (is_town_or_city ?to)
                (not (at ?to))  ; Prevents flying to the current location
            )
            :effect (and
                (at ?to)
                (not (at ?from))
            )
        )
    )
    """
    with open('pokemon_domain.pddl', 'w') as f:
        f.write(domain_pddl)

def generate_pddl_problem(desc: PokemonGameDesc):
    # Extract objects, init, and goal
    locations = set()
    conditions = set()
    connections = []
    grants = []
    requirements = []
    requires_grant = []

    # Gather conditions from node grants and edge requirements
    for node, attrs in desc.graph.nodes(data=True):
        node_formatted = format_name(node)
        locations.add(node_formatted)
        grants_conditions = attrs.get('grants_conditions', [])
        for grant in grants_conditions:
            condition = format_name(grant['condition'])
            conditions.add(condition)
            grants.append((node_formatted, condition))
            required_conditions = grant.get('required_conditions', [])
            for req in required_conditions:
                req_condition = format_name(req)
                conditions.add(req_condition)
                requires_grant.append((node_formatted, condition, req_condition))

    for u, v, attrs in desc.graph.edges(data=True):
        u_formatted = format_name(u)
        v_formatted = format_name(v)
        locations.update([u_formatted, v_formatted])
        # Add both directions for bidirectional movement
        connections.append((u_formatted, v_formatted))
        connections.append((v_formatted, u_formatted))  # Add reverse connection

        edge_condition = attrs.get('condition')
        if edge_condition:
            if isinstance(edge_condition, list):
                for cond in edge_condition:
                    cond_formatted = format_name(cond)
                    conditions.add(cond_formatted)
                    requirements.append((u_formatted, v_formatted, cond_formatted))
                    requirements.append((v_formatted, u_formatted, cond_formatted))  # Reverse
            else:
                cond_formatted = format_name(edge_condition)
                conditions.add(cond_formatted)
                requirements.append((u_formatted, v_formatted, cond_formatted))
                requirements.append((v_formatted, u_formatted, cond_formatted))  # Reverse

    conditions.add(format_name('Fly'))
    conditions.add(format_name('Fly out of battle'))

    formatted_towns_and_cities = [format_name(town) for town in desc.towns_and_cities]

    # Prepare the PDDL problem file content
    problem_pddl = "(define (problem pokemon_problem)\n"
    problem_pddl += "    (:domain pokemon)\n"

    # Objects
    problem_pddl += "    (:objects\n"
    problem_pddl += "        " + " ".join(sorted(locations)) + " - location\n"
    problem_pddl += "        " + " ".join(sorted(conditions)) + " - condition\n"
    problem_pddl += "    )\n"

    # Initial state
    problem_pddl += "    (:init\n"
    problem_pddl += f"        (at {format_name(desc.starting_town)})\n"
    for u, v in connections:
        problem_pddl += f"        (connected {u} {v})\n"
    for loc, cond in grants:
        problem_pddl += f"        (grants {loc} {cond})\n"
    for u, v, cond in requirements:
        problem_pddl += f"        (requires {u} {v} {cond})\n"
    for loc, cond, req in requires_grant:
        problem_pddl += f"        (requires_grant {loc} {cond} {req})\n"
    for town in formatted_towns_and_cities:
        problem_pddl += f"        (is_town_or_city {town})\n"
    problem_pddl += "    )\n"

    # Goal state
    badges = desc.badges
    problem_pddl += "    (:goal\n"
    problem_pddl += "        (and\n"
    problem_pddl += f"            (at {format_name(desc.end_goal)})\n"
    for badge in badges:
        problem_pddl += f"            (has {format_name(badge)})\n"
    problem_pddl += "        )\n"
    problem_pddl += "    )\n"
    problem_pddl += ")\n"

    with open(f'pokemon_problem_{desc.game_name}.pddl', 'w') as f:
        f.write(problem_pddl)