# (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.

"""
Collection of utilities to automatically remove code elements or
section and to perform Dead Code Elimination.
"""

import operator as op
from itertools import chain

from loki.analyse import dataflow_analysis_attached
from loki.batch import Transformation
from loki.expression import simplify, symbols as sym, symbolic_op
from loki.ir import nodes as ir, Transformer, FindNodes, FindVariables, FindInlineCalls, SubstituteExpressions
from loki.ir.pragma_utils import (
    is_loki_pragma, pragma_regions_attached, get_pragma_parameters
)
from loki.program_unit import ProgramUnit
from loki.tools import flatten, as_tuple, OrderedSet
__all__ = [
    'RemoveCodeTransformation',
    'do_remove_dead_code', 'RemoveDeadCodeTransformer',
    'do_remove_marked_regions', 'RemoveRegionTransformer',
    'do_remove_calls', 'RemoveCallsTransformer', 'do_remove_unused_vars'
]


class RemoveCodeTransformation(Transformation):
    """
    A :any:`Transformation` that provides named call and import
    removal, code removal of pragma-marked regions and Dead Code
    Elimination for batch processing via the :any:`Scheduler`.

    The transformation will apply the following methods in order:

    * :any:`do_remove_calls`
    * :any:`do_remove_marked_regions`
    * :any:`do_remove_dead_code`
    * :any:`do_remove_unused_vars`

    Parameters
    ----------
    remove_marked_regions : boolean
        Flag to trigger the use of :meth:`remove_marked_regions`;
        default: ``True``
    mark_with_comment : boolean
        Flag to trigger the insertion of a marker comment when
        removing a region; default: ``True``.
    replacement_call : optional, str
        Name of the "abort" subroutine to call if a replacement call
        is to be inserted in :meth:`do_remove_marked_regions`.
    replacement_msg : optional, str
        Optional error message that will be passed as argument to
        the replacement call in :meth:`do_remove_marked_regions`.
    replacement_module : optional, str
        Optional name of the module from which to import the
        replacement subroutine in :meth:`do_remove_marked_regions`.
    remove_dead_code : boolean
        Flag to trigger the use of :meth:`remove_dead_code`;
        default: ``False``
    use_simplify : boolean
        Use :any:`simplify` when branch pruning in during
        :meth:`remove_dead_code`.
    call_names : list of str
        List of subroutine names against which to match
        :any:`CallStatement` nodes during :meth:`remove_calls`.
    intrinsic_names : list of str
        List of module names against which to match :any:`GenericStmt`
        nodes during :meth:`remove_calls`.
    remove_imports : boolean
        Flag indicating whether to remove symbols from :any:`Import`
        objects during :meth:`remove_calls`; default: ``True``
    kernel_only : boolean
        Only apply the configured removal to subroutines marked as
        "kernel"; default: ``False``
    remove_unused_args : boolean
        Remove unused dummy arguments from routines.
    remove_unused_vars : boolean
        Remove unused variables/locals from routines.
    remove_only_arrays : boolean
        Whether to only remove unused arrays from routines
        or all variables/locals.
    """

    _key = 'RemoveCodeTransformation'

    # Recurse to subroutines in ``contains`` clause
    recurse_to_internal_procedures = True
    reverse_traversal = True

    def __init__(
            self, remove_marked_regions=True, mark_with_comment=True,
            replacement_call=None, replacement_msg=None, replacement_module=None,
            remove_dead_code=False, use_simplify=True,
            call_names=None, intrinsic_names=None,
            remove_imports=True, kernel_only=False,
            remove_unused_args=False, remove_unused_vars=False,
            remove_only_arrays=True
    ):
        self.remove_marked_regions = remove_marked_regions
        self.mark_with_comment = mark_with_comment
        self.replacement_call = replacement_call
        self.replacement_msg = replacement_msg
        self.replacement_module = replacement_module

        self.remove_dead_code = remove_dead_code
        self.use_simplify = use_simplify

        self.call_names = as_tuple(call_names)
        self.intrinsic_names = as_tuple(intrinsic_names)
        self.remove_imports = remove_imports

        self.kernel_only = kernel_only
        self.remove_unused_args = remove_unused_args

        self.remove_unused_vars = remove_unused_vars
        self.remove_only_arrays = remove_only_arrays

    def transform_subroutine(self, routine, **kwargs):

        if kwargs.get('role') == 'kernel' or not self.kernel_only:
            # Apply named node removal to strip specific calls
            if self.call_names or self.intrinsic_names:
                do_remove_calls(
                    routine, call_names=self.call_names,
                    intrinsic_names=self.intrinsic_names,
                    remove_imports=self.remove_imports
                )

            # Apply marked region removal
            if self.remove_marked_regions:
                do_remove_marked_regions(
                    routine, mark_with_comment=self.mark_with_comment,
                    replacement_call=self.replacement_call,
                    replacement_msg=self.replacement_msg,
                    replacement_module=self.replacement_module
                )

            # Apply Dead Code Elimination
            if self.remove_dead_code:
                do_remove_dead_code(routine, use_simplify=self.use_simplify)

            if self.remove_unused_vars:
                do_remove_unused_vars(routine, remove_only_arrays=self.remove_only_arrays)

        if self.remove_unused_args and (item := kwargs['item']):
            # collect unused args from successors
            successors = kwargs['sub_sgraph'].successors(item=item)
            unused_args_map = {}
            for successor in successors:
                succ_data = successor.trafo_data.get(self._key, {})
                succ_ir = successor.ir
                unused_args_map[succ_ir] = succ_data.get(succ_ir.name, {}).get('unused_args', {})

            variable_map = routine.symbol_map
            used_or_defined_symbols = OrderedSet()
            if item.config.get('remove_unused_args', True) and kwargs['role'] == 'kernel':
                # internal procedures are transformed after the parent routine,
                # so determine their unused args eagerly to update call sites here
                for child in routine.subroutines:
                    child_symbol_map = child.symbol_map
                    child_used_symbols = get_used_or_defined_symbols(child)

                    # Collect symbols used or defined in contained subroutines, whilst ignoring
                    # those that mask parent symbols
                    used_or_defined_symbols |= OrderedSet(
                        variable_map[child_symbol.name_parts[0]].clone(dimensions=None)
                        for child_symbol in child_used_symbols
                        if child_symbol.name_parts[0] in variable_map
                        and child_symbol.name_parts[0] not in child_symbol_map
                    )

                    child_unused_args, _ = find_unused_dummy_args_and_vars(
                        child, used_or_defined_symbols=child_used_symbols
                    )
                    unused_args_map[child] = child_unused_args
                    item.trafo_data.setdefault(self._key, {})[child.name] = {'unused_args': child_unused_args}

            do_remove_unused_call_args(routine, unused_args_map)

            if item.config.get('remove_unused_args', True) and kwargs['role'] == 'kernel':
                # find unused args, which for contained subroutines may already be known
                item_data = item.trafo_data.get(self._key, {})
                unused_args = item_data.get(routine.name, {}).get('unused_args', {})

                if not unused_args:
                    used_or_defined_symbols |= get_used_or_defined_symbols(routine)
                    unused_args, _ = find_unused_dummy_args_and_vars(
                        routine, used_or_defined_symbols=used_or_defined_symbols
                    )

                    # store unused args keyed by routine name to avoid
                    # internal procedures overwriting the parent's data
                    item.trafo_data.setdefault(self._key, {})[routine.name] = {'unused_args': unused_args}

                do_remove_unused_dummy_args(routine, unused_args)


def do_remove_unused_dummy_args(routine, unused_args):
    """
    Utility routine to remove unused dummy arguments from
    a given routine.

    Parameters
    ----------
    routine : :any:`Subroutine`
        A :any:`Subroutine` whose unused dummy arguments will be removed.
    unused_args : dict
       A dict mapping the unused dummy argument symbol to its position in the
       routine's argument list.This must be retrieved using the
       :any:`find_unused_dummy_args_and_vars` utility.
    """

    routine.variables = [a for a in routine.variables
                         if not a.name.lower() in unused_args]

def do_remove_unused_vars(routine, unused_vars=None, remove_only_arrays=True):
    """
    Utility routine to remove unused variables (or only local arrays) from a given routine.

    Parameters
    ----------
    routine : :any:`Subroutine`
        A :any:`Subroutine` whose unused dummy arguments will be removed.
    unused_args : dict, optional
        A list of unused vars. This can be retrieved using the
       :any:`find_unused_dummy_args_and_vars` utility.
    remove_only_arrays : bool, optional
        Whether to only remove arrays or all variables/temporaries
        that are unused within the routine.
    """
    if unused_vars is None:
        _, unused_vars = find_unused_dummy_args_and_vars(routine)
    if remove_only_arrays:
        unused_vars = [var for var in unused_vars if isinstance(var, sym.Array)]
    routine.variables = [var for var in routine.variables
                         if not var.name.lower() in unused_vars]

def do_remove_unused_call_args(routine, unused_args_map):
    """
    Utility routine to remove unused arguments from all the
    :any:`CallStatement`s and :any:`InlineCall`s in a given routine.

    Parameters
    ----------
    routine : :any:`Subroutine`
        A :any:`Subroutine` whose call statements and inline calls will be updated.
    unused_args_map : dict
       A dict mapping the :any:`Subroutine` corresponding to the :any:`CallStatement`
       or :any:`InlineCall`, accessed via the ``routine`` property, to its unused arguments.
       The unused arguments must be retrieved using the :any:`find_unused_dummy_args`
       utility.
    """

    inline_call_map = {}
    for call in chain(FindNodes(ir.CallStatement).visit(routine.body), FindInlineCalls().visit(routine.body)):
        if not call.routine or not unused_args_map.get(call.routine, None):
            continue

        unused_positions = {c for c in unused_args_map[call.routine].values() if c < len(call.arguments)}
        unused_kwargs = [(kw, arg) for kw, arg in call.kwarguments if kw.lower() in unused_args_map[call.routine]]

        new_args = [arg for i, arg in enumerate(call.arguments) if i not in unused_positions]
        new_kwargs = [(kw, arg) for kw, arg in call.kwarguments if not (kw, arg) in unused_kwargs]

        if isinstance(call, ir.CallStatement):
            call._update(arguments=as_tuple(new_args), kwarguments=as_tuple(new_kwargs))
        else:
            inline_call_map[call] = call.clone(parameters=as_tuple(new_args), kw_parameters=dict(new_kwargs))

    if inline_call_map:
        routine.body = SubstituteExpressions(inline_call_map).visit(routine.body)


def get_used_or_defined_symbols(routine):
    """
    Collect symbols used or defined in a routine, including array-shape references.
    """
    variable_map = routine.symbol_map
    routine_arg_names = routine.argnames
    with dataflow_analysis_attached(routine):
        used_or_defined_symbols = routine.body.uses_symbols | routine.body.defines_symbols

        # We search for symbols used to define array sizes of symbols referenced
        # in the body, as well as local arrays declared in the routine.
        used_or_defined_array_shapes = [s.shape for s in used_or_defined_symbols if isinstance(s, sym.Array)]
        local_var_shapes = [
            var.shape for var in routine.variables
            if isinstance(var, sym.Array) and var.name.lower() not in routine_arg_names
        ]
        used_or_defined_symbols |= FindVariables(unique=True).visit(
            used_or_defined_array_shapes + local_var_shapes
        )

        used_or_defined_symbols |= OrderedSet(
            variable_map.get(v.name_parts[0], v).clone(dimensions=None)
            for v in used_or_defined_symbols
        )

    return used_or_defined_symbols


def find_unused_dummy_args_and_vars(routine, used_or_defined_symbols=None):
    """
    Utility routine to find all the unused arguments in a :any:`Subroutine`.

    Parameters
    ----------
    routine : :any:`Subroutine`
        A :any:`Subroutine` to search for unused dummy arguments.
    used_or_defined_symbols : set, optional
        The set of symbols used or defined in the routine.

    Return
    ------
    unused_args : dict
       A dict mapping the unused dummy argument symbol to its position in the
       routine's argument list.
    """

    if used_or_defined_symbols is None:
        used_or_defined_symbols = get_used_or_defined_symbols(routine)

    unused_args = {a.clone(dimensions=None): c for c, a in enumerate(routine.arguments)
                   if not a.name.lower() in used_or_defined_symbols}
    routine_arg_names = routine.argnames
    local_vars = [var for var in routine.variables if var.name.lower() not in routine_arg_names]
    unused_vars = [var.clone(dimensions=None) for var in local_vars
            if not var.name.lower() in used_or_defined_symbols]

    return unused_args, unused_vars


def do_remove_dead_code(routine, use_simplify=True):
    """
    Perform Dead Code Elimination on the given :any:`Subroutine` object.

    Parameters
    ----------
    routine : :any:`Subroutine`
        The subroutine to which to apply dead code elimination.
    simplify : boolean
        Use :any:`simplify` when evaluating expressions for branch pruning.
    """

    transformer = RemoveDeadCodeTransformer(use_simplify=use_simplify)
    routine.body = transformer.visit(routine.body)


class RemoveDeadCodeTransformer(Transformer):
    """
    :any:`Transformer` class that removes provably unreachable code paths.

    The primary modification performed is to prune individual code branches
    under :any:`Conditional` nodes.

    Parameters
    ----------
    use_simplify : boolean
        Use :any:`simplify` when evaluating expressions for branch pruning.
    """

    def __init__(self, use_simplify=True, **kwargs):
        super().__init__(**kwargs)
        self.use_simplify = use_simplify

    def visit_Conditional(self, o, **kwargs):
        condition = self.visit(o.condition, **kwargs)
        body = as_tuple(flatten(as_tuple(self.visit(o.body, **kwargs))))
        else_body = as_tuple(flatten(as_tuple(self.visit(o.else_body, **kwargs))))

        if self.use_simplify:
            condition = simplify(condition)

        if condition == 'True':
            return body

        if condition == 'False':
            return else_body

        has_elseif = o.has_elseif and else_body and isinstance(else_body[0], ir.Conditional)
        return self._rebuild(o, tuple((condition,) + (body,) + (else_body,)), has_elseif=has_elseif)

    def visit_MultiConditional(self, o, **kwargs):
        # Get select expression and simplify if requested
        expr = self.visit(o.expr, **kwargs)
        expr = simplify(expr) if self.use_simplify else expr

        values = self.visit(o.values, **kwargs)
        bodies = self.visit(o.bodies, **kwargs)
        else_body = self.visit(o.else_body, **kwargs)

        for val, body in zip(values, bodies):
            # Equate select expression against case values
            for v in val:
                if symbolic_op(expr, op.eq, v):
                    return body

        if expr == 'False':
            # Simplify to default if always false
            return else_body

        return self._rebuild(o, tuple((expr,) + (values,) + (bodies,) + (else_body,)), name=o.name)


def do_remove_marked_regions(
        routine, mark_with_comment=True, replacement_call=None,
        replacement_msg=None, replacement_module=None
):
    """
    Utility routine to remove code regions marked with
    ``!$loki remove`` pragmas from a subroutine's body.

    Optionally, any removed region might be marked with a
    comment and/or a simple single-argument "abort" call. For this,
    a subroutine name and message can be specified and an optional
    check for an import can also be defined to ensure the interface
    for the abort procedure is available. To bypass the replacement
    call insertion for individual pragma regions use
    ``!$loki remove no-replacement-call``.

    Parameters
    ----------
    routine : :any:`Subroutine`
        The subroutine to which to apply dead code elimination.
    mark_with_comment : boolean
        Flag to trigger the insertion of a marker comment when
        removing a region; default: ``True``.
    replacement_call : optional, str
        Name of the "abort" subroutine to call if a replacement call
        is to be inserted.
    replacement_msg : optional, str
        Optional error message that will be passed as a single
        argument to the replacement call.
    replacement_module : optional, str
        Optional name of the module from which to import the
        replacement subroutine. This will only be inserted if a
        replacement was perfored and will not replace existing imports
        of the same module or symbols.
    """

    transformer = RemoveRegionTransformer(
        mark_with_comment=mark_with_comment,
        replacement_call=replacement_call,
        replacement_msg=replacement_msg,
    )

    with pragma_regions_attached(routine, keyword='loki'):
        routine.body = transformer.visit(routine.body, scope=routine)

    if transformer.replacement_done and replacement_module:
        # Get newly inject procedure symbol for the replacement call
        callsym = sym.ProcedureSymbol(replacement_call, scope=routine)

        # Inject import of replacement module if it does not exist
        import_map = {i.module: i for i in routine.imports}
        if imprt := import_map.get(replacement_module):
            if not any(s == replacement_call for s in imprt.symbols):
                imprt._update(symbols=imprt.symbols + (callsym,))
        else:
            routine.spec.prepend(ir.Import(
                module=f'{replacement_module}', symbols=(callsym,), c_import=False)
            )


class RemoveRegionTransformer(Transformer):
    """
    A :any:`Transformer` that removes code regions marked with
    ``!$loki remove`` pragmas.

    This :any:`Transformer` only removes :any:`PragmaRegion` nodes,
    and thus requires the IR tree to have pragma regions attached, for
    example via :meth:`pragma_regions_attached`.

    When removing a marked code region the transformer may leave a
    comment or a replacement call to trigger "abort" errors in the
    source to mark the previous location.

    Parameters
    ----------
    mark_with_comment : boolean
        Flag to trigger the insertion of a marker comment when
        removing a region; default: ``True``.
    replacement_call : optional, str
        Name of the "abort" subroutine to call if a replacement call
        is to be inserted.
    replacement_msg : optional, str
        Optional error message that will be passed as a single
        argument to the replacmeent call.
    """

    def __init__(
            self, mark_with_comment=True, replacement_call=None, replacement_msg=None, **kwargs
    ):
        super().__init__(**kwargs)

        self.mark_with_comment = mark_with_comment

        # Replace section with call to trigger abort messages!
        self.replacement_call = replacement_call
        self.replacement_msg = replacement_msg
        self.replacement_done = False

    def visit_PragmaRegion(self, o, **kwargs):
        """ Remove :any:`PragmaRegion` nodes with ``!$loki remove`` pragmas """

        # Skip if the bypass clause is present
        bypass = 'no-replacement-call' in get_pragma_parameters(o.pragma, starts_with='remove')
        if is_loki_pragma(o.pragma, starts_with='remove'):

            # Leave a comment to mark the removed region in source
            replacement = []
            if self.mark_with_comment:
                replacement.append(ir.Comment(text='! [Loki] Removed content of pragma-marked region!'))

            if self.replacement_call and not bypass:
                # Get the outer scope, to avoid picking associates
                routine = kwargs['scope']
                while not isinstance(routine, ProgramUnit):
                    routine = routine.parent

                # If requested add a call to a simple subroutine with an error message arg
                replacement.append(ir.CallStatement(
                    name=sym.ProcedureSymbol(self.replacement_call, scope=routine),
                    arguments=sym.Literal(str(self.replacement_msg.format(routine.name)))
                ))
                # Set a flag to trigger import injections
                self.replacement_done = True

            return as_tuple(replacement)

        # Recurse into the pragama region and rebuild
        rebuilt = tuple(self.visit(i, **kwargs) for i in o.children)
        return self._rebuild(o, rebuilt)


def do_remove_calls(
        routine, call_names=None, intrinsic_names=None, remove_imports=True
):
    """
    Utility routine to remove all :any:`CallStatement` nodes
    to specific named subroutines in a :any:`Subroutine`.

    For more information, see :any:`RemoveCallsTransformer`.

    Parameters
    ----------
    call_names : list of str
        List of subroutine names against which to match
        :any:`CallStatement` nodes.
    intrinsic_names : list of str
        List of module names against which to match :any:`GenericStmt`
        nodes.
    remove_imports : boolean
        Flag indicating whether to remove the respective procedure
        symbols from :any:`Import` objects; default: ``True``.
    """

    transformer = RemoveCallsTransformer(
        call_names=call_names, intrinsic_names=intrinsic_names,
        remove_imports=remove_imports
    )
    routine.spec = transformer.visit(routine.spec)
    routine.body = transformer.visit(routine.body)


class RemoveCallsTransformer(Transformer):
    """
    A :any:`Transformer` that removes all :any:`CallStatement` nodes
    to specific named subroutines.

    This :any:`Transformer` will by default also remove the enclosing
    inline-conditional when encountering calls of the form ```if
    (flag) call named_procedure()``.

    This :any:`Transformer` will also attempt to match and remove
    :any:`GenericStmt` nodes against a given list of name strings.  This
    allows removing intrinsic calls like ``write (*,*) "..."``.

    In addition, this :any:`Transformer` can also attempt to match and
    remove :any:`Import` nodes if given a list of strings to
    match. This can be used to remove the associated imports of the
    removed subroutines.

    Parameters
    ----------
    call_names : list of str
        List of subroutine names against which to match
        :any:`CallStatement` nodes.
    intrinsic_names : list of str
        List of module names against which to match :any:`GenericStmt`
        nodes.
    remove_imports : boolean
        Flag indicating whether to remove the respective procedure
        symbols from :any:`Import` objects; default: ``True``.
    """

    def __init__(
            self, call_names=None, intrinsic_names=None,
            remove_imports=True, **kwargs
    ):
        super().__init__(**kwargs)

        self.call_names = as_tuple(call_names)
        self.intrinsic_names = as_tuple(intrinsic_names)
        self.remove_imports = remove_imports

    def visit_CallStatement(self, o, **kwargs):
        """ Match and remove :any:`CallStatement` nodes against name patterns """
        if o.name in self.call_names:
            return None

        rebuilt = tuple(self.visit(i, **kwargs) for i in o.children)
        return self._rebuild(o, rebuilt)

    def visit_Conditional(self, o, **kwargs):
        """ Remove inline-conditionals after recursing into their body """

        # First, recurse into condition and bodies
        cond, body, else_body = tuple(self.visit(i, **kwargs) for i in o.children)

        # Capture and remove newly empty inline conditionals
        if o.inline and len(body) == 0:
            return None

        return self._rebuild(o, (cond, body, else_body))

    def visit_GenericStmt(self, o, **kwargs):
        """ Match and remove :any:`GenericStmt` nodes against name patterns """
        if self.intrinsic_names:
            if any(str(c).lower() in str(o.text).lower() for c in self.intrinsic_names):
                return None

        rebuilt = tuple(self.visit(i, **kwargs) for i in o.children)
        return self._rebuild(o, rebuilt)

    def visit_Import(self, o, **kwargs):
        """ Remove the symbol of any named calls from Import nodes """

        symbols_found = any(s in self.call_names for s in o.symbols)
        if self.remove_imports and symbols_found:
            new_symbols = tuple(s for s in o.symbols if s not in self.call_names)
            return o.clone(symbols=new_symbols) if new_symbols else None

        rebuilt = tuple(self.visit(i, **kwargs) for i in o.children)
        return self._rebuild(o, rebuilt)
