#
# This file is part of pyperplan.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>
#

"""
Classes and functions that traverse the PDDL abstract syntax tree (AST)
generated by pddl_parser.py, thereby creating the PDDL data structure.
"""

import pddl.pddl as pddl


class SemanticError(Exception):
    """Exception indicating an error during traversal of AST."""

    def __init__(self, value):
        """Constructor of SemanticError.

        Keyword arguments:
        value -- the error message
        """
        self.value = value

    def __str__(self):
        return repr(self.value)


class Visitable:
    """
    The Visitable class is part of the Visitor Pattern. Every AST node created
    by the Parser derives from this class.

    The accept-method calls the appropriate method of the visitor.
    """
    def __init__(self, vname=None):
        """Constructor of the Visitable.

        Keyword arguments:
        vname -- the name of the of the callback that will be called on the
                 visitor instance
        """
        self._visitorName = vname

    def accept(self, visitor):
        if self._visitorName == None:
            raise ValueError('Error: visit method of uninitialized visitor '
                             'called!')
        # get the appropriate method of the visitor instance
        m = getattr(visitor, self._visitorName)
        # ensure that the method is callable
        if not hasattr(m, '__call__'):
            raise ValueError('Error: cannot call undefined method: %s on '
                             'visitor' % self._visitorName)
        # and finally call the callback
        m(self)


class PDDLVisitor:
    """
    The standard tree traversal PDDL Visitor from which other Visitors derive.

    In each node, the Visitor just applies itself to all its children.
    """
    def visit_domain_def(self, node):
        node.requirements.accept(self)
        if node.types != None:
            for t in node.types:
                t.accept(self)
        if node.constants != None:
            for c in node.constants:
                c.accept(self)
        node.predicates.accept(self)
        if node.actions != None:
            for a in node.actions:
                a.accept(self)

    def visit_problem_def(self, node):
        for o in node.objects:
            o.accept(self)
        node.init.accept(self)
        node.goal.accept(self)

    def visit_predicates_stmt(self, node):
        for p in node.predicates:
            p.accept(self)

    def visit_action_stmt(self, node):
        for p in node.parameters:
            p.accept(self)
        node.precond.accept(self)
        node.effect.accept(self)

    def visit_formula(self, node):
        for c in node.children:
            c.accept(self)

    def visit_type(self, node):
        return  # nothing to do right now

    def visit_effect_stmt(self, node):
        node.formula.accept(self)

    def visit_precondition_stmt(self, node):
        node.formula.accept(self)

    def visit_requirements_stmt(self, node):
        for k in node.keywords:
            k.accept(self)

    def visit_predicate(self, node):
        for p in node.parameters:
            p.accept(self)

    def visit_variable(self, node):
        pass

    def visit_init_stmt(self, node):
        for p in node.predicates:
            p.accept(self)

    def visit_goal_stmt(self, node):
        node.formula.accept(self)

    def visit_predicate_instance(self, node):
        return  # nothing to do right now

    def visit_object(self, node):
        return  # nothing to do right now

    def visit_keyword(self, node):
        return  # nothing to do right now


class TraversePDDLDomain(PDDLVisitor):
    """The PDDL-domain Visitor.

    Expands the functionality of the PDDLVisitor to traversal of ASTs that
    represent a PDDL-domain file. It results in the PDDL data structure
    (pddl.py) representation of the domain file.
    """

    def get_in(self, node):
        """
        Helper method to access a global hash in which information for each
        node in the AST can be stored.
        """
        return self._nodeHash[node]

    def set_in(self, node, val):
        """
        Helper method to write a global hash in which information for each node
        in the AST can be stored.
        """
        self._nodeHash[node] = val

    def __init__(self):
        self._types = dict()
        self._predicates = dict()
        self._nodeHash = dict()
        self._requirements = set()
        self._actions = dict()
        self.domain = None
        self._objectType = pddl.Type('object', None)
        self._constants = dict()

    def visit_domain_def(self, node):
        """Visits a PDDL domain definition."""
        explicitObjectDef = False

        # Requirements statement is optional.
        if node.requirements:
            node.requirements.accept(self)

        # Visit all type definitions.
        if node.types != None:
            for t in node.types:
                if t.name == 'object':
                    explicitObjectDef = True
                t.accept(self)
                type = self.get_in(t)
                self._types[type.name] = type
        # Add the default object type to the type definitions,
        # if it was not explicitly created.
        if not explicitObjectDef:
            self._types['object'] = self._objectType

        # Link all types to their parent types directly.
        for t in self._types.values():
            # Object type has no parent.
            if t.name == 'object':
                continue
            if not t.parent in self._types:
                raise SemanticError('Error unknown parent type: ' + t.parent)
            t.parent = self._types[t.parent]

        # Visit all predicates.
        node.predicates.accept(self)

        # Visit all actions.
        if node.actions != None:
            for a in node.actions:
                a.accept(self)
                action = self.get_in(a)
                if action.name in self._actions:
                    raise SemanticError('Error: action with name ' +
                                        action.name +
                                        ' has already been defined')
                self._actions[action.name] = action

        # Visit all constants.
        if node.constants != None:
            for c in node.constants:
                c.accept(self)

        # Finally generate PDDL domain data structure.
        self.domain = pddl.Domain(node.name, self._types, self._predicates,
                                  self._actions, self._constants)

    def visit_object(self, node):
        """Visits a PDDL object definition."""
        type_name = node.typeName
        if type_name == None:
            type_name = 'object'
        if not type_name in self._types:
            raise SemanticError('Error: unknown type ' + type_name +
                                ' used in object definition!')
        if node.name in self._constants:
            raise SemanticError('Error: multiple defines of object with '
                                'name ' + node.name)
        # Add constant with its corresponding type to the constants dict.
        self._constants[node.name] = self._types[type_name]

    def visit_type(self, node):
        """Visits a PDDL type definition."""
        # Store matching parent type in node
        # (if none is given, it's always object)
        if node.parent == None:
            self.set_in(node, pddl.Type(node.name, 'object'))
        else:
            self.set_in(node, pddl.Type(node.name, node.parent))

    def visit_requirements_stmt(self, node):
        """Visits a PDDL requirement statement."""
        # Visit all requirement keywords...
        for k in node.keywords:
            k.accept(self)
            requirementName = self.get_in(k)
            # ... and add them to the requirement list.
            self._requirements.add(requirementName)

    def visit_keyword(self, node):
        """Visits a PDDL keyword."""
        # Nothing to do but to store its name in the node.
        self.set_in(node, node.name)

    def visit_predicates_stmt(self, node):
        """Visits a PDDL predicate statement."""
        # Visit all predicates in the predicate statement.
        for p in node.predicates:
            p.accept(self)
            predicate = self.get_in(p)
            # Check for duplicate predicate definitions.
            if predicate.name in self._predicates:
                raise SemanticError('Error predicate with name ' +
                                    predicate.name +
                                    ' has already been defined')
            # Add to predicate list.
            self._predicates[predicate.name] = predicate

    def visit_predicate(self, node):
        """Visits a PDDL predicate."""
        signature = list()
        # Visit all predicate parameters.
        for v in node.parameters:
            v.accept(self)
            signatureTuple = self.get_in(v)
            # Append each parameter to the predicate signature.
            signature.append(signatureTuple)
        # Create new PDDL predicate and store it in node.
        self.set_in(node, pddl.Predicate(node.name, signature))

    def visit_variable(self, node):
        """Visits a PDDL variable."""
        # If there is no type given, its always of type 'object'.
        if not node.typed:
            self.set_in(node, (node.name, [self._types['object']]))
        else:
            # Visit all type declarations of the variable.
            typelist = list()
            for t in node.types:
                # Check whether they have been defined.
                if not t in self._types:
                    raise SemanticError('Error unknown type ' + t +
                                        ' used in predicate definition')
                typelist.append(self._types[t])
            # Store variable information (var_name, tuple(types)) in node.
            self.set_in(node, (node.name, tuple(typelist)))

    def visit_action_stmt(self, node):
        """Visits a PDDL action statement."""
        signature = list()
        # Visit all parameters and create signature.
        for v in node.parameters:
            v.accept(self)
            signatureTuple = self.get_in(v)
            signature.append(signatureTuple)

        # Visit the precondition statement.
        node.precond.accept(self)
        precond = self.get_in(node.precond)

        # Visit the effect statement.
        node.effect.accept(self)
        effect = self.get_in(node.effect)

        # Create new PDDL action and store in node.
        self.set_in(node, pddl.Action(node.name, signature, precond, effect))

    def add_precond(self, precond, c):
        """Helper function for visit_precondition_stmt.

        Keyword arguments:
        precond -- a list of preconditions
        c -- the formula representing a precondition we want to add to the list
        """
        from pddl.parser import Variable
        predDef = self._predicates[c.key]
        signature = list()
        count = 0
        # Check for correct number of arguments.
        if len(c.children) != len(predDef.signature):
            raise SemanticError('Error: wrong number of arguments for '
                                'predicate ' + c.key + ' in precondition of '
                                'action')
        # Apply to all arguments.
        for v in c.children:
            if isinstance(v.key, Variable):
                signature.append((v.key.name, predDef.signature[count][1]))
            else:
                signature.append((v.key, predDef.signature[count][1]))
            count += 1

        # Add predicate to precondition list.
        precond.append(pddl.Predicate(c.key, signature))

    def visit_precondition_stmt(self, node):
        """ Visits a PDDL precondition statement."""
        precond = list()
        formula = node.formula
        # For now we only allow and in the precondition.
        if formula.key == 'and':
            # Apply to all predicates in precondition.
            for c in formula.children:
                if not isinstance(c.key, str):
                    raise SemanticError('Error predicate with non str key: ' +
                                        ''.join([c2.key.name + ' '
                                                for c2 in formula.children]))
                # Check whether predicate was defined.
                if not c.key in self._predicates:
                    raise SemanticError('Error unknown predicate ' + c.key +
                                        ' used in precondition of action')
                # Call helper.
                self.add_precond(precond, c)
        else:
            # If not 'and' we only allow a single predicate in precondition.
            if not formula.key in self._predicates:
                raise SemanticError('Error: predicate in precondition is not '
                                    'in CNF')
            # Call helper.
            self.add_precond(precond, formula)
        self.set_in(node, precond)

    def add_effect(self, effect, c):
        """Helper function for visit_effect_stmt.

        Keyword arguments:
        effect -- instance of the effect data structure
        c -- the formula representing the effect that we want to add to the
             addlist or dellist
        """
        # Needed for instance check.
        from pddl.parser import Variable
        nextPredicate = None
        isNegative = False
        if c.key == 'not':
            # This is a negative effect, only one child allowed.
            if len(c.children) != 1:
                raise SemanticError('Error not statement with multiple '
                                    'children in effect of action')
            nextPredicate = c.children[0]
            isNegative = True
        else:
            nextPredicate = c
        # Check whether predicate was defined previously.
        if not nextPredicate.key in self._predicates:
            raise SemanticError('Error: unknown predicate %s used in effect '
                                'of action' % nextPredicate.key)
        if nextPredicate == None:
            raise SemanticError('Error: NoneType predicate used in effect of '
                                'action')
        predDef = self._predicates[nextPredicate.key]
        signature = list()
        count = 0
        # Check whether predicate is used with the correct signature.
        if len(nextPredicate.children) != len(predDef.signature):
            raise SemanticError('Error: wrong number of arguments for '
                                'predicate ' + nextPredicate.key +
                                ' in effect of action')
        # Apply to all parameters.
        for v in nextPredicate.children:
            if isinstance(v.key, Variable):
                signature.append((v.key.name, predDef.signature[count][1]))
            else:
                signature.append((v.key, predDef.signature[count][1]))
            count += 1

        # Add a new effect to the positive or negative effects respectively.
        if isNegative:
            effect.dellist.add(pddl.Predicate(nextPredicate.key, signature))
        else:
            effect.addlist.add(pddl.Predicate(nextPredicate.key, signature))

    def visit_effect_stmt(self, node):
        """ Visits a PDDL effect statement."""
        formula = node.formula
        effect = pddl.Effect()
        # For now we only allow 'and' in the effect.
        if formula.key == 'and':
            for c in formula.children:
                # Call helper.
                self.add_effect(effect, c)
        else:
            # Call helper.
            self.add_effect(effect, formula)
        # Store effect in node.
        self.set_in(node, effect)


class TraversePDDLProblem(PDDLVisitor):
    """The PDDL-problem Visitor.

    Expands the functionality of the PDDLVisitor to traversal of ASTs that
    represent a PDDL-problem file. It results in the PDDL data structure
    (pddl.py) representation of the problem file.
    """
    def get_in(self, node):
        """
        Helper method to access a global hash in which information for each
        node in the AST can be stored.
        """
        return self._nodeHash[node]

    def set_in(self, node, val):
        """
        Helper method to write a global hash in which information for each node
        in the AST can be stored.
        """
        self._nodeHash[node] = val

    def get_problem(self):
        """Getter for the resulting pddl-problem data structure."""
        return self._problemDef
    problemDef = property(get_problem)

    def __init__(self, domain):
        """Constructor for pddl-problem visitor.

        Keyword arguments:
        domain -- the coressponding pddl-domain datastrucutre
        """
        self._domain = domain
        self._nodeHash = dict()
        self._objects = dict()
        self._problemDef = None

    def visit_problem_def(self, node):
        """Visits a PDDL-problem definition."""
        # Check whether the in the problem file referenced domain name matches
        # the supplied domain data structure.
        if node.domainName != self._domain.name:
            raise SemanticError('Error trying to parse problem file with '
                                'domain: %s together with a domain file that '
                                'specifies domain: %s' %
                                (node.domainName, self._domain.name))
        # Apply to all object definitions.
        for o in node.objects:
            o.accept(self)

        # Apply to the initial state definition.
        node.init.accept(self)
        init_list = self.get_in(node.init)

        # Apply to the goal state definition.
        node.goal.accept(self)
        goal_list = self.get_in(node.goal)

        # Create the problem data structure.
        self._problemDef = pddl.Problem(node.name, self._domain, self._objects,
                                        init_list, goal_list)

    def visit_object(self, node):
        """ Visits a PDDL-problem object definition."""
        type_def = None
        # Check for multiple definition of objects.
        if node.name in self._objects:
            raise SemanticError('Error multiple defines of object with name ' +
                                node.name)
        # Untyped objects get the standard type 'object'.
        if node.typeName == None:
            type_def = self._domain.types['object']
        else:
            # Check whether used type was introduced in domain file.
            if not node.typeName in self._domain.types:
                raise SemanticError('Error: unknown type ' + node.typeName +
                                    ' used in object definition!')
            type_def = self._domain.types[node.typeName]
        self._objects[node.name] = type_def

    def visit_init_stmt(self, node):
        """ Visits a PDDL-problem initial state statement."""
        initList = list()
        # Apply to all predicates in the statement.
        for p in node.predicates:
            p.accept(self)
            pred = self.get_in(p)
            initList.append(pred)
        self.set_in(node, initList)

    def add_goal(self, goal, c):
        """Helper function for visit_goal_stmt.

        Keyword arguments:
        goal -- a list of goals
        c -- a formula representing a goal we want to add to the goal list
        """
        # Check whether predicate was introduced in domain file.
        if not c.key in self._domain.predicates:
            raise SemanticError('Error: unknown predicate ' + c.key +
                                ' in goal definition')
        # Get predicate from the domain data structure.
        predDef = self._domain.predicates[c.key]
        signature = list()
        count = 0
        # Check whether the predicate uses the correct signature.
        if len(c.children) != len(predDef.signature):
            raise SemanticError('Error: wrong number of arguments for '
                                'predicate ' + c.key + ' in goal')
        for v in c.children:
            signature.append((v.key, predDef.signature[count][1]))
            count += 1
        # Add the predicate to the goal.
        goal.append(pddl.Predicate(c.key, signature))

    def visit_goal_stmt(self, node):
        """ Visits a PDDL-problem goal state statement."""
        formula = node.formula
        goal = list()
        # For now we only allow 'and' in the goal.
        if formula.key == 'and':
            for c in formula.children:
                if not isinstance(c.key, str):
                    raise SemanticError('Error predicate with non str key: ' +
                                        ''.join([c2.key.name + ' '
                                        for c2 in formula.children]))
                # Call helper.
                self.add_goal(goal, c)
        else:
            # Only a single predicate is allowed then (s.a.)
            if not formula.key in self._domain.predicates:
                raise SemanticError('Error: predicate in goal definition is '
                                    'not in CNF')
            # Call helper.
            self.add_goal(goal, formula)
        self.set_in(node, goal)

    def visit_predicate_instance(self, node):
        """ Visits a PDDL-problem predicate instance."""
        signature = list()
        # Visit all parameters.
        for o in node.parameters:
            o_type = None
            # Check whether predicate was introduced in objects or domain
            # constants.
            if not (o in self._objects or o in self._domain.constants):
                raise SemanticError('Error: object ' + o + ' referenced in '
                                    'problem definition - but not defined')
            elif o in self._objects:
                o_type = self._objects[o]
            elif o in self._domain.constants:
                o_type = self._domain.constants[o]
            signature.append((o, (o_type)))
        self.set_in(node, pddl.Predicate(node.name, signature))
