# (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 loki.expression import LokiIdentityMapper
from loki.ir import (
    FindNodes, Assignment, StatementFunction, SubstituteExpressions
)
from loki.logging import detail
from loki.types import BasicType


__all__ = ['InlineSubstitutionMapper']


class InlineSubstitutionMapper(LokiIdentityMapper):
    """
    An expression mapper that defines symbolic substitution for inlining.
    """

    def map_algebraic_leaf(self, expr, *args, **kwargs):
        raise NotImplementedError

    def map_scalar(self, expr, *args, **kwargs):
        parent = self.rec(expr.parent, *args, **kwargs) if expr.parent is not None else None

        scope = kwargs.get('scope') or expr.scope
        # We're re-scoping an imported symbol
        if expr.scope != scope:
            return expr.clone(scope=scope, type=expr.type.clone(), parent=parent)
        return expr.clone(parent=parent)

    map_deferred_type_symbol = map_scalar

    def map_array(self, expr, *args, **kwargs):
        if expr.dimensions:
            dimensions = self.rec(expr.dimensions, *args, **kwargs)
        else:
            dimensions = None
        parent = self.rec(expr.parent, *args, **kwargs) if expr.parent is not None else None

        scope = kwargs.get('scope') or expr.scope
        # We're re-scoping an imported symbol
        if expr.scope != scope:
            return expr.clone(scope=scope, type=expr.type.clone(), parent=parent, dimensions=dimensions)
        return expr.clone(parent=parent, dimensions=dimensions)

    def map_procedure_symbol(self, expr, *args, **kwargs):
        parent = self.rec(expr.parent, *args, **kwargs) if expr.parent is not None else None

        scope = kwargs.get('scope') or expr.scope
        # We're re-scoping an imported symbol
        if expr.scope != scope:
            return expr.clone(scope=scope, type=expr.type.clone(), parent=parent)
        return expr.clone(parent=parent)

    def map_inline_call(self, expr, *args, **kwargs):
        if expr.procedure_type in (None, BasicType.DEFERRED) or expr.procedure_type.is_intrinsic:
            # Unkonw inline call, potentially an intrinsic
            # We still need to recurse and ensure re-scoping
            return super().map_inline_call(expr, *args, **kwargs)

        # if it is an inline call to a Statement Function
        if isinstance(expr.routine, StatementFunction):
            function = expr.routine
            # Substitute all arguments through the elemental body
            arg_map = dict(expr.arg_iter())
            fbody = SubstituteExpressions(arg_map).visit(function.rhs)
            return fbody

        function = expr.procedure_type.procedure
        v_result = [v for v in function.variables if v == function.name][0]

        scope = kwargs.get('scope') or expr.function.scope
        if scope and function.name in scope.interface_map:
            # Inline call to a function that is provided via an interface
            # We don't have the function body available for inlining
            detail(f'Cannot inline {expr.function.name} into {scope.name}. Only interface available.')
            return super().map_inline_call(expr, *args, **kwargs)

        # Substitute all arguments through the elemental body
        arg_map = dict(expr.arg_iter())
        fbody = SubstituteExpressions(arg_map).visit(function.body)

        # Extract the RHS of the final result variable assignment
        stmts = [s for s in FindNodes(Assignment).visit(fbody) if s.lhs == v_result]
        assert len(stmts) == 1
        rhs = self.rec(stmts[0].rhs, *args, **kwargs)
        return rhs
