"""
Simplify and optimise sequences of statements by rewriting and pulling out loop invariants.
"""

import ast
import copy
import itertools
from collections import OrderedDict
from functools import reduce

from brian2.core.functions import DEFAULT_CONSTANTS, DEFAULT_FUNCTIONS
from brian2.core.preferences import prefs
from brian2.core.variables import AuxiliaryVariable
from brian2.parsing.bast import (
    BrianASTRenderer,
    brian_ast,
    brian_dtype_from_dtype,
    dtype_hierarchy,
)
from brian2.parsing.rendering import NodeRenderer
from brian2.utils.stringtools import get_identifiers, word_substitute

from .statements import Statement

# Default namespace has all the standard functions and constants in it
defaults_ns = {k: v.pyfunc for k, v in DEFAULT_FUNCTIONS.items()}
defaults_ns.update({k: v.value for k, v in DEFAULT_CONSTANTS.items()})


__all__ = ["optimise_statements", "ArithmeticSimplifier", "Simplifier"]


def evaluate_expr(expr, ns):
    """
    Try to evaluate the expression in the given namespace

    Returns either (value, True) if successful, or (expr, False) otherwise.

    Examples
    --------
    >>> assumptions = {'exp': DEFAULT_FUNCTIONS['exp'].pyfunc,
    ...                'inf': DEFAULT_CONSTANTS['inf'].value}
    >>> evaluate_expr('1/2', assumptions)
    (0.5, True)
    >>> evaluate_expr('exp(-inf)', assumptions)
    (0.0, True)
    >>> evaluate_expr('sin(2*pi*freq*t)', assumptions)
    ('sin(2*pi*freq*t)', False)
    >>> evaluate_expr('1/0', assumptions)
    ('1/0', False)
    """
    try:
        val = eval(expr, ns)
        return val, True
    except (NameError, ArithmeticError):
        return expr, False


def expression_complexity(expr, variables):
    return brian_ast(expr, variables).complexity


def optimise_statements(scalar_statements, vector_statements, variables, blockname=""):
    """
    Optimise a sequence of scalar and vector statements

    Performs the following optimisations:

    1. Constant evaluations (e.g. exp(0) to 1). See `evaluate_expr`.
    2. Arithmetic simplifications (e.g. 0*x to 0). See `ArithmeticSimplifier`, `collect`.
    3. Pulling out loop invariants (e.g. v*exp(-dt/tau) to a=exp(-dt/tau) outside the loop and v*a inside).
       See `Simplifier`.
    4. Boolean simplifications (allowing the replacement of expressions with booleans with a sequence of if/thens).
       See `Simplifier`.

    Parameters
    ----------
    scalar_statements : sequence of Statement
        Statements that only involve scalar values and should be evaluated in the scalar block.
    vector_statements : sequence of Statement
        Statements that involve vector values and should be evaluated in the vector block.
    variables : dict of (str, Variable)
        Definition of the types of the variables.
    blockname : str, optional
        Name of the block (used for LIO constant prefixes to avoid name clashes)

    Returns
    -------
    new_scalar_statements : sequence of Statement
        As above but with loop invariants pulled out from vector statements
    new_vector_statements : sequence of Statement
        Simplified/optimised versions of statements
    """
    boolvars = {
        k: v
        for k, v in variables.items()
        if hasattr(v, "dtype") and brian_dtype_from_dtype(v.dtype) == "boolean"
    }
    # We use the Simplifier class by rendering each expression, which generates new scalar statements
    # stored in the Simplifier object, and these are then added to the scalar statements.
    simplifier = Simplifier(variables, scalar_statements, extra_lio_prefix=blockname)
    new_vector_statements = []
    for stmt in vector_statements:
        # Carry out constant evaluation, arithmetic simplification and loop invariants
        new_expr = simplifier.render_expr(stmt.expr)
        new_stmt = Statement(
            stmt.var,
            stmt.op,
            new_expr,
            stmt.comment,
            dtype=stmt.dtype,
            constant=stmt.constant,
            subexpression=stmt.subexpression,
            scalar=stmt.scalar,
        )
        # Now check if boolean simplification can be carried out
        complexity_std = expression_complexity(new_expr, simplifier.variables)
        idents = get_identifiers(new_expr)
        used_boolvars = [var for var in boolvars if var in idents]
        if len(used_boolvars):
            # We want to iterate over all the possible assignments of boolean variables to values in (True, False)
            bool_space = [[False, True] for _ in used_boolvars]
            expanded_expressions = {}
            complexities = {}
            for bool_vals in itertools.product(*bool_space):
                # substitute those values into the expr and simplify (including potentially pulling out new
                # loop invariants)
                subs = {var: str(val) for var, val in zip(used_boolvars, bool_vals)}
                curexpr = word_substitute(new_expr, subs)
                curexpr = simplifier.render_expr(curexpr)
                key = tuple((var, val) for var, val in zip(used_boolvars, bool_vals))
                expanded_expressions[key] = curexpr
                complexities[key] = expression_complexity(curexpr, simplifier.variables)
            # See Statement for details on these
            new_stmt.used_boolean_variables = used_boolvars
            new_stmt.boolean_simplified_expressions = expanded_expressions
            new_stmt.complexity_std = complexity_std
            new_stmt.complexities = complexities
        new_vector_statements.append(new_stmt)
    # Generate additional scalar statements for the loop invariants
    new_scalar_statements = copy.copy(scalar_statements)
    for expr, name in simplifier.loop_invariants.items():
        dtype_name = simplifier.loop_invariant_dtypes[name]
        if dtype_name == "boolean":
            dtype = bool
        elif dtype_name == "integer":
            dtype = int
        else:
            dtype = prefs.core.default_float_dtype
        new_stmt = Statement(
            name,
            ":=",
            expr,
            "",
            dtype=dtype,
            constant=True,
            subexpression=False,
            scalar=True,
        )
        new_scalar_statements.append(new_stmt)
    return new_scalar_statements, new_vector_statements


def _replace_with_zero(zero_node, node):
    """
    Helper function to return a "zero node" of the correct type.

    Parameters
    ----------
    zero_node : `ast.Constant`
        The node to replace
    node : `ast.Node`
        The node that determines the type

    Returns
    -------
    zero_node : `ast.Constant`
        The original ``zero_node`` with its value replaced by 0 or 0.0.
    """
    # must not change the dtype of the output,
    # e.g. handle 0/float->0.0 and 0.0/int->0.0
    zero_node.dtype = node.dtype
    if node.dtype == "integer":
        zero_node.value = 0
    else:
        zero_node.value = prefs.core.default_float_dtype(0.0)
    return zero_node


class ArithmeticSimplifier(BrianASTRenderer):
    """
    Carries out the following arithmetic simplifications:

    1. Constant evaluation (e.g. exp(0)=1) by attempting to evaluate the expression in an "assumptions namespace"
    2. Binary operators, e.g. 0*x=0, 1*x=x, etc. You have to take care that the dtypes match here, e.g.
       if x is an integer, then 1.0*x shouldn't be replaced with x but left as 1.0*x.

    Parameters
    ----------
    variables : dict of (str, Variable)
        Usual definition of variables.
    assumptions : sequence of str
        Additional assumptions that can be used in simplification, each assumption is a string statement.
        These might be the scalar statements for example.
    """

    def __init__(self, variables):
        BrianASTRenderer.__init__(self, variables, copy_variables=False)
        self.assumptions = []
        self.assumptions_ns = dict(defaults_ns)
        self.bast_renderer = BrianASTRenderer(variables, copy_variables=False)

    def render_node(self, node):
        """
        Assumes that the node has already been fully processed by BrianASTRenderer
        """
        if not hasattr(node, "simplified"):
            node = super().render_node(node)
            node.simplified = True
        # can't evaluate vector expressions, so abandon in this case
        if not node.scalar:
            return node
        # No evaluation necessary for simple names or numbers
        if node.__class__.__name__ in ["Name", "NameConstant", "Num", "Constant"]:
            return node
        # Don't evaluate stateful nodes (e.g. those containing a rand() call)
        if not node.stateless:
            return node
        # try fully evaluating using assumptions
        expr = NodeRenderer().render_node(node)
        val, evaluated = evaluate_expr(expr, self.assumptions_ns)
        if evaluated:
            if node.dtype == "boolean":
                val = bool(val)
                if hasattr(ast, "Constant"):
                    newnode = ast.Constant(val)
                elif hasattr(ast, "NameConstant"):
                    newnode = ast.NameConstant(val)
                else:
                    # None is the expression context, we don't use it so we just set to None
                    newnode = ast.Name(repr(val), None)
            elif node.dtype == "integer":
                val = int(val)
            else:
                val = prefs.core.default_float_dtype(val)
            if node.dtype != "boolean":
                newnode = ast.Constant(val)
            newnode.dtype = node.dtype
            newnode.scalar = True
            newnode.stateless = node.stateless
            newnode.complexity = 0
            return newnode
        return node

    def render_BinOp(self, node):
        if node.dtype == "float":  # only try to collect float type nodes
            if node.op.__class__.__name__ in [
                "Mult",
                "Div",
                "Add",
                "Sub",
            ] and not hasattr(node, "collected"):
                newnode = self.bast_renderer.render_node(collect(node))
                newnode.collected = True
                return self.render_node(newnode)
        left = node.left = self.render_node(node.left)
        right = node.right = self.render_node(node.right)
        node = super().render_BinOp(node)
        op = node.op
        # Handle multiplication by 0 or 1
        if op.__class__.__name__ == "Mult":
            for operand, other in [(left, right), (right, left)]:
                if operand.__class__.__name__ in ["Num", "Constant"]:
                    op_value = operand.value
                    if op_value == 0:
                        # Do not remove stateful functions
                        if node.stateless:
                            return _replace_with_zero(operand, node)
                    if op_value == 1:
                        # only simplify this if the type wouldn't be cast by the operation
                        if (
                            dtype_hierarchy[operand.dtype]
                            <= dtype_hierarchy[other.dtype]
                        ):
                            return other
        # Handle division by 1, or 0/x
        elif op.__class__.__name__ == "Div":
            if (
                left.__class__.__name__ in ["Num", "Constant"] and left.value == 0
            ):  # 0/x
                if node.stateless:
                    # Do not remove stateful functions
                    return _replace_with_zero(left, node)
            if (
                right.__class__.__name__ in ["Num", "Constant"] and right.value == 1
            ):  # x/1
                # only simplify this if the type wouldn't be cast by the operation
                if dtype_hierarchy[right.dtype] <= dtype_hierarchy[left.dtype]:
                    return left
        elif op.__class__.__name__ == "FloorDiv":
            if (
                left.__class__.__name__ in ["Num", "Constant"] and left.value == 0
            ):  # 0//x
                if node.stateless:
                    # Do not remove stateful functions
                    return _replace_with_zero(left, node)
            # Only optimise floor division by 1 if both numbers are integers,
            # for floating point values, floor division by 1 changes the value,
            # and division by 1.0 can change the type for an integer value
            if (
                left.dtype == right.dtype == "integer"
                and right.__class__.__name__ in ["Num", "Constant"]
                and right.value == 1
            ):  # x//1
                return left
        # Handle addition of 0
        elif op.__class__.__name__ == "Add":
            for operand, other in [(left, right), (right, left)]:
                if (
                    operand.__class__.__name__ in ["Num", "Constant"]
                    and operand.value == 0
                ):
                    # only simplify this if the type wouldn't be cast by the operation
                    if dtype_hierarchy[operand.dtype] <= dtype_hierarchy[other.dtype]:
                        return other
        # Handle subtraction of 0
        elif op.__class__.__name__ == "Sub":
            if right.__class__.__name__ in ["Num", "Constant"] and right.value == 0:
                # only simplify this if the type wouldn't be cast by the operation
                if dtype_hierarchy[right.dtype] <= dtype_hierarchy[left.dtype]:
                    return left

        # simplify e.g. 2*float to 2.0*float to make things more explicit: not strictly necessary
        # but might be useful for some codegen targets
        if node.dtype == "float" and op.__class__.__name__ in [
            "Mult",
            "Add",
            "Sub",
            "Div",
        ]:
            for subnode in [node.left, node.right]:
                if subnode.__class__.__name__ in ["Num", "Constant"] and not (
                    subnode.value is True or subnode.value is False
                ):
                    subnode.dtype = "float"
                    subnode.value = prefs.core.default_float_dtype(subnode.value)
        return node


class Simplifier(BrianASTRenderer):
    """
    Carry out arithmetic simplifications (see `ArithmeticSimplifier`) and loop invariants

    Parameters
    ----------
    variables : dict of (str, Variable)
        Usual definition of variables.
    scalar_statements : sequence of Statement
        Predefined scalar statements that can be used as part of simplification

    Notes
    -----

    After calling `render_expr` on a sequence of expressions (coming from vector statements typically),
    this object will have some new attributes:

    ``loop_invariants`` : OrderedDict of (expression, varname)
        varname will be of the form ``_lio_N`` where ``N`` is some integer, and the expressions will be
        strings that correspond to scalar-only expressions that can be evaluated outside of the vector
        block.
    ``loop_invariant_dtypes`` : dict of (varname, dtypename)
        dtypename will be one of ``'boolean'``, ``'integer'``, ``'float'``.
    """

    def __init__(self, variables, scalar_statements, extra_lio_prefix=""):
        BrianASTRenderer.__init__(self, variables, copy_variables=False)
        self.loop_invariants = OrderedDict()
        self.loop_invariant_dtypes = {}
        self.value = 0
        self.node_renderer = NodeRenderer()
        self.arithmetic_simplifier = ArithmeticSimplifier(variables)
        self.scalar_statements = scalar_statements
        if extra_lio_prefix is None:
            extra_lio_prefix = ""
        if len(extra_lio_prefix):
            extra_lio_prefix = f"{extra_lio_prefix}_"
        self.extra_lio_prefix = extra_lio_prefix

    def render_expr(self, expr):
        node = brian_ast(expr, self.variables)
        node = self.arithmetic_simplifier.render_node(node)
        node = self.render_node(node)
        return self.node_renderer.render_node(node)

    def render_node(self, node):
        """
        Assumes that the node has already been fully processed by BrianASTRenderer
        """
        # can we pull this out?
        if node.scalar and node.complexity > 0:
            expr = self.node_renderer.render_node(
                self.arithmetic_simplifier.render_node(node)
            )
            if expr in self.loop_invariants:
                name = self.loop_invariants[expr]
            else:
                self.value += 1
                name = f"_lio_{self.extra_lio_prefix}{str(self.value)}"
                self.loop_invariants[expr] = name
                self.loop_invariant_dtypes[name] = node.dtype
                numpy_dtype = {
                    "boolean": bool,
                    "integer": int,
                    "float": prefs.core.default_float_dtype,
                }[node.dtype]
                self.variables[name] = AuxiliaryVariable(
                    name, dtype=numpy_dtype, scalar=True
                )
            # None is the expression context, we don't use it so we just set to None
            newnode = ast.Name(name, None)
            newnode.scalar = True
            newnode.dtype = node.dtype
            newnode.complexity = 0
            newnode.stateless = node.stateless
            return newnode
        # otherwise, render node as usual
        return super().render_node(node)


def reduced_node(terms, op):
    """
    Reduce a sequence of terms with the given operator

    For examples, if terms were [a, b, c] and op was multiplication then the reduction would be (a*b)*c.

    Parameters
    ----------
    terms : list
        AST nodes.
    op : AST node
        Could be `ast.Mult` or `ast.Add`.

    Examples
    --------
    >>> import ast
    >>> nodes = [ast.Name(id='x'), ast.Name(id='y'), ast.Name(id='z')]
    >>> ast.unparse(reduced_node(nodes, ast.Mult))
    'x * y * z'
    >>> nodes = [ast.Name(id='x')]
    >>> ast.unparse(reduced_node(nodes, ast.Add))
    'x'
    """
    # Remove None terms
    terms = [term for term in terms if term is not None]
    if not len(terms):
        return None
    return reduce(lambda left, right: ast.BinOp(left, op(), right), terms)


def cancel_identical_terms(primary, inverted):
    """
    Cancel terms in a collection, e.g. a+b-a should be cancelled to b

    Simply renders the nodes into expressions and removes whenever there is a common expression
    in primary and inverted.

    Parameters
    ----------
    primary : list of AST nodes
        These are the nodes that are positive with respect to the operator, e.g.
        in x*y/z it would be [x, y].
    inverted : list of AST nodes
        These are the nodes that are inverted with respect to the operator, e.g.
        in x*y/z it would be [z].

    Returns
    -------
    primary : list of AST nodes
        Primary nodes after cancellation
    inverted : list of AST nodes
        Inverted nodes after cancellation
    """
    nr = NodeRenderer()
    expressions = {node: nr.render_node(node) for node in primary}
    expressions.update({node: nr.render_node(node) for node in inverted})
    new_primary = []
    inverted_expressions = [expressions[term] for term in inverted]
    for term in primary:
        expr = expressions[term]
        if expr in inverted_expressions and term.stateless:
            new_inverted = []
            for iterm in inverted:
                if expressions[iterm] == expr:
                    expr = ""  # handled
                else:
                    new_inverted.append(iterm)
            inverted = new_inverted
            inverted_expressions = [expressions[term] for term in inverted]
        else:
            new_primary.append(term)
    return new_primary, inverted


def collect(node):
    """
    Attempts to collect commutative operations into one and simplifies them.

    For example, if x and y are scalars, and z is a vector, then (x*z)*y should
    be rewritten as (x*y)*z to minimise the number of vector operations. Similarly,
    ((x*2)*3)*4 should be rewritten as x*24.

    Works for either multiplication/division or addition/subtraction nodes.

    The final output is a subexpression of the following maximal form:

        (((numerical_value*(product of scalars))/(product of scalars))*(product of vectors))/(product of vectors)

    Any possible cancellations will have been done.

    Parameters
    ----------
    node : Brian AST node
        The node to be collected/simplified.

    Returns
    -------
    node : Brian AST node
        Simplified node.
    """
    node.collected = True
    orignode_dtype = node.dtype
    # we only work on */ or +- ops, which are both BinOp
    if node.__class__.__name__ != "BinOp":
        return node
    # primary would be the * or + nodes, and inverted would be the / or - nodes
    terms_primary = []
    terms_inverted = []
    # we handle both multiplicative and additive nodes in the same way by using these variables
    if node.op.__class__.__name__ in ["Mult", "Div"]:
        op_primary = ast.Mult
        op_inverted = ast.Div
        op_null = prefs.core.default_float_dtype(1.0)  # the identity for the operator
        op_py_primary = lambda x, y: x * y
        op_py_inverted = lambda x, y: x / y
    elif node.op.__class__.__name__ in ["Add", "Sub"]:
        op_primary = ast.Add
        op_inverted = ast.Sub
        op_null = prefs.core.default_float_dtype(0.0)
        op_py_primary = lambda x, y: x + y
        op_py_inverted = lambda x, y: x - y
    else:
        return node
    if node.dtype == "integer":
        op_null_with_dtype = int(op_null)
    else:
        op_null_with_dtype = op_null
    # recursively collect terms into the terms_primary and terms_inverted lists
    collect_commutative(node, op_primary, op_inverted, terms_primary, terms_inverted)
    x = op_null
    # extract the numerical nodes and fully evaluate
    remaining_terms_primary = []
    remaining_terms_inverted = []
    for term in terms_primary:
        if term.__class__.__name__ == "Num":
            x = op_py_primary(x, term.n)
        elif term.__class__.__name__ == "Constant":
            x = op_py_primary(x, term.value)
        else:
            remaining_terms_primary.append(term)
    for term in terms_inverted:
        if term.__class__.__name__ == "Num":
            x = op_py_inverted(x, term.n)
        elif term.__class__.__name__ == "Constant":
            x = op_py_inverted(x, term.value)
        else:
            remaining_terms_inverted.append(term)
    # if the fully evaluated node is just the identity/null element then we
    # don't have to make it into an explicit term
    if x != op_null:
        num_node = ast.Constant(x)
    else:
        num_node = None
    terms_primary = remaining_terms_primary
    terms_inverted = remaining_terms_inverted
    node = num_node
    for scalar in (True, False):
        primary_terms = [term for term in terms_primary if term.scalar == scalar]
        inverted_terms = [term for term in terms_inverted if term.scalar == scalar]
        primary_terms, inverted_terms = cancel_identical_terms(
            primary_terms, inverted_terms
        )

        # produce nodes that are the reduction of the operator on these subsets
        prod_primary = reduced_node(primary_terms, op_primary)
        prod_inverted = reduced_node(inverted_terms, op_primary)

        # construct the simplest version of the fully simplified node (only doing operations where necessary)
        node = reduced_node([node, prod_primary], op_primary)
        if prod_inverted is not None:
            if node is None:
                node = ast.Constant(op_null_with_dtype)
            node = ast.BinOp(node, op_inverted(), prod_inverted)

    if node is None:  # everything cancelled
        node = ast.Constant(op_null_with_dtype)
    if (
        hasattr(node, "dtype")
        and dtype_hierarchy[node.dtype] < dtype_hierarchy[orignode_dtype]
    ):
        node = ast.BinOp(ast.Constant(op_null_with_dtype), op_primary(), node)
    node.collected = True
    return node


def collect_commutative(
    node, primary, inverted, terms_primary, terms_inverted, add_to_inverted=False
):
    # This function is called recursively, so we use add_to_inverted to keep track of whether or not
    # we're working in the numerator/denominator (for multiplicative nodes, equivalent for additive).
    op_primary = node.op.__class__ is primary
    # this should only be called with node a BinOp of type primary or inverted
    # left_exact is the condition that we can collect terms (we can do it with floats or add/sub,
    # but not integer mult/div - the reason being that for C-style division e.g. 3/(4/3)!=(3*3)/4
    left_exact = node.left.dtype == "float" or (
        hasattr(node.left, "op") and node.left.op.__class__.__name__ in ["Add", "Sub"]
    )
    if (
        node.left.__class__.__name__ == "BinOp"
        and node.left.op.__class__ in [primary, inverted]
        and left_exact
    ):
        collect_commutative(
            node.left,
            primary,
            inverted,
            terms_primary,
            terms_inverted,
            add_to_inverted=add_to_inverted,
        )
    else:
        if add_to_inverted:
            terms_inverted.append(node.left)
        else:
            terms_primary.append(node.left)
    right_exact = node.right.dtype == "float" or (
        hasattr(node.right, "op") and node.right.op.__class__.__name__ in ["Add", "Sub"]
    )
    if (
        node.right.__class__.__name__ == "BinOp"
        and node.right.op.__class__ in [primary, inverted]
        and right_exact
    ):
        if node.op.__class__ is primary:
            collect_commutative(
                node.right,
                primary,
                inverted,
                terms_primary,
                terms_inverted,
                add_to_inverted=add_to_inverted,
            )
        else:
            collect_commutative(
                node.right,
                primary,
                inverted,
                terms_primary,
                terms_inverted,
                add_to_inverted=not add_to_inverted,
            )
    else:
        if (not add_to_inverted and op_primary) or (add_to_inverted and not op_primary):
            terms_primary.append(node.right)
        else:
            terms_inverted.append(node.right)