# (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 collections import defaultdict
import enum
from functools import reduce
from math import gcd, floor
import operator as _op
import pymbolic.primitives as pmbl

from loki.expression.mappers import LokiIdentityMapper
from loki.expression.evaluation import LokiEvaluationMapper
import loki.expression.symbols as sym
from loki.tools import as_tuple

__all__ = [
    'is_constant', 'symbolic_op', 'simplify', 'accumulate_polynomial_terms',
    'Simplification', 'SimplifyMapper', 'is_dimension_constant', 'ceil_division',
    'get_pyrange', 'iteration_number', 'iteration_index'
]


def is_minus_prefix(expr):
    """
    Return `True` if the given expression prefixes a nested expression with a minus sign,
    else return `False`.

    It essentially means that `expr == Product((-1, ...))`.
    """
    if isinstance(expr, sym.Product) and expr.children:
        return pmbl.is_zero(expr.children[0]+1)
    return False


def strip_minus_prefix(expr):
    """
    Return the expression without the minus prefix.

    Raises a `ValueError` if the expression is not prefixed by a minus.
    """
    if not is_minus_prefix(expr):
        raise ValueError('Given expression does not have a minus prefix.')
    children = expr.children[1:]
    if len(children) == 1:
        return children[0]
    return sym.Product(as_tuple(children))


def is_constant(expr):
    """
    Return `True` if the given expression reduces to a constant value, else return `False`.
    """
    if is_minus_prefix(expr):
        return is_constant(strip_minus_prefix(expr))
    return pmbl.is_constant(expr)


def is_dimension_constant(d):
    """Establish if a given dimension symbol is a compile-time constant"""
    if isinstance(d, sym.IntLiteral):
        return True

    if isinstance(d, sym.RangeIndex):
        if d.lower:
            return is_dimension_constant(d.lower) and is_dimension_constant(d.upper)
        return is_dimension_constant(d.upper)

    if isinstance(d, sym.Scalar) and isinstance(d.initial , sym.IntLiteral):
        return True

    return False


def symbolic_op(expr1, op, expr2):
    """
    Evaluate `expr1 <op> expr2` (or equivalently, `op(expr1, expr2)`) and
    return the result.

    `op` can be any binary operation such as the rich comparison operators
    from the `operator` library.

    While calling this function largely equivalent to applying the operator
    directly, it is to be understood as a convenience layer that applies,
    depending on the operator, a number of symbolically neutral manipulations.
    Currently, this only applies to comparison operators (such as `eq`, `ne`,
    `lt, `le`, `gt`, `ge`). Since expression nodes do not imply an order,
    such comparisons would fail even if a symbolic meaning can be derived.

    For that reason, these operations are reformulated as the difference
    between the two expressions and compared against `0`. For example:
    ```
    # n < n + 1
    Scalar('n') < Sum((Scalar('n'), IntLiteral(1)))
    ```
    raises a `TypeError` but
    ```
    # n < n + 1
    symbolic_op(Scalar('n'), operator.lt, Sum((Scalar('n'), IntLiteral(1))))
    ```
    returns `True`.

    This is done by transforming this expression into
    ```
    # n - (n + 1) < 0
    Sum((Scalar('n'), Product(-1, Sum((Scalar('n'), IntLiteral(1)))))) < 0
    ```
    and then calling `simplify` on the left hand side to obtain
    ```
    # -1 < 0
    Product(-1, IntLiteral(1)) < 0
    ```
    In combination with stripping the minus prefix this yields the result.
    """
    if op in (_op.eq, _op.ne, _op.lt, _op.le, _op.gt, _op.ge):
        expr1, expr2 = simplify(expr1 - expr2), 0
        if is_minus_prefix(expr1):
            # Strip minus prefix to possibly yield constant expression
            if op in (_op.eq, _op.ne):
                return symbolic_op(strip_minus_prefix(expr1), op, expr2)
            return not symbolic_op(strip_minus_prefix(expr1), op, expr2)
    return op(expr1, expr2)


def distribute_product(expr):
    """
    Flatten (nested) products into a sum of products.

    This converts for example `a * (b + c) * (d + e)` to
    `a * b * d + a * c * d + a * b * e + a * c * e`.
    """
    def _retval(numerator, denominator):
        if not denominator:
            return numerator
        if len(denominator) == 1:
            return sym.Quotient(numerator, denominator[0])
        return sym.Quotient(numerator, sym.Product(as_tuple(denominator)))

    if not isinstance(expr, sym.Product):
        return expr

    queue = list(expr.children)
    denominator = []
    done = [[]]

    while queue:
        item = queue.pop(0)

        if isinstance(item, sym.IntLiteral) and item.value == 1:
            continue

        if isinstance(item, sym.Product):
            # Prepend children to maintain order of operands
            queue = list(item.children) + queue
        elif isinstance(item, sym.Quotient):
            # Enqueue the numerator and save the denominator for later
            queue = [item.numerator] + queue
            denominator += [item.denominator]
        elif isinstance(item, sym.Sum):
            # This is the distribution part
            old_done, done = done, []
            for child in item.children:
                done += [l + [child] for l in old_done]
        else:
            # Some other factor that we simply carry over
            done = [l + [item] for l in done]

    if not done:
        return _retval(sym.IntLiteral(1), denominator)

    # Form the new products, eliminating multiple `-1` in the process
    children = []
    for components in done:
        is_neg = False
        if -1 in components:
            is_neg = sum(1 for v in components if v == -1) % 2 == 1
            components = [v for v in components if v != -1]

        if not components:
            components = sym.IntLiteral(1)
        elif len(components) == 1:
            components = components[0]
        else:
            components = sym.Product(as_tuple(components))

        if is_neg:
            components = sym.Product((-1, components))
        children.append(components)

    if len(children) == 1:
        return _retval(children[0], denominator)
    return _retval(sym.Sum(as_tuple(children)), denominator)


def distribute_quotient(expr):
    """
    Flatten (nested) quotients into a sum of quotients.

    This converts for example `(a/b + c) / d` to `a / (b*d) + c / d`.
    """
    if not isinstance(expr, sym.Quotient):
        return expr

    if is_minus_prefix(expr.numerator):
        q = sym.Quotient(strip_minus_prefix(expr.numerator), expr.denominator)
        return sym.Product((-1, distribute_quotient(q)))

    if is_minus_prefix(expr.denominator):
        q = sym.Quotient(expr.numerator, strip_minus_prefix(expr.denominator))
        return sym.Product((-1, distribute_quotient(q)))

    queue = [expr.numerator]
    done = []

    while queue:
        item = queue.pop(0)

        if isinstance(item, sym.IntLiteral) and item.value == 0:
            continue

        if isinstance(item, sym.Sum):
            # Prepend children to maintain order of operands
            queue = list(item.children) + queue
        elif isinstance(item, sym.Quotient):
            done += [distribute_quotient(sym.Quotient(item.numerator, item.denominator * expr.denominator))]
        else:
            # Convert to a quotient
            done += [sym.Quotient(item, expr.denominator)]

    if not done:
        return sym.IntLiteral(1)
    if len(done) == 1:
        return done[0]
    return sym.Sum(as_tuple(done))


def flatten_expr(expr):
    """
    Flatten an expression by flattening any sub-sums and distributing products and quotients.

    This converts for example `a + (b - (c + d) * e)` to `a + b - c * e - d * e`.

    This is an (enhanced) re-implementation of the original `flattened_sum` routine from
    Pymbolic to account for the Loki-specific expression nodes and expand the flattening to
    distributing products.
    """
    queue = list(as_tuple(expr))
    done = []

    while queue:
        item = queue.pop(0)

        if pmbl.is_zero(item):
            continue

        if isinstance(item, sym.Product):
            item = distribute_product(item)

        if isinstance(item, sym.Quotient):
            item = distribute_quotient(item)

        if isinstance(item, sym.Sum):
            # Prepend children to maintain order of operands
            queue = list(item.children) + queue
        else:
            done.append(item)

    if not done:
        return sym.IntLiteral(0)
    if len(done) == 1:
        return done[0]
    return sym.Sum(as_tuple(done))


def sum_literals(expr, int_arithmetic=True, fp_arithmetic=False):
    """
    Sum up the values of all numeric literals in the sum and return the reduced sum.

    Parameters
    ----------
    expr : :any:`Sum`
         The sum comprising constant and non-constant sub-expressions.
    int_arithmetic : bool, default True
        Perform arithmetic on integer literals
    fp_arithmetic : bool, default False
        Perform arithmetic on floating point literals
    """
    def _process(child):
        if (int_arithmetic or fp_arithmetic) and isinstance(child, sym.IntLiteral):
            return child.value, False, None
        if fp_arithmetic and isinstance(child, sym.FloatLiteral):
            return float(child.value), True, None
        if is_minus_prefix(child):
            value, has_float, stripped_child = _process(strip_minus_prefix(child))
            if value != 0:
                return -value, has_float, stripped_child
        return 0, False, child

    if not isinstance(expr, sym.Sum):
        return expr

    transformed_components = list(zip(*[_process(child) for child in expr.children]))
    value = sum(transformed_components[0])
    has_float = any(transformed_components[1])
    if not int_arithmetic and not has_float:
        return expr
    remaining_components = [ch for ch in transformed_components[2] if ch is not None]
    if value != 0:
        remaining_components = [sym.Literal(value)] + remaining_components

    if not remaining_components:
        return sym.Literal(0.0) if has_float else sym.IntLiteral(0)
    if len(remaining_components) == 1:
        return remaining_components[0]
    return sym.Sum(as_tuple(remaining_components))


def separate_coefficients(expr, int_arithmetic=True, fp_arithmetic=False):
    """
    Helper routine that separates components of a product into constant coefficients
    and remaining factors.

    Parameters
    ----------
    expr : :any:`Product`
         The product comprising constant and non-constant sub-expressions.
    int_arithmetic : bool, default True
        Perform arithmetic on integer literals
    fp_arithmetic : bool, default False
        Perform arithmetic on floating point literals

    Returns
    -------
    (int, list)
         The constant coefficient and remaining non-constant sub-expressions.
    """
    def _process(child):
        if (int_arithmetic or fp_arithmetic) and isinstance(child, int):
            return child, False, None
        if (int_arithmetic or fp_arithmetic) and isinstance(child, sym.IntLiteral):
            return child.value, False, None
        if fp_arithmetic and isinstance(child, sym.FloatLiteral):
            return float(child.value), True, None
        if is_minus_prefix(child):
            # We recurse here as products that are only there to change the sign
            # should not introduce a layer in the expression tree.
            value, has_float, component = _process(child.children[1])
            return -value, has_float, component
        return 1, False, child

    if (int_arithmetic or fp_arithmetic) and isinstance(expr, sym.IntLiteral):
        return expr.value, False, []
    if fp_arithmetic and isinstance(expr, sym.FloatLiteral):
        return float(expr.value), True, []
    if not isinstance(expr, sym.Product):
        return 1, False, [expr]

    if is_minus_prefix(expr):
        value, has_float, remaining_components = separate_coefficients(
            strip_minus_prefix(expr), int_arithmetic=int_arithmetic, fp_arithmetic=fp_arithmetic
        )
        return -value, has_float, remaining_components

    transformed_components = list(zip(*[_process(child) for child in expr.children]))
    value = reduce(_op.mul, transformed_components[0], 1)
    has_float = any(transformed_components[1])
    if not int_arithmetic and not has_float:
        return 1, False, [expr]
    remaining_components = [ch for ch in transformed_components[2] if ch is not None]
    return value, has_float, remaining_components


def mul_literals(expr, int_arithmetic=True, fp_arithmetic=False):
    """
    Multiply all `IntLiteral` in the given `Product` and return the reduced expression.

    Parameters
    ----------
    expr : :any:`Product`
         The product comprising constant and non-constant sub-expressions.
    int_arithmetic : bool, default True
        Perform arithmetic on integer literals
    fp_arithmetic : bool, default False
        Perform arithmetic on floating point literals
    """
    if not isinstance(expr, sym.Product):
        return expr

    value, has_float, remaining_components = separate_coefficients(
        expr, int_arithmetic=int_arithmetic, fp_arithmetic=fp_arithmetic
    )
    if value == 0:
        return sym.Literal(value)
    if abs(value) != 1:
        remaining_components = [sym.Literal(abs(value))] + remaining_components

    if not remaining_components:
        ret = sym.Literal(1.0) if has_float else sym.IntLiteral(1)
    elif len(remaining_components) == 1:
        ret = remaining_components[0]
    else:
        ret = sym.Product(as_tuple(remaining_components))

    if value < 0:
        return sym.Product((-1, ret))
    return ret


def div_literals(expr, fp_arithmetic=False):
    """
    Reduce fractions where the denominator is a `IntLiteral`.

    Parameters
    ----------
    expr : :any:`Expression`
         The expression comprising constant and non-constant sub-expressions.
    fp_arithmetic : bool, default False
        Perform arithmetic on floating point literals
    """
    if not isinstance(expr, sym.Quotient):
        return expr

    if is_minus_prefix(expr.numerator):
        q = sym.Quotient(strip_minus_prefix(expr.numerator), expr.denominator)
        return sym.Product((-1, div_literals(q, fp_arithmetic=fp_arithmetic)))

    if is_minus_prefix(expr.denominator):
        q = sym.Quotient(expr.numerator, strip_minus_prefix(expr.denominator))
        return sym.Product((-1, div_literals(q, fp_arithmetic=fp_arithmetic)))

    if isinstance(expr.numerator, sym.FloatLiteral) or isinstance(expr.denominator, sym.FloatLiteral):
        if not fp_arithmetic:
            return expr
        return sym.Literal(float(expr.numerator.value) / float(expr.denominator.value))

    if not isinstance(expr.denominator, sym.IntLiteral):
        return expr

    if isinstance(expr.numerator, sym.IntLiteral):
        div = gcd(expr.numerator.value, expr.denominator.value)
        numerator = sym.IntLiteral(expr.numerator.value / div)
        denominator = sym.IntLiteral(expr.denominator.value / div)

    elif isinstance(expr.numerator, sym.Product):
        value, _, remaining_components = separate_coefficients(expr.numerator, fp_arithmetic=fp_arithmetic)
        div = gcd(value, expr.denominator.value)
        numerator = mul_literals(
            sym.Product((sym.IntLiteral(value / div), *remaining_components)), fp_arithmetic=fp_arithmetic
        )
        denominator = sym.IntLiteral(expr.denominator.value / div)

    else:
        numerator, denominator = expr.numerator, expr.denominator

    if denominator == 1:
        return numerator
    return sym.Quotient(numerator, denominator)


def accumulate_polynomial_terms(expr):
    """
    Collect all occurences of each base and determine the constant coefficient
    in a list of expressions.

    Note that this works for any non-constant sub-expression as "base" for summands and thus
    this can be applied not only to polynomials.

    :param list components: list of expressions, e.g., components of a :py:class:`sym.Sum`.
    :returns: mapping of base and corresponding coefficient
    :rtype: dict
    """
    if isinstance(expr, sym.Sum):
        components = expr.children
    else:
        components = as_tuple(expr)

    summands = defaultdict(int)  # map (base, coefficient) pairs
    for item in components:
        if isinstance(item, sym.Product):
            value, _, remaining_components = separate_coefficients(item)
            if value == 0:
                continue
            if not remaining_components:
                summands[1] += value
            else:
                # We sort the components using their string representation
                summands[as_tuple(sorted(remaining_components, key=str))] += value
        elif isinstance(item, int):
            summands[1] += item
        elif isinstance(item, sym.IntLiteral):
            summands[1] += item.value
        else:
            summands[as_tuple(item)] += 1

    return dict(summands)


def collect_coefficients(expr):
    """
    Simplify a polynomial-type expression by combining all occurences of a non-constant
    subexpression into a single summand.

    :param list components: list of expressions, e.g., components of a :py:class:`sym.Sum`.
    :returns: reduced list of expressions.
    :rtype: list
    """
    def _get_coefficient(value):
        if value == 1:
            return []
        if value == -1:
            return [-1]
        if value < 0:
            return [-1, sym.IntLiteral(abs(value))]
        return [sym.IntLiteral(abs(value))]

    summands = accumulate_polynomial_terms(expr)
    components = []

    # Treat the constant part separately to make sure this is flat
    constant = summands.pop(1, 0)
    if constant < 0:
        components += [sym.Product((-1, sym.IntLiteral(abs(constant))))]
    elif constant > 0:
        components += [sym.IntLiteral(constant)]

    # Insert the remaining summands
    for base, factor in summands.items():
        if factor == 0:
            continue
        if factor == 1 and len(base) == 1:
            components.append(base[0])
        else:
            components.append(sym.Product(as_tuple(_get_coefficient(factor) + list(base))))

    if not components:
        return sym.IntLiteral(0)
    if len(components) == 1:
        return components[0]
    return sym.Sum(as_tuple(components))


class Simplification(enum.Flag):
    """
    The selection of available simplification techniques that can be used to simplify expressions.
    Multiple techniques can be combined using bitwise logical operations, for example:
    ```
    Flatten | IntegerArithmetic
    ALL & ~Flatten
    ```

    Attributes:
        Flatten             Flatten sub-sums and distribute products.
        IntegerArithmetic   Perform arithmetic on integer literals (addition and multiplication).
        FloatingPointArithmetic Perform arithmetic on floating point literals (addition and multiplication).
        CollectCoefficients Combine summands as far as possible.
        LogicEvaluation     Resolve logically fully determinate expressions, like ``1 == 1`` or ``1 == 6``
        ALL                 All of the above.
    """
    Flatten = enum.auto()
    IntegerArithmetic = enum.auto()
    FloatingPointArithmetic = enum.auto()
    CollectCoefficients = enum.auto()
    LogicEvaluation = enum.auto()

    # pylint: disable-next=unsupported-binary-operation
    ALL = Flatten | IntegerArithmetic | FloatingPointArithmetic | CollectCoefficients | LogicEvaluation


class SimplifyMapper(LokiIdentityMapper):
    """
    A mapper that attempts to symbolically simplify an expression.

    It applies all enabled simplifications from `Simplification` to a expression.
    """
    # pylint: disable=abstract-method

    def __init__(self, enabled_simplifications=Simplification.ALL):
        super().__init__()

        self.enabled_simplifications = enabled_simplifications

    def map_sum(self, expr, *args, **kwargs):
        new_expr = sym.Sum(as_tuple([self.rec(child, *args, **kwargs) for child in expr.children]))

        if self.enabled_simplifications & Simplification.Flatten:
            new_expr = flatten_expr(new_expr)

        if self.enabled_simplifications & (Simplification.IntegerArithmetic | Simplification.FloatingPointArithmetic):
            new_expr = sum_literals(
                new_expr,
                int_arithmetic=self.enabled_simplifications & Simplification.IntegerArithmetic,
                fp_arithmetic=self.enabled_simplifications & Simplification.FloatingPointArithmetic
            )

        if self.enabled_simplifications & Simplification.CollectCoefficients:
            new_expr = collect_coefficients(new_expr)

        if new_expr != expr:
            return self.rec(new_expr, *args, **kwargs)
        return expr

    def map_product(self, expr, *args, **kwargs):
        new_expr = sym.Product(as_tuple([self.rec(child, *args, **kwargs) for child in expr.children]))

        if self.enabled_simplifications & Simplification.Flatten:
            new_expr = flatten_expr(new_expr)

        if self.enabled_simplifications & (Simplification.IntegerArithmetic | Simplification.FloatingPointArithmetic):
            new_expr = mul_literals(
                new_expr,
                int_arithmetic=self.enabled_simplifications & Simplification.IntegerArithmetic,
                fp_arithmetic=self.enabled_simplifications & Simplification.FloatingPointArithmetic
            )

        if new_expr != expr:
            return self.rec(new_expr, *args, **kwargs)
        return expr

    def map_quotient(self, expr, *args, **kwargs):
        numerator = self.rec(expr.numerator, *args, **kwargs)
        denominator = self.rec(expr.denominator, *args, **kwargs)
        new_expr = sym.Quotient(numerator, denominator)

        if self.enabled_simplifications & Simplification.Flatten:
            new_expr = flatten_expr(new_expr)

        if self.enabled_simplifications & (Simplification.IntegerArithmetic | Simplification.FloatingPointArithmetic):
            new_expr = div_literals(
                new_expr, fp_arithmetic=self.enabled_simplifications & Simplification.FloatingPointArithmetic
            )

        if new_expr != expr:
            return self.rec(new_expr, *args, **kwargs)
        return expr

    def map_power(self, expr, *args, **kwargs):
        base = self.rec(expr.base, *args, **kwargs)
        exponent = self.rec(expr.exponent, *args, **kwargs)

        if self.enabled_simplifications & Simplification.IntegerArithmetic:
            literal_types = (sym.IntLiteral, sym.FloatLiteral)
            base_value = float(base.value) if isinstance(base, sym.FloatLiteral) else getattr(base, 'value', None)
            exponent_value = (
                float(exponent.value) if isinstance(exponent, sym.FloatLiteral) else getattr(exponent, 'value', None)
            )

            if isinstance(base, literal_types) and base_value == 1:
                return base
            if isinstance(exponent, literal_types):
                if exponent_value == 0:
                    return sym.IntLiteral(1)
                if exponent_value == 1:
                    return base
                if isinstance(base, literal_types) and exponent_value > 0 and float(exponent_value).is_integer():
                    return sym.Literal(base_value ** int(exponent_value))

        return sym.Power(base, exponent)

    map_parenthesised_add = map_sum
    map_parenthesised_mul = map_product
    map_parenthesised_div = map_quotient

    def map_comparison(self, expr, *args, **kwargs):
        def get_constant_value(expr):
            if is_minus_prefix(expr):
                return -1 * strip_minus_prefix(expr).value
            return expr.value

        left = self.rec(expr.left, *args, **kwargs)
        right = self.rec(expr.right, *args, **kwargs)

        op_map = {'==': _op.eq, '>': _op.gt, '>=': _op.ge, '!=': _op.ne,
                '<': _op.lt, '<=': _op.le}
        if self.enabled_simplifications & Simplification.LogicEvaluation:
            if is_constant(left) and is_constant(right):
                left = get_constant_value(left)
                right = get_constant_value(right)
                if op_map[expr.operator](left, right):
                    return sym.LogicLiteral('True')
                return sym.LogicLiteral('False')

        return sym.Comparison(operator=expr.operator, left=left, right=right)

    def map_logical_and(self, expr, *args, **kwargs):
        children = tuple(self.rec(child, *args, **kwargs) for child in expr.children)
        if self.enabled_simplifications & Simplification.LogicEvaluation:
            if any(c == 'False' for c in children):
                return sym.LogicLiteral('False')
            if any(c == 'True' for c in children):
                # Trim all literals and return .true. if all were .true.
                children = tuple(c for c in children if not c == 'True')

        return sym.LogicalAnd(children) if len(children) > 0 else sym.LogicLiteral('True')

    def map_logical_or(self, expr, *args, **kwargs):
        children = tuple(self.rec(child, *args, **kwargs) for child in expr.children)
        if self.enabled_simplifications & Simplification.LogicEvaluation:
            if any(c == 'True' for c in children):
                return sym.LogicLiteral('True')
            if any(c == 'False' for c in children):
                # Trim all literals and return .false. if all were .false.
                children = tuple(c for c in children if not c == 'False')

        return sym.LogicalOr(children) if len(children) > 0 else sym.LogicLiteral('False')

    def map_logical_not(self, expr, *args, **kwargs):
        child = self.rec(expr.child, *args, **kwargs)
        if self.enabled_simplifications & Simplification.LogicEvaluation:
            if child == 'True':
                return sym.LogicLiteral(False)
            if child == 'False':
                return sym.LogicLiteral(True)

        return sym.LogicalNot(child)

    def map_string_concat(self, expr, *args, **kwargs):
        children = tuple(self.rec(child, *args, **kwargs) for child in expr.children)

        new_children = []
        string_parts = []
        for child in children:
            if isinstance(child, sym.StringLiteral):
                string_parts.append(child.value)
            else:
                if string_parts:
                    new_children.append(sym.StringLiteral(''.join(string_parts)))
                    string_parts = []
                new_children.append(child)

        if string_parts:
            new_children.append(sym.StringLiteral(''.join(string_parts)))

        return new_children[0] if len(new_children) == 1 else sym.StringConcat(tuple(new_children))


def simplify(expr, enabled_simplifications=Simplification.ALL):
    """
    Simplify the given expression by applying selected simplifications.
    """
    return SimplifyMapper(enabled_simplifications=enabled_simplifications)(expr)


def ceil_division(iexpr1: pmbl.Expression, iexpr2: pmbl.Expression) -> pmbl.Expression:
    """
    Returns ceiled division expression of two integer expressions iexpr1/iexpr2.
    """
    expr = sym.Sum(children=(sym.Quotient(numerator=sym.Sum(children=(iexpr1, sym.IntLiteral(-1))),
                                          denominator=iexpr2),
                             sym.IntLiteral(1)))
    return simplify(expr, enabled_simplifications=Simplification.IntegerArithmetic)


def get_pyrange(loop_range: sym.LoopRange):
    """
    Returns a python range corresponding to a LoopRange of IntLiterals.
    """
    LEM = LokiEvaluationMapper()
    if loop_range.step is None:
        return range(LEM(loop_range.start), floor(LEM(loop_range.stop))+1)
    return range(LEM(loop_range.start), floor(LEM(loop_range.stop))+1, LEM(loop_range.step))



def iteration_number(iter_idx, loop_range: sym.LoopRange) -> pmbl.Expression:
    """
    Returns the normalized iteration number of the iteration variable

    Given the loop iteration index for an iteration in a loop defined by the
    :any:´LoopRange´ this method returns the normalized iteration index given by
    iter_num = (iter_idx - start + step)/step = (iter_idx-start)/step + 1

    Parameters
    ----------
    iter_idx : :any:`Variable`, :any:`Expression`, or :any:`IntLiteral`
        corresponding to a valid iteration index for the parameter `loop_range`
    loop_range: :any:`LoopRange`
    """
    if loop_range.step is None:
        expr = sym.Sum((sym.Sum((iter_idx, -loop_range.start)), sym.IntLiteral(1)))

    else:
        expr = sym.Sum(
            (sym.Quotient(sym.Sum((iter_idx, -loop_range.start)), loop_range.step),
             sym.IntLiteral(1)))
    return simplify(expr, enabled_simplifications=Simplification.IntegerArithmetic)


def iteration_index(iter_num, loop_range: sym.LoopRange) -> pmbl.Expression:
    """
    Returns the iteration index of the loop based on the iteration number

    Given the normalized iteration number for an iteration in a loop defined by the
    :any:´LoopRange´ this method returns the iteration index given by
    iter_idx = (iter_num-1)*step+start

    Parameters
    ----------
    iter_num : :any:`Variable`, :any:`Expression`, or :any:`sym.IntLiteral`
        corresponding to a valid iteration number for the parameter `loop_range`
    loop_range: :any:`LoopRange`
    """
    if loop_range.step is None:
        expr = sym.Sum((iter_num, sym.IntLiteral(-1), loop_range.start))

    else:
        expr = sym.Sum((sym.Product((sym.Sum((iter_num, sym.IntLiteral(-1))), loop_range.step)),
                    loop_range.start))
    return simplify(expr, enabled_simplifications=Simplification.IntegerArithmetic)
