# Copyright 2021 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.
#


from fractions import Fraction
import re
import sys
import unified_planning as up
import unified_planning.environment
import unified_planning.walkers as walkers
from unified_planning.model import DurativeAction, InstantaneousAction, Fluent, Parameter, Object
from unified_planning.model.types import _UserType, _RealType, _IntType
from typing import IO, Dict, List, Optional, cast, Union
from io import StringIO

ANML_KEYWORDS = {'action', 'and', 'constant', 'duration', 'else', 'fact', 'fluent', 'function', 'goal', 'in', 'instance', 'motivated',
        'predicate', 'symbol', 'variable', 'when', 'with', 'decomposition', 'use', 'coincident',
        'comprise', 'comprises', 'contain', 'contains', 'exists', 'forall', 'implies', 'iff', 'not', 'or', 'ordered', 'unordered',
        'xor', 'UNDEFINED', 'all', 'end', 'false', 'infinity', 'object', 'start', 'true', 'boolean', 'float', 'rational', 'integer', 'string',
        'type', 'set', 'subset', 'powerset', 'intersect', 'union', 'elt'}

# The following map is used to mangle the invalid names by their class.
INITIAL_LETTER: Dict[type, str] = {InstantaneousAction: 'a', DurativeAction: 'a', Fluent: 'f', Parameter: 'p', Object: 'o'}


class ConverterToANMLString(walkers.DagWalker):
    '''Expression converter to an ANML string.'''

    def __init__(self, names_mapping: Dict[Union['up.model.Type', 'up.model.Action', 'up.model.Parameter', 'up.model.Fluent', 'up.model.Object'], str], env: 'up.environment.Environment'):
        walkers.DagWalker.__init__(self)
        self._names_mapping = names_mapping
        self.simplifier = walkers.Simplifier(env)

    def convert(self, expression):
        '''Converts the given expression to a ANML string.'''
        return self.walk(self.simplifier.simplify(expression)) #NOTE maybe the converter could remove the first and last char, if they are '(' and ')'

    def walk_exists(self, expression, args):
        assert len(args) == 1
        vars_string_gen = (f'{_get_anml_name(v.type, self._names_mapping)} {_get_anml_name(v, self._names_mapping)}' for v in expression.variables())
        return f'(exists({", ".join(vars_string_gen)}) {{ {args[0]} }})'

    def walk_forall(self, expression, args):
        assert len(args) == 1
        vars_string_gen = (f'{_get_anml_name(v.type, self._names_mapping)} {_get_anml_name(v, self._names_mapping)}' for v in expression.variables())
        return f'(forall({", ".join(vars_string_gen)}) {{ {args[0]} }})'

    def walk_variable_exp(self, expression, args):
        assert len(args) == 0
        return _get_anml_name(expression.variable(), self._names_mapping)

    def walk_and(self, expression, args):
        assert len(args) > 1
        return f'({" and ".join(args)})'

    def walk_or(self, expression, args):
        assert len(args) > 1
        return f'({" or ".join(args)})'

    def walk_not(self, expression, args):
        assert len(args) == 1
        return f'(not {args[0]})'

    def walk_implies(self, expression, args):
        assert len(args) == 2
        return f'({args[0]} implies {args[1]})'

    def walk_iff(self, expression, args):
        assert len(args) == 2
        return f'({args[0]} iff {args[1]})'

    def walk_fluent_exp(self, expression, args):
        if len(args) == 0:
            return self._names_mapping[expression.fluent()]
        else:
            return f'{self._names_mapping[expression.fluent()]}({", ".join(args)})'

    def walk_param_exp(self, expression, args):
        assert len(args) == 0
        return _get_anml_name(expression.parameter(), self._names_mapping)

    def walk_object_exp(self, expression, args):
        assert len(args) == 0
        return _get_anml_name(expression.object(), self._names_mapping)

    def walk_bool_constant(self, expression, args):
        assert len(args) == 0
        if expression.bool_constant_value():
            return 'true'
        return 'false'

    def walk_real_constant(self, expression, args):
        assert len(args) == 0
        frac = cast(Fraction, expression.constant_value())
        return f'({frac.numerator}/{frac.denominator})'

    def walk_int_constant(self, expression, args):
        assert len(args) == 0
        return str(expression.constant_value())

    def walk_plus(self, expression, args):
        assert len(args) > 1
        return f"({' + '.join(args)})"

    def walk_minus(self, expression, args):
        assert len(args) == 2
        return f'({args[0]} - {args[1]})'

    def walk_times(self, expression, args):
        assert len(args) > 1
        return f"({' * '.join(args)})"

    def walk_div(self, expression, args):
        assert len(args) == 2
        return f'({args[0]} / {args[1]})'

    def walk_le(self, expression, args):
        assert len(args) == 2
        return f'({args[0]} <= {args[1]})'

    def walk_lt(self, expression, args):
        assert len(args) == 2
        return f'({args[0]} < {args[1]})'

    def walk_equals(self, expression, args):
        assert len(args) == 2
        return f'({args[0]} == {args[1]})'


class ANMLWriter:
    '''This class can be used to write a Problem in ANML.'''

    def __init__(self, problem: 'up.model.Problem'):
        self.problem = problem

    def _write_problem(self, out: IO[str]):
        names_mapping: Dict[Union['up.model.Type', 'up.model.Action', 'up.model.Parameter', 'up.model.Fluent', 'up.model.Object'], str] = {}
        # Init names_mapping.
        names_mapping[self.problem.env.type_manager.BoolType()] = 'boolean'
        names_mapping[self.problem.env.type_manager.IntType()] = 'integer'
        names_mapping[self.problem.env.type_manager.RealType()] = 'float'
        for t in self.problem.user_types:
            ut = cast(_UserType, t)
            if _is_valid_anml_name(ut.name): # No renaming needed
                names_mapping[t] = ut.name
        for a in self.problem.actions:
            if _is_valid_anml_name(a.name): # No renaming needed
                names_mapping[a] = a.name
        for f in self.problem.fluents:
            if _is_valid_anml_name(f.name): # No renaming needed
                names_mapping[f] = f.name
        for o in self.problem.all_objects:
            if _is_valid_anml_name(o.name): # No renaming needed
                names_mapping[o] = o.name

        for t in self.problem.user_types:
            anml_type_name = _get_anml_name(t, names_mapping)
            out.write(f'type {anml_type_name}')
            if cast(_UserType, t).father is None:
                out.write(';\n')
            else:
                # For construction in the Problem, the father of a UserType is always added before the UserType itself.
                father = cast(_UserType, t).father
                assert father is not None
                assert names_mapping[father] is not None
                out.write(f' < {names_mapping[father]};\n')
        static_fluents = self.problem.get_static_fluents()
        for f in self.problem.fluents:
            parameters = [f'{_get_anml_name(ap.type, names_mapping)} {_get_anml_name(ap, names_mapping)}' for ap in f.signature]
            params_written = f'({", ".join(parameters)})' if len(parameters) > 0 else ''
            if f in static_fluents:
                out.write(f'constant {_get_anml_name(f.type, names_mapping)} {_get_anml_name(f, names_mapping)}{params_written};\n')
            else:
                out.write(f'fluent {_get_anml_name(f.type, names_mapping)} {_get_anml_name(f, names_mapping)}{params_written};\n')


        converter = ConverterToANMLString(names_mapping, self.problem.env)

        for a in self.problem.actions:
            if isinstance(a, up.model.InstantaneousAction):
                parameters = [f'{_get_anml_name(ap.type, names_mapping)} {_get_anml_name(ap, names_mapping)}' for ap in a.parameters]
                out.write(f'action {_get_anml_name(a, names_mapping)}({", ".join(parameters)}) {{\n')
                for p in a.preconditions:
                    out.write(f'   [ start ] {converter.convert(p)};\n')
                for e in a.effects:
                    out.write(f'   {self._convert_effect(e, converter, None, 3)}')
                out.write('};\n')
            elif isinstance(a, DurativeAction):
                parameters = [f'{_get_anml_name(ap.type, names_mapping)} {_get_anml_name(ap, names_mapping)}' for ap in a.parameters]
                out.write(f'action {_get_anml_name(a, names_mapping)}({", ".join(parameters)}) {{\n')
                left_bound = ' > ' if a.duration.is_left_open() else ' >= '
                right_bound = ' < ' if a.duration.is_right_open() else ' <= '
                out.write(f'   duration{left_bound}{converter.convert(a.duration.lower)} and ')
                out.write(f'duration{right_bound}{converter.convert(a.duration.upper)};\n')
                for i, cl in a.conditions.items():
                    for c in cl:
                        out.write(f'   {self._convert_anml_interval(i)} {converter.convert(c)};\n')
                for ti, el in a.effects.items():
                    for e in el:
                        out.write(f'   {self._convert_effect(e, converter, ti, 3)}')
                out.write('};\n')
            else:
                raise NotImplementedError

        for t in self.problem.user_types: # Define objects
            obj_names = [_get_anml_name(o, names_mapping) for o in self.problem.objects(t) if o.type == t]
            if len(obj_names) > 0:
                out.write(f'instance {_get_anml_name(t, names_mapping)} {", ".join(obj_names)};\n')

        for fe, v in self.problem.initial_values.items():
            assert fe.is_fluent_exp()
            if fe.fluent() in static_fluents:
                out.write(f'{converter.convert(fe)} := {converter.convert(v)};\n')
            else:
                out.write(f'[ start ] {converter.convert(fe)} := {converter.convert(v)};\n')

        for ti, el in self.problem.timed_effects.items():
            for e in el:
                out.write(self._convert_effect(e, converter, ti))

        for g in self.problem.goals:
            out.write(f'[ end ] {converter.convert(g)};\n')

        for i, gl in self.problem.timed_goals.items():
            for g in gl:
                out.write(f'{self._convert_anml_interval(i)} {converter.convert(g)};\n')

    def print_problem(self):
        '''Prints to std output the ANML problem.'''
        self._write_problem(sys.stdout)

    def get_problem(self) -> str:
        '''Returns the ANML problem.'''
        out = StringIO()
        self._write_problem(out)
        return out.getvalue()

    def write_problem(self, filename: str):
        '''Dumps to file the ANML problem.'''
        with open(filename, 'w') as f:
            self._write_problem(f)

    def _convert_effect(self, effect: 'up.model.Effect',
                        converter: ConverterToANMLString,
                        timing: 'up.model.Timing' = None,
                        spaces_from_left: int = 0) -> str:
        results: List[str] = []
        anml_timing = self._convert_anml_timing(timing) if timing is not None else 'start'
        if effect.is_conditional():
            results.append(f'when [ {anml_timing} ] {converter.convert(effect.condition)}\n{spaces_from_left*" "}{{')
        results.append(f'[ {anml_timing} ] ')
        results.append(converter.convert(effect.fluent))
        if effect.is_assignment():
            results.append(' := ')
        elif effect.is_increase():
            results.append(f' :increase ')
        elif effect.is_decrease():
            results.append(f' :decrease ')
        else:
            raise NotImplementedError
        results.append(f'{converter.convert(effect.value)};\n')
        if effect.is_conditional():
            results.append(f'{spaces_from_left*" "}}}\n')
        return ''.join(results)

    def _convert_anml_timing(self, timing: 'up.model.Timing') -> str:
        time = 'start' if timing.is_from_start() else 'end'
        if timing.delay > 0:
            return f'{time} + {str(timing.delay)}'
        elif timing.delay == 0:
            return time
        else: #timing.delay < 0
            return f'{time} - {str(timing.delay * (-1))}'

    def _convert_anml_interval(self, interval: 'up.model.TimeInterval') -> str:
        left_bracket = '(' if interval.is_left_open() else '['
        right_bracket = ')' if interval.is_left_open() else ']'
        return f'{left_bracket} {self._convert_anml_timing(interval.lower)}, {self._convert_anml_timing(interval.upper)} {right_bracket}'

def _is_valid_anml_name(name: str) -> bool:
    regex = re.compile(r'^[a-zA-Z][a-zA-Z0-9_]*')
    if re.match(regex, name) is None or name in ANML_KEYWORDS: # If the name does not start with an alphabetic char or is a keyword
        return False
    return True

def _get_anml_valid_name(item: Union['up.model.Type', 'up.model.Action', 'up.model.Parameter', 'up.model.Fluent', 'up.model.Object']) -> str:
    '''This function returns a valid ANML name.'''
    if isinstance(item, up.model.Type):
        assert item.is_user_type()
        name = cast(_UserType, item).name
    else:
        name = item.name
    regex = re.compile(r'^[a-zA-Z]+.*')
    if re.match(regex, name) is None: # If the name does not start with an alphabetic char, we make it start with one.
        name = f'{INITIAL_LETTER.get(type(item), "x")}_{name}'
    name = re.sub('[^0-9a-zA-Z_]', '_', name) #Substitute non-valid elements with "_"
    while name in ANML_KEYWORDS: # If the name is in the keywords, apply an underscore at the end until it is not a keyword anymore.
        name = f'{name}_'
    return name

def _get_anml_name(item: Union['up.model.Type', 'up.model.Action', 'up.model.Parameter', 'up.model.Fluent', 'up.model.Object'],
    names_mapping: Dict[Union['up.model.Type', 'up.model.Action', 'up.model.Parameter', 'up.model.Fluent', 'up.model.Object'], str]) -> str:
    '''Important note: This method updates the names_mapping '''
    new_name: Optional[str] = names_mapping.get(item, None)
    if new_name is None: # The type is not in the dictionary, so his name must be added
        if isinstance(item, up.model.Type) and item.is_int_type():
            num_type = cast(_IntType, item)
            left_bound = '(-infinity' if num_type.lower_bound is None else f'[{str(num_type.lower_bound)}'
            right_bound = 'infinity)' if num_type.upper_bound is None else f'{str(num_type.upper_bound)}]'
            new_name = f'integer {left_bound}, {right_bound}'
        elif isinstance(item, up.model.Type) and item.is_real_type():
            num_real_type = cast(_RealType, item)
            if num_real_type.lower_bound is None:
                left_bound = '(-infinity'
            elif num_real_type.lower_bound.denominator == 1:
                left_bound = f'[{str(num_real_type.lower_bound)}.0'
            else:
                left_bound = f'[{str(num_real_type.lower_bound)}'
            if num_real_type.upper_bound is None:
                right_bound = 'infinity)'
            elif num_real_type.upper_bound.denominator == 1:
                right_bound = f'{str(num_real_type.upper_bound)}.0]'
            else:
                right_bound = f'{str(num_real_type.upper_bound)}]'
            new_name = f'float {left_bound}, {right_bound}'
        else: # We mangle the name and get a fresh one
            new_name = _get_anml_valid_name(item)
            test_name = new_name # Init values
            count = 0
            while test_name in names_mapping.values(): # Loop until we find a fresh name
                test_name = f'{new_name}_{str(count)}'
                count += 1
            new_name = test_name
            assert _is_valid_anml_name(new_name)
        names_mapping[item] = new_name # Once a fresh valid name is found, update the map.
    return cast(str, new_name)
