# Copyright 2021-2023 AIPlan4EU project
#
# 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 module defines the state invariants remover class."""


import unified_planning as up
import unified_planning.engines as engines
from unified_planning.engines.mixins.compiler import CompilationKind, CompilerMixin
from unified_planning.engines.results import CompilerResult
from unified_planning.model import Problem, ProblemKind
from unified_planning.engines.compilers.utils import (
    add_invariant_condition_apply_function_to_problem_expressions,
    replace_action,
)
from typing import Optional
from functools import partial


class StateInvariantsRemover(engines.engine.Engine, CompilerMixin):
    """
    Global constraints remover class: this class offers the capability
    to transform a :class:`~unified_planning.model.Problem` with state innvariants
    into a `Problem` without state innvariants.
    This capability is offered by the :meth:`~unified_planning.engines.compilers.GlobalConstraintsRemover.compile`
    method, that returns a :class:`~unified_planning.engines.CompilerResult` in which the :meth:`problem <unified_planning.engines.CompilerResult.problem>` field
    is the compiled Problem.

    This is done by setting the state innvariants as action's preconditions and goals, so that at every step the constraints are checked.

    This `Compiler` supports only the the `STATE_INVARIANTS_REMOVING` :class:`~unified_planning.engines.CompilationKind`.
    """

    def __init__(self):
        engines.engine.Engine.__init__(self)
        CompilerMixin.__init__(self, CompilationKind.STATE_INVARIANTS_REMOVING)

    @property
    def name(self):
        return "gcrm"

    @staticmethod
    def METHOD_NAME() -> ProblemKind:
        METHOD_NAME = ProblemKind()
        METHOD_NAME.set_problem_class("ACTION_BASED")
        METHOD_NAME.set_typing("FLAT_TYPING")
        METHOD_NAME.set_typing("HIERARCHICAL_TYPING")
        METHOD_NAME.set_numbers("CONTINUOUS_NUMBERS")
        METHOD_NAME.set_numbers("DISCRETE_NUMBERS")
        METHOD_NAME.set_numbers("BOUNDED_TYPES")
        METHOD_NAME.set_problem_type("SIMPLE_NUMERIC_PLANNING")
        METHOD_NAME.set_problem_type("GENERAL_NUMERIC_PLANNING")
        METHOD_NAME.set_fluents_type("NUMERIC_FLUENTS")
        METHOD_NAME.set_fluents_type("OBJECT_FLUENTS")
        METHOD_NAME.set_conditions_kind("NEGATIVE_CONDITIONS")
        METHOD_NAME.set_conditions_kind("DISJUNCTIVE_CONDITIONS")
        METHOD_NAME.set_conditions_kind("EQUALITIES")
        METHOD_NAME.set_conditions_kind("EXISTENTIAL_CONDITIONS")
        METHOD_NAME.set_conditions_kind("UNIVERSAL_CONDITIONS")
        METHOD_NAME.set_effects_kind("CONDITIONAL_EFFECTS")
        METHOD_NAME.set_effects_kind("INCREASE_EFFECTS")
        METHOD_NAME.set_effects_kind("DECREASE_EFFECTS")
        METHOD_NAME.set_effects_kind("STATIC_FLUENTS_IN_BOOLEAN_ASSIGNMENTS")
        METHOD_NAME.set_effects_kind("STATIC_FLUENTS_IN_NUMERIC_ASSIGNMENTS")
        METHOD_NAME.set_effects_kind("STATIC_FLUENTS_IN_OBJECT_ASSIGNMENTS")
        METHOD_NAME.set_effects_kind("FLUENTS_IN_BOOLEAN_ASSIGNMENTS")
        METHOD_NAME.set_effects_kind("FLUENTS_IN_NUMERIC_ASSIGNMENTS")
        METHOD_NAME.set_effects_kind("FLUENTS_IN_OBJECT_ASSIGNMENTS")
        METHOD_NAME.set_effects_kind("FORALL_EFFECTS")
        METHOD_NAME.set_time("CONTINUOUS_TIME")
        METHOD_NAME.set_time("DISCRETE_TIME")
        METHOD_NAME.set_time("INTERMEDIATE_CONDITIONS_AND_EFFECTS")
        METHOD_NAME.set_time("EXTERNAL_CONDITIONS_AND_EFFECTS")
        METHOD_NAME.set_time("TIMED_EFFECTS")
        METHOD_NAME.set_time("TIMED_GOALS")
        METHOD_NAME.set_time("DURATION_INEQUALITIES")
        METHOD_NAME.set_simulated_entities("SIMULATED_EFFECTS")
        METHOD_NAME.set_constraints_kind("STATE_INVARIANTS")
        METHOD_NAME.set_constraints_kind("TRAJECTORY_CONSTRAINTS")
        METHOD_NAME.set_quality_metrics("ACTIONS_COST")
        METHOD_NAME.set_actions_cost_kind("STATIC_FLUENTS_IN_ACTIONS_COST")
        METHOD_NAME.set_actions_cost_kind("FLUENTS_IN_ACTIONS_COST")
        METHOD_NAME.set_quality_metrics("PLAN_LENGTH")
        METHOD_NAME.set_quality_metrics("OVERSUBSCRIPTION")
        METHOD_NAME.set_quality_metrics("TEMPORAL_OVERSUBSCRIPTION")
        METHOD_NAME.set_quality_metrics("MAKESPAN")
        METHOD_NAME.set_quality_metrics("FINAL_VALUE")
        return METHOD_NAME

    @staticmethod
    def supports(problem_kind):
        return problem_kind <= StateInvariantsRemover.METHOD_NAME()

    @staticmethod
    def supports_compilation(compilation_kind: CompilationKind) -> bool:
        return compilation_kind == CompilationKind.STATE_INVARIANTS_REMOVING

    @staticmethod
    def resulting_problem_kind(
        problem_kind: ProblemKind, compilation_kind: Optional[CompilationKind] = None
    ) -> ProblemKind:
        new_kind = ProblemKind(problem_kind.features)
        if new_kind.has_state_invariants():
            new_kind.unset_constraints("STATE_INVARIANTS")
        return new_kind

    def _compile(
        self,
        problem: "up.model.AbstractProblem",
        compilation_kind: "up.engines.CompilationKind",
    ) -> CompilerResult:
        """
        Takes an instance of a :class:`~unified_planning.model.Problem` and the wanted :class:`~unified_planning.engines.CompilationKind`
        and returns a :class:`~unified_planning.engines.results.CompilerResult` where the :meth:`problem<unified_planning.engines.results.CompilerResult.problem>`
        field does not have state innvariants.

        :param problem: The instance of the :class:`~unified_planning.model.Problem` that must be returned without state innvariants.
        :param compilation_kind: The :class:`~unified_planning.engines.CompilationKind` that must be applied on the given problem;
            only :class:`~unified_planning.engines.CompilationKind.STATE_INVARIANTS_REMOVING` is supported by this compiler
        :return: The resulting :class:`~unified_planning.engines.results.CompilerResult` data structure.
        """
        assert isinstance(problem, Problem)
        env = problem.environment
        em = env.expression_manager
        new_problem = Problem(f"{problem.name}_{self.name}", env)
        new_problem.add_objects(problem.all_objects)
        new_problem.add_fluents(problem.fluents)

        new_to_old = add_invariant_condition_apply_function_to_problem_expressions(
            problem,
            new_problem,
            em.And(problem.state_invariants).simplify(),
        )

        # Add to the problem all the trajectory constraints that are not an always
        new_traj_constraints = []
        for tc in new_problem.trajectory_constraints:
            if tc.is_and():
                for a in tc.args:
                    if not a.is_always():
                        new_traj_constraints.append(a)
            elif tc.is_forall():
                if not tc.arg(0).is_always():
                    new_traj_constraints.append(tc)
            elif not tc.is_always():
                new_traj_constraints.append(tc)
        new_problem.clear_trajectory_constraints()
        for ntc in new_traj_constraints:
            new_problem.add_trajectory_constraint(ntc)

        return CompilerResult(
            new_problem, partial(replace_action, map=new_to_old), self.name
        )