# Copyright 2022 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
This file contains functions to convert PushWorld puzzles into PDDL problems and
domains. This PDDL representation involves the following design choices:


1) Optimizations for object movements and collision checking

    To minimize the computational cost of advancing the state during a search process,
    the state only contains a single position for each object. If instead the state
    stored all occupied positions of each object (since objects can have shapes that
    occupy multiple positions), moving an object would require updating all of its
    occupied positions, which would be slower.

    To support this optimization, all colliding positions between pairs of objects
    are precomputed and stored in the PDDL problem. Because colliding positions are
    precomputed, planners can detect collisions between a pair of objects in O(1) time.


2) Decomposing simultaneous object movements into multiple actions

    Any number of objects can push each other in a chain. For example, objects A, B,
    and C could be in contact such that when the agent pushes A, then A simultaneously
    pushes B, then B simultaneously pushes C, etc. To move object chains of arbitrary
    length with a single action would require derived predicates to label all objects
    that should be pushed, but derived predicates are not widely supported by existing
    classical planners (as of September 2022). To maximize compatibility with existing
    planners, the PDDL representation generated by this file uses one `push` action to
    move one object at a time. Pushing multiple objects simultaneously (i.e. due to a
    single movement of the agent) therefore requires multiple PDDL actions. This
    approach has the same big-O complexity as derived predicates for checking
    collisions to determine which objects to push.


3) Memory optimizations for Best-First Width Search

    The Best-First Width Search (BFWS) planner from
    https://github.com/nirlipo/BFWS-public contains two versions: One that uses
    `libff` (referring to the Fast Forward heuristic) from LAPKT to preprocess the PDDL
    representation, and another that uses Python code from Fast Downward (FD).
    On PushWorld puzzles, the FD preprocessing is orders of magnitude slower than the
    FF preprocessing, so benchmarks use the FF version. However, `libff` is written
    so that it allocates memory for (total number of constants)^(max predicate arity),
    which exceeds 30 Gb in typical puzzles. To reduce memory usage, a `for_bfws` mode
    is available to reduce the maximum predicate arity from 4 to 3, at the expense of
    more conditions in predicates.

"""
import itertools
import os
from typing import Tuple

import numpy as np

from pushworld.config import (
    BENCHMARK_PUZZLES_PATH,
    DOMAIN_SUFFIX,
    PROBLEM_SUFFIX,
    PUZZLE_EXTENSION,
)
from pushworld.puzzle import AGENT_IDX, PushWorldPuzzle, points_overlap
from pushworld.utils.filesystem import map_files_with_extension

domain_template = """(define
    (domain {problem_name})
    (:requirements :typing :strips :conditional-effects :negative-preconditions)

    (:types
        position - object
        direction - object

        ; Any object that can move and push other objects
        moveable-object - object

        ; The object that can be directly controlled
        agent-object - moveable-object

        ; A pair of objects, only used for BFWS. This is an auxiliary type to reduce
        ; the arity of the `in-collision` predicate. BFWS uses the `libff` library from
        ; LAPKT, and `libff` allocates memory for (num constants)^(predicate arity),
        ; which can exceed 100 Gb in some problems when `in-collision` has arity=4.
        object-pair - object
    )

    (:constants
        agent - agent-object
        up down left right - direction
        {moveable_objects} - moveable-object
        {object_pair_names}
    )

    (:predicates
        (should-move ?obj - moveable-object ?dir - direction)
        (has-moved ?obj - moveable-object)
        (at ?obj - moveable-object ?pos - position)
        (connected ?from - position ?to - position ?dir - direction)
        (wall-collision ?obj - moveable-object ?next-pos - position)
        {in_collision_predicate}
        {is_pair_predicate}
    )

    (:action move-agent
        :parameters (?dir - direction)
        :precondition (and
{no_objects_should_move}
        )
        :effect (and
            (should-move agent ?dir)
            (forall
                (?obj - moveable-object)
                (not (has-moved ?obj)))
        )
    )

    (:action push
        :parameters ( ?obj - moveable-object ?dir - direction ?pos - position ?next-pos - position
        )
        :precondition (and
            (should-move ?obj ?dir)
            (not (has-moved ?obj))
            (at ?obj ?pos)
            (connected ?pos ?next-pos ?dir)
            (not (wall-collision ?obj ?next-pos))
        )
        :effect (and
            (not (at ?obj ?pos))
            (at ?obj ?next-pos)
            (has-moved ?obj)
            (not (should-move ?obj ?dir))
            (forall (?other-obj - moveable-object)
                (when
                    (and
                        (not (has-moved ?other-obj)){push_condition}
                    )
                    (should-move ?other-obj ?dir)
                )
            )
        )
    )
)"""

push_condition = """
                        (exists (?other-pos - position)
                            (and
                                (at ?other-obj ?other-pos)
                                (in-collision ?obj ?next-pos ?other-obj ?other-pos)
                            )
                        )
"""

bfws_push_condition = """
                        (exists (?pair - object-pair ?other-pos - position)
                            (and
                                (at ?other-obj ?other-pos)
                                (is-pair ?pair ?obj ?other-obj)
                                (in-collision ?pair ?next-pos ?other-pos)
                            )
                        )
"""

is_pair_predicate = """
        ; Only used for BFWS
        (is-pair
            ?pair - object-pair
            ?obj - moveable-object
            ?other-obj - moveable-object
        )
"""

in_collision_predicate = """
        (in-collision
            ?obj - moveable-object
            ?pos - position
            ?other-obj - moveable-object
            ?other-pos - position
        )
"""

bfws_in_collision_predicate = """
        (in-collision
            ?pair - object-pair
            ?pos - position
            ?other-pos - position
        )
"""

problem_template = """(define
    (problem {problem_name})
    (:domain {problem_name})

    (:objects
{positions_decl}
    )

    (:init
{initial_object_positions}
{wall_collisions}
        ; Define the grid of positions. (0, 0) is the top-left.
{position_connections}
        ; Enumerate all collisions between all pairs of objects in all positions,
        ; skipping positions that are already forbidden by walls. The "agent" object
        ; never occurs on the right because it is always the first pusher.
{object_collisions}
{object_pairs}
    )

    (:goal
        (and
{position_goals}
{no_objects_should_move}
        )
    )
)"""


def puzzle_to_pddl(
    name: str, puzzle: PushWorldPuzzle, for_bfws: bool = False
) -> Tuple[str, str]:
    """Converts the given puzzle into a PDDL description.

    Args:
        name: The name of the puzzle.
        puzzle: The puzzle to convert to PDDL.
        for_bfws: If True, the PDDL is adjusted to minimize the memory usage of the
            `ff-version` of the Best-First Width Search planner from:
            https://github.com/nirlipo/BFWS-public

    Returns:
        A (PDDL domain, PDDL problem) tuple.
    """
    moveable_objects = " ".join(
        f"m{i}" for i in range(len(puzzle.initial_state)) if i != AGENT_IDX
    )

    width, height = puzzle.dimensions

    positions_decl = "\n".join(
        (" " * 8 + " ".join([f"pos{x}-{y}" for x in range(width - 2)]) + " - position")
        for y in range(height - 2)
    )

    initial_object_positions = ""
    object_names = []
    object_sizes = []
    for i, obj in enumerate(puzzle.movable_objects):
        obj_name = "agent" if i == AGENT_IDX else f"m{i}"
        object_names.append(obj_name)
        object_sizes.append(np.max(list(obj.cells), axis=0) + 1)
        pos = obj.position
        initial_object_positions += (
            f"        (at {obj_name} pos{pos[0]-1}-{pos[1]-1})\n"
        )

    wall_collisions = ""
    collision_free_positions = []
    for i, obj in enumerate(puzzle.movable_objects):
        size = object_sizes[i]
        name = object_names[i]
        wall_positions = (
            puzzle.agent_wall_positions if i == AGENT_IDX else puzzle.wall_positions
        )

        obj_free_positions = []
        collision_free_positions.append(obj_free_positions)

        for x, y in itertools.product(
            range(width - 1 - size[0]),
            range(height - 1 - size[1]),
        ):
            if points_overlap(obj.cells, wall_positions, (x + 1, y + 1)):
                wall_collisions += " " * 8 + f"(wall-collision {name} pos{x}-{y})\n"
            else:
                obj_free_positions.append((x, y))

    position_connections = ""
    for y in range(height - 2):
        for x in range(width - 2):
            if x > 0:
                position_connections += (
                    f"        (connected pos{x}-{y} pos{x-1}-{y} left)\n"
                )
            if x + 1 < width - 2:
                position_connections += (
                    f"        (connected pos{x}-{y} pos{x+1}-{y} right)\n"
                )
            if y > 0:
                position_connections += (
                    f"        (connected pos{x}-{y} pos{x}-{y-1} up)\n"
                )
            if y + 1 < height - 2:
                position_connections += (
                    f"        (connected pos{x}-{y} pos{x}-{y+1} down)\n"
                )

    object_collisions_list = []
    for i, obj_a in enumerate(puzzle.movable_objects):
        name_a = object_names[i]

        for j, obj_b in enumerate(puzzle.movable_objects[i + 1 :], start=i + 1):
            name_b = object_names[j]
            offsets_to_collisions = {}

            for (x_a, y_a), (x_b, y_b) in itertools.product(
                collision_free_positions[i], collision_free_positions[j]
            ):
                offset = (x_a - x_b, y_a - y_b)
                if offset in offsets_to_collisions:
                    in_collision = offsets_to_collisions[offset]
                else:
                    in_collision = points_overlap(obj_a.cells, obj_b.cells, offset)
                    offsets_to_collisions[offset] = in_collision

                if in_collision:
                    if for_bfws:
                        a_to_b_collision = (
                            f"(in-collision {name_a}-{name_b} pos{x_a}-{y_a} "
                            f"pos{x_b}-{y_b})"
                        )
                    else:
                        a_to_b_collision = (
                            f"(in-collision {name_a} pos{x_a}-{y_a} {name_b} "
                            f"pos{x_b}-{y_b})"
                        )

                    object_collisions_list.append(a_to_b_collision)

                    if i != AGENT_IDX:
                        # Also consider the swapped order
                        if for_bfws:
                            b_to_a_collision = (
                                f"(in-collision {name_b}-{name_a} pos{x_b}-{y_b} "
                                f"pos{x_a}-{y_a})"
                            )
                        else:
                            b_to_a_collision = (
                                f"(in-collision {name_b} pos{x_b}-{y_b} {name_a} "
                                f"pos{x_a}-{y_a})"
                            )

                        object_collisions_list.append(b_to_a_collision)

    object_collisions = ("\n" + " " * 8).join(object_collisions_list)

    position_goals = ""
    for obj_name, (x, y) in zip(object_names[1:], puzzle.goal_state):
        position_goals += f"            (at {obj_name} pos{x-1}-{y-1})\n"

    no_objects_should_move = []
    for obj in ["agent"] + moveable_objects.split():
        for direction in ["left", "right", "up", "down"]:
            no_objects_should_move.append(
                " " * 12 + f"(not (should-move {obj} {direction}))"
            )
    no_objects_should_move = "\n".join(no_objects_should_move)

    if for_bfws:
        object_pairs_list = []
        object_pair_names_list = []
        for name_a in object_names:
            for name_b in object_names:
                if name_a != name_b:
                    pair_name = f"{name_a}-{name_b}"
                    object_pair_names_list.append(pair_name)
                    object_pairs_list.append(f"(is-pair {pair_name} {name_a} {name_b})")

        object_pairs = "\n        ".join(object_pairs_list)
        object_pair_names = " ".join(object_pair_names_list) + " - object-pair"

    else:
        object_pairs = ""
        object_pair_names = ""

    problem_name = name.replace(" ", "_")

    return (
        domain_template.format(
            problem_name=problem_name,
            moveable_objects=moveable_objects,
            is_pair_predicate=is_pair_predicate if for_bfws else "",
            in_collision_predicate=(
                bfws_in_collision_predicate if for_bfws else in_collision_predicate
            ),
            object_pair_names=object_pair_names,
            no_objects_should_move=no_objects_should_move,
            push_condition=bfws_push_condition if for_bfws else push_condition,
        ),
        problem_template.format(
            problem_name=problem_name,
            positions_decl=positions_decl,
            initial_object_positions=initial_object_positions,
            wall_collisions=wall_collisions,
            position_connections=position_connections,
            object_collisions=object_collisions,
            object_pairs=object_pairs,
            position_goals=position_goals,
            no_objects_should_move=no_objects_should_move,
        ),
    )


def convert_all_puzzles_to_pddl(
    pddl_path: str, puzzle_path: str = BENCHMARK_PUZZLES_PATH, for_bfws: bool = False
) -> None:
    """Iterates over all PushWorld puzzles found in the given `puzzle_path` directory
    and converts them into pairs of PDDL domain and problem files.

    This function recursively descends into all subdirectories in the puzzle path and
    recreates the subdirectory structure in the directory of saved PDDL files.

    Args:
        pddl_path: The path to a directory in which to save the resulting PDDL files.
            This directory is created if it does not exist.
        puzzle_path: The path to a PushWorld puzzle file or to a directory
            that contains puzzle files, possibly nested in subdirectories.
        for_bfws: See `puzzle_to_pddl`.
    """
    for puzzle_file_path, pddl_file_path in map_files_with_extension(
        input_file_or_directory_path=puzzle_path,
        input_extension=PUZZLE_EXTENSION,
        output_directory_path=pddl_path,
    ):
        puzzle = PushWorldPuzzle(puzzle_file_path)
        puzzle_name = os.path.splitext(os.path.split(puzzle_file_path)[1])[0]
        pddl_domain, pddl_problem = puzzle_to_pddl(puzzle_name, puzzle, for_bfws)

        with open(pddl_file_path + DOMAIN_SUFFIX, "w") as domain_file:
            domain_file.write(pddl_domain)

        with open(pddl_file_path + PROBLEM_SUFFIX, "w") as problem_file:
            problem_file.write(pddl_problem)
