# (C) Copyright 2018- ECMWF.
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.

from sys import intern
import re
import pytools.lex
from pymbolic.parser import Parser as ParserBase
from pymbolic.mapper import Mapper
import pymbolic.primitives as pmbl
from pymbolic.parser import (
    _openpar, _closepar, _minus, FinalizedTuple, _PREC_UNARY,
    _PREC_TIMES, _PREC_PLUS, _PREC_CALL, _times, _plus
)
try:
    from fparser.two.Fortran2003 import Intrinsic_Name

    FORTRAN_INTRINSIC_PROCEDURES = Intrinsic_Name.function_names
    """list of intrinsic fortran procedure(s) names"""
except ImportError:
    FORTRAN_INTRINSIC_PROCEDURES = ()

from loki.expression import symbols as sym, operations as sym_ops
from loki.tools.util import CaseInsensitiveDict
from loki.types.scope import Scope

__all__ = ['ExpressionParser', 'parse_expr', 'FORTRAN_INTRINSIC_PROCEDURES']


class PymbolicMapper(Mapper):
    """
    Pymbolic expression to Loki expression mapper.

    Convert pymbolic expressions to Loki expressions.
    """
    # pylint: disable=abstract-method,unused-argument

    def map_product(self, expr, *args, **kwargs):
        children = tuple(self.rec(child, *args, **kwargs) for child in expr.children)
        if isinstance(expr, sym_ops.ParenthesisedMul):
            return sym_ops.ParenthesisedMul(children)
        return sym.Product(children)

    def map_sum(self, expr, *args, **kwargs):
        children = tuple(self.rec(child, *args, **kwargs) for child in expr.children)
        if isinstance(expr, sym_ops.ParenthesisedAdd):
            return sym_ops.ParenthesisedAdd(children)
        return sym.Sum(children)

    def map_power(self, expr, *args, **kwargs):
        base=self.rec(expr.base, *args, **kwargs)
        exponent=self.rec(expr.exponent, *args, **kwargs)
        if isinstance(expr, sym_ops.ParenthesisedPow):
            return sym_ops.ParenthesisedPow(base=base, exponent=exponent)
        return sym.Power(base=base, exponent=exponent)

    def map_quotient(self, expr, *args, **kwargs):
        numerator=self.rec(expr.numerator, *args, **kwargs)
        denominator=self.rec(expr.denominator, *args, **kwargs)
        if isinstance(expr, sym_ops.ParenthesisedDiv):
            return sym_ops.ParenthesisedDiv(numerator=numerator, denominator=denominator)
        return sym.Quotient(numerator=numerator, denominator=denominator)

    def map_comparison(self, expr, *args, **kwargs):
        return sym.Comparison(left=self.rec(expr.left, *args, **kwargs),
                operator=expr.operator,
                right=self.rec(expr.right, *args, **kwargs))

    def map_logical_and(self, expr, *args, **kwargs):
        return sym.LogicalAnd(tuple(self.rec(child, *args, **kwargs) for child in expr.children))

    def map_logical_or(self, expr, *args, **kwargs):
        return sym.LogicalOr(tuple(self.rec(child, *args, **kwargs) for child in expr.children))

    def map_logical_not(self, expr, *args, **kwargs):
        return sym.LogicalNot(self.rec(expr.child, *args, **kwargs))

    def map_constant(self, expr, *args, **kwargs):
        if expr == -1:
            return expr
        if isinstance(expr, (sym.FloatLiteral, sym.IntLiteral, sym.StringLiteral, sym.LogicLiteral)):
            if isinstance(expr, sym.IntLiteral) and expr.value < 0:
                return sym.Product((-1, sym.IntLiteral(abs(expr.value))))
            return expr
        if isinstance(expr, bool):
            return sym.LogicLiteral('true' if expr else 'false')
        return sym.Literal(expr)

    map_logic_literal = map_constant
    map_string_literal = map_constant
    map_intrinsic_literal = map_constant

    map_int_literal = map_constant

    map_float_literal = map_int_literal
    map_variable_symbol = map_constant
    map_deferred_type_symbol = map_constant

    def map_meta_symbol(self, expr, *args, **kwargs):
        return sym.Variable(name=str(expr.name))
    map_scalar = map_meta_symbol
    map_array = map_meta_symbol

    def map_slice(self, expr, *args, **kwargs):
        children = tuple(self.rec(child, *args, **kwargs) if child is not None else child for child in expr.children)
        if len(children) == 1 and children[0] is None:
            # this corresponds to ':' (sym.RangeIndex((None, None)))
            children = (None, None)
        return sym.RangeIndex(children)

    map_range = map_slice
    map_range_index = map_slice
    map_loop_range = map_slice

    def map_variable(self, expr, *args, **kwargs):
        parent = kwargs.pop('parent', None)
        return sym.Variable(name=expr.name, parent=parent)

    def map_algebraic_leaf(self, expr, *args, **kwargs):
        if str(expr).isnumeric():
            return self.map_constant(expr)
        if isinstance(expr, pmbl.Call):
            if expr.function.name.lower() in ('real', 'int'):
                return sym.Cast(expr.function.name, [self.rec(param, *args, **kwargs) for param in expr.parameters][0])
            if expr.function.name.upper() in FORTRAN_INTRINSIC_PROCEDURES:
                return sym.InlineCall(function=sym.Variable(name=expr.function.name),
                        parameters=tuple(self.rec(param, *args, **kwargs) for param in expr.parameters))
            parent = kwargs.pop('parent', None)
            dimensions = tuple(self.rec(param, *args, **kwargs) for param in expr.parameters)
            if not dimensions:
                return sym.InlineCall(function=sym.Variable(name=expr.function.name, parent=parent),
                        parameters=dimensions)
            return sym.Variable(name=expr.function.name, parent=parent,
                    dimensions=tuple(self.rec(param, *args, **kwargs) for param in expr.parameters))
        try:
            return self.map_variable(expr, *args, **kwargs)
        except Exception as e:
            print(f"Exception: {e}")
            return expr

    def map_call_with_kwargs(self, expr, *args, **kwargs):
        parent = kwargs.pop('parent', None)
        name = sym.Variable(name=expr.function.name, parent=parent)
        parameters = tuple(self.rec(param, *args, **kwargs) for param in expr.parameters)
        kw_parameters = {key: self.rec(value, *args, **kwargs) for key, value\
                in CaseInsensitiveDict(expr.kw_parameters).items()}
        if expr.function.name.lower() in ('real', 'int'):
            return sym.Cast(name, parameters, kind=kw_parameters['kind'])
        return sym.InlineCall(function=name, parameters=parameters, kw_parameters=kw_parameters)

    def map_tuple(self, expr, *args, **kwargs):
        return tuple(self.rec(elem, *args, **kwargs) for elem in expr)

    def map_list(self, expr, *args, **kwargs):
        return sym.LiteralList([self.rec(elem, *args, **kwargs) for elem in expr])

    def map_remainder(self, expr, *args, **kwargs):
        # this should never happen as '%' is overwritten to represent derived types
        raise NotImplementedError

    def map_lookup(self, expr, *args, **kwargs):
        # construct derived type(s) variables
        parent = kwargs.pop('parent', None)
        parent = self.rec(expr.aggregate, parent=parent)
        return self.rec(expr.name, parent=parent)


class ExpressionParser(ParserBase):
    """
    String Parser based on :any:`pymbolic.parser.Parser` for
    parsing expressions from strings.

    The Loki String Parser utilises and extends pymbolic's parser to incorporate
    Fortran specific syntax and to map pymbolic expressions to Loki expressions, utilising 
    the mapper :any:`PymbolicMapper`.

    **Further**, in order to ensure correct ordering of Fortran Statements as documented
    in `'WD 1539-1 J3/23-007r1 (Draft Fortran 2023)' <https://j3-fortran.org/doc/year/10/10-007.pdf#page=155>`_,
    pymbolic's parsing logic needed to be slightly adapted.

    Pymbolic references:

    * `GitHub: pymbolic <https://github.com/inducer/pymbolic>`_
    * `pymbolic/parser.py <https://github.com/inducer/pymbolic/blob/main/pymbolic/parser.py>`_
    * `pymbolic's parser documentation <https://documen.tician.de/pymbolic/utilities.html>`_

    .. note::
       **Example:**
        Using the expression parser and possibly evaluate them

        .. code-block::

            >>> from loki import parse_expr
            >>> # parse numerical expressions
            >>> ir = parse_expr('3 + 2**4')
            >>> ir
            Sum((IntLiteral(3, None), Power(IntLiteral(2, None), IntLiteral(4, None))))
            >>> # or expressions with variables
            >>> ir = parse_expr('a*b')
            >>> ir
            Product((DeferredTypeSymbol('a', None, None, <SymbolAttributes BasicType.DEFERRED>),\
                    DeferredTypeSymbol('b', None, None, <SymbolAttributes BasicType.DEFERRED>)))
            >>> # and provide a scope e.g, with some routine defining a and b as 'real's
            >>> ir = parse_expr('a*b', scope=routine)
            >>> ir
            Product((Scalar('a', None, None, None), Scalar('b', None, None, None)))
            >>> # further, it is possible to evaluate expressions
            >>> ir = parse_expr('a*b + 1', evaluate=True, context={'a': 2, 'b': 3})
            >>> ir
            >>> IntLiteral(7, None)
            >>> # even with functions implemented in Python
            >>> def add(a, b):
            >>>     return a + b
            >>> ir = parse_expr('a + add(a, b)', evaluate=True, context={'a': 2, 'b': 3, 'add': add})
            >>> ir
            >>> IntLiteral(7, None)

    .. automethod:: __call__
    """

    _f_true = intern("f_true")
    _f_false = intern("f_false")
    _f_lessequal = intern('_f_lessequal')
    _f_less = intern('_f_less')
    _f_greaterequal = intern('_f_greaterequal')
    _f_greater = intern('_f_greater')
    _f_equal = intern('_f_equal')
    _f_notequal = intern('_f_notequal')
    _f_and = intern("and")
    _f_or = intern("or")
    _f_not = intern("not")
    _f_float = intern("f_float")
    _f_int = intern("f_int")
    _f_string = intern("f_string")
    _f_openbracket = intern("openbracket")
    _f_closebracket = intern("closebracket")
    _f_derived_type = intern("dot")

    lex_table = [
            (_f_true, pytools.lex.RE(r"\.true\.", re.IGNORECASE)),
            (_f_false, pytools.lex.RE(r"\.false\.", re.IGNORECASE)),
            (_f_lessequal, pytools.lex.RE(r"\.le\.", re.IGNORECASE)),
            (_f_less, pytools.lex.RE(r"\.lt\.", re.IGNORECASE)),
            (_f_greaterequal, pytools.lex.RE(r"\.ge\.", re.IGNORECASE)),
            (_f_greater, pytools.lex.RE(r"\.gt\.", re.IGNORECASE)),
            (_f_equal, pytools.lex.RE(r"\.eq\.", re.IGNORECASE)),
            (_f_notequal, ("|", pytools.lex.RE(r"\.ne\.", re.IGNORECASE), pytools.lex.RE(r"/=", re.IGNORECASE))),
            (_f_and, pytools.lex.RE(r"\.and\.", re.IGNORECASE)),
            (_f_or, pytools.lex.RE(r"\.or\.", re.IGNORECASE)),
            (_f_not, pytools.lex.RE(r"\.not\.", re.IGNORECASE)),
            (_f_float, ("|", pytools.lex.RE(r"[0-9]+\.[0-9]*([eEdD][+-]?[0-9]+)?(_([\w$]+|[0-9]+))+$", re.IGNORECASE))),
            (_f_int, pytools.lex.RE(r"[0-9]+?(_[a-zA-Z]*)", re.IGNORECASE)),
            (_f_string, ("|", pytools.lex.RE(r'\".*\"', re.IGNORECASE),
                pytools.lex.RE(r"\'.*\'", re.IGNORECASE))),
            (_f_openbracket, pytools.lex.RE(r"\(/")),
            (_f_closebracket, pytools.lex.RE(r"/\)")),
            (_f_derived_type, pytools.lex.RE(r"\%")),
            ] + ParserBase.lex_table
    """
    Extend :any:`pymbolic.parser.Parser.lex_table` to accomodate for Fortran specifix syntax/expressions.
    """

    ParserBase._COMP_TABLE.update({
         _f_lessequal: "<=",
         _f_less: "<",
         _f_greaterequal: ">=",
         _f_greater: ">",
         _f_equal: "==",
         _f_notequal: "!="
         })

    @staticmethod
    def _parenthesise(expr):
        """
        Utility method to parenthesise specific expressions.

        E.g., from :any:`pymbolic.primitives.Sum` to 
        :any:`ParenthesisedAdd`.
        """
        if isinstance(expr, pmbl.Sum):
            return sym_ops.ParenthesisedAdd(expr.children)
        if isinstance(expr, pmbl.Product):
            return sym_ops.ParenthesisedMul(expr.children)
        if isinstance(expr, pmbl.Quotient):
            return sym_ops.ParenthesisedDiv(numerator=expr.numerator,
                    denominator=expr.denominator)
        if isinstance(expr, pmbl.Power):
            return sym_ops.ParenthesisedPow(base=expr.base, exponent=expr.exponent)
        return expr

    def parse_prefix(self, pstate):
        pstate.expect_not_end()

        if pstate.is_next(_minus):
            pstate.advance()
            left_exp = pmbl.Product((-1, self.parse_expression(pstate, _PREC_UNARY)))
            return left_exp
        if pstate.is_next(_openpar):
            pstate.advance()

            if pstate.is_next(_closepar):
                left_exp = ()
            else:
                # This is parsing expressions separated by commas, so it
                # will return a tuple. Kind of the lazy way out.
                left_exp = self.parse_expression(pstate)
                # NECESSARY to ensure correct ordering!
                left_exp = self._parenthesise(left_exp)
            pstate.expect(_closepar)
            pstate.advance()
            if isinstance(left_exp, tuple):
                # These could just be plain parentheses.

                # Finalization prevents things from being appended
                # to containers after their closing delimiter.
                left_exp = FinalizedTuple(left_exp)
            return left_exp
        return super().parse_prefix(pstate)

    def parse_postfix(self, pstate, min_precedence, left_exp):

        did_something = False
        if pstate.is_next(self._f_derived_type) and _PREC_CALL > min_precedence:
            pstate.advance()
            right_exp = self.parse_expression(pstate, _PREC_PLUS)
            left_exp = pmbl.Lookup(left_exp, right_exp)
            did_something = True
        elif pstate.is_next(_times) and _PREC_TIMES > min_precedence:
            pstate.advance()
            right_exp = self.parse_expression(pstate, _PREC_PLUS)
            # NECESSARY to ensure correct ordering!
            # pylint: disable=unidiomatic-typecheck
            if type(right_exp) is pmbl.Quotient:
                left_exp = pmbl.Quotient(numerator=pmbl.Product((left_exp, right_exp.numerator)),
                        denominator=right_exp.denominator)
            # pylint: disable=unidiomatic-typecheck
            elif type(right_exp) is pmbl.Product:
                left_exp = pmbl.Product((sym.Product((left_exp, right_exp.children[0])), right_exp.children[1]))
            else:
                left_exp = pmbl.Product((left_exp, right_exp))
            did_something = True
        elif pstate.is_next(_plus) and _PREC_PLUS > min_precedence:
            pstate.advance()
            right_exp = self.parse_expression(pstate, _PREC_PLUS)
            left_exp = pmbl.Sum((left_exp, right_exp))
            did_something = True
        elif pstate.is_next(_minus) and _PREC_PLUS > min_precedence:
            pstate.advance()
            right_exp = self.parse_expression(pstate, _PREC_PLUS)
            right_exp = pmbl.Product((-1, right_exp))
            left_exp = pmbl.Sum((left_exp, right_exp))
            did_something = True
        else:
            return super().parse_postfix(pstate, min_precedence, left_exp)
        return left_exp, did_something

    def parse_terminal(self, pstate):
        if pstate.is_next(self._f_float):
            return self.parse_f_float(pstate.next_str_and_advance())
        if pstate.is_next(self._f_int):
            return self.parse_f_int(pstate.next_str_and_advance())
        if pstate.is_next(self._f_string):
            return self.parse_f_string(pstate.next_str_and_advance())
        if pstate.is_next(self._f_true):
            assert pstate.next_str_and_advance().lower() == ".true."
            return sym.LogicLiteral('.TRUE.')
        if pstate.is_next(self._f_false):
            assert pstate.next_str_and_advance().lower() == ".false."
            return sym.LogicLiteral('.FALSE.')
        return super().parse_terminal(pstate)

    def __call__(self, expr_str, scope=None, evaluate=False, strict=False, context=None):
        """
        Call Loki String Parser to convert expression(s) represented in a string to Loki expression(s)/IR.

        Parameters
        ----------
        expr_str : str
            The expression as a string
        scope : :any:`Scope`
            The scope to which symbol names inside the expression belong
        evaluate : bool, optional
            Whether to evaluate the expression or not (default: `False`)
        strict : bool, optional
            Whether to raise exception for unknown variables/symbols when
            evaluating an expression (default: `False`)
        context : dict, optional
            Symbol context, defining variables/symbols/procedures to help/support
            evaluating an expression

        Returns
        -------
        :any:`Expression`
            The expression tree corresponding to the expression
        """
        from loki.ir import AttachScopes  # pylint: disable=import-outside-toplevel,cyclic-import
        from loki.expression.evaluation import eval_expr # pylint: disable=import-outside-toplevel,cyclic-import
        result = super().__call__(expr_str)
        ir = PymbolicMapper()(result)
        if evaluate:
            ir = eval_expr(ir, context=context, strict=strict)
        return AttachScopes().visit(ir, scope=scope or Scope())

    def parse_float(self, s):
        """
        Parse float literals.

        Do not cast to float via 'float()' in order to keep the original
        notation, e.g., do not convert 1E-3 to 0.003.
        """
        return sym.FloatLiteral(value=s.replace("d", "e").replace("D", "e"))

    def parse_f_float(self, s):
        """
        Parse "Fortran-style" float literals.

        E.g., ``3.1415_my_real_kind``.
        """
        stripped = s.split('_', 1)
        if len(stripped) == 2:
            return sym.FloatLiteral(value=self.parse_float(stripped[0]), kind=sym.Variable(name=stripped[1].lower()))
        return self.parse_float(stripped[0])

    def parse_f_int(self, s):
        """
        Parse "Fortran-style" int literals.

        E.g., ``1_my_int_kind``.
        """
        stripped = s.split('_', 1)
        value = int(stripped[0].replace("d", "e").replace("D", "e"))
        return sym.IntLiteral(value=value, kind=sym.Variable(name=stripped[1].lower()))

    def parse_f_string(self, s):
        """
        Parse string literals.
        """
        return sym.StringLiteral(s)


parse_expr = ExpressionParser()
"""
An instance of :any:`ExpressionParser` that allows parsing expression strings into a Loki expression tree.
See :any:`ExpressionParser.__call__` for a description of the available arguments.
"""
