"""
Module defining `CodeString`, a class for a string of code together with
information about its namespace. Only serves as a parent class, its subclasses
`Expression` and `Statements` are the ones that are actually used.
"""

from collections.abc import Hashable

import sympy

from brian2.parsing.sympytools import str_to_sympy, sympy_to_str
from brian2.utils.logger import get_logger
from brian2.utils.stringtools import get_identifiers

__all__ = ["Expression", "Statements"]

logger = get_logger(__name__)


class CodeString(Hashable):
    """
    A class for representing "code strings", i.e. a single Python expression
    or a sequence of Python statements.

    Parameters
    ----------
    code : str
        The code string, may be an expression or a statement(s) (possibly
        multi-line).

    """

    def __init__(self, code):
        self._code = code.strip()

        # : Set of identifiers in the code string
        self.identifiers = get_identifiers(code)

    code = property(lambda self: self._code, doc="The code string")

    def __str__(self):
        return self.code

    def __repr__(self):
        return f"{self.__class__.__name__}({self.code!r})"

    def __eq__(self, other):
        if not isinstance(other, CodeString):
            return NotImplemented
        return self.code == other.code

    def __ne__(self, other):
        return not self == other

    def __hash__(self):
        return hash(self.code)


class Statements(CodeString):
    """
    Class for representing statements.

    Parameters
    ----------
    code : str
        The statement or statements. Several statements can be given as a
        multi-line string or separated by semicolons.

    Notes
    -----
    Currently, the implementation of this class does not add anything to
    `~brian2.equations.codestrings.CodeString`, but it should be used instead
    of that class for clarity and to allow for future functionality that is
    only relevant to statements and not to expressions.
    """

    pass


class Expression(CodeString):
    """
    Class for representing an expression.

    Parameters
    ----------
    code : str, optional
        The expression. Note that the expression has to be written in a form
        that is parseable by sympy. Alternatively, a sympy expression can be
        provided (in the ``sympy_expression`` argument).
    sympy_expression : sympy expression, optional
        A sympy expression. Alternatively, a plain string expression can be
        provided (in the ``code`` argument).
    """

    def __init__(self, code=None, sympy_expression=None):
        if code is None and sympy_expression is None:
            raise TypeError("Have to provide either a string or a sympy expression")
        if code is not None and sympy_expression is not None:
            raise TypeError(
                "Provide a string expression or a sympy expression, not both"
            )

        if code is None:
            code = sympy_to_str(sympy_expression)
        else:
            # Just try to convert it to a sympy expression to get syntax errors
            # for incorrect expressions
            str_to_sympy(code)
        super().__init__(code=code)

    stochastic_variables = property(
        lambda self: {
            variable
            for variable in self.identifiers
            if variable == "xi" or variable.startswith("xi_")
        },
        doc="Stochastic variables in this expression",
    )

    def split_stochastic(self):
        """
        Split the expression into a stochastic and non-stochastic part.

        Splits the expression into a tuple of one `Expression` objects f (the
        non-stochastic part) and a dictionary mapping stochastic variables
        to `Expression` objects. For example, an expression of the form
        ``f + g * xi_1 + h * xi_2`` would be returned as:
        ``(f, {'xi_1': g, 'xi_2': h})``
        Note that the `Expression` objects for the stochastic parts do not
        include the stochastic variable itself.

        Returns
        -------
        (f, d) : (`Expression`, dict)
            A tuple of an `Expression` object and a dictionary, the first
            expression being the non-stochastic part of the equation and
            the dictionary mapping stochastic variables (``xi`` or starting
            with ``xi_``) to `Expression` objects. If no stochastic variable
            is present in the code string, a tuple ``(self, None)`` will be
            returned with the unchanged `Expression` object.
        """
        stochastic_variables = []
        for identifier in self.identifiers:
            if identifier == "xi" or identifier.startswith("xi_"):
                stochastic_variables.append(identifier)

        # No stochastic variable
        if not len(stochastic_variables):
            return (self, None)

        stochastic_symbols = [
            sympy.Symbol(variable, real=True) for variable in stochastic_variables
        ]

        # Note that collect only works properly if the expression is expanded
        collected = (
            str_to_sympy(self.code).expand().collect(stochastic_symbols, evaluate=False)
        )

        f_expr = None
        stochastic_expressions = {}
        for var, s_expr in collected.items():
            expr = Expression(sympy_expression=s_expr)
            if var == 1:
                if any(s_expr.has(s) for s in stochastic_symbols):
                    raise AssertionError(
                        "Error when separating expression "
                        f"'{self.code}' into stochastic and non-"
                        "stochastic term: non-stochastic "
                        f"part was determined to be '{s_expr}' but "
                        "contains a stochastic symbol."
                    )
                f_expr = expr
            elif var in stochastic_symbols:
                stochastic_expressions[str(var)] = expr
            else:
                raise ValueError(
                    f"Expression '{self.code}' cannot be separated into "
                    "stochastic and non-stochastic "
                    "term"
                )

        if f_expr is None:
            f_expr = Expression("0.0")

        return f_expr, stochastic_expressions

    def _repr_pretty_(self, p, cycle):
        """
        Pretty printing for ipython.
        """
        if cycle:
            raise AssertionError("Cyclical call of 'CodeString._repr_pretty'")
        # Make use of sympy's pretty printing
        p.pretty(str_to_sympy(self.code))

    def __eq__(self, other):
        if not isinstance(other, Expression):
            return NotImplemented
        return self.code == other.code

    def __ne__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return hash(self.code)


def is_constant_over_dt(expression, variables, dt_value):
    """
    Check whether an expression can be considered as constant over a time step.
    This is *not* the case when the expression either:

    1. contains the variable ``t`` (except as the argument of a function that
       can be considered as constant over a time step, e.g. a `TimedArray` with
       a dt equal to or greater than the dt used to evaluate this expression)
    2. refers to a stateful function such as ``rand()``.

    Parameters
    ----------
    expression : `sympy.Expr`
        The (sympy) expression to analyze
    variables : dict
        The variables dictionary.
    dt_value : float or None
        The length of a timestep (without units), can be ``None`` if the
        time step is not yet known.

    Returns
    -------
    is_constant : bool
        Whether the expression can be considered to be constant over a time
        step.
    """
    t_symbol = sympy.Symbol("t", real=True, positive=True)
    if expression == t_symbol:
        return False  # The full expression is simply "t"
    func_name = str(expression.func)
    func_variable = variables.get(func_name, None)
    if func_variable is not None and not func_variable.stateless:
        return False
    for arg in expression.args:
        if arg == t_symbol and dt_value is not None:
            # We found "t" -- if it is not the only argument of a locally
            # constant function we bail out
            if not (
                func_variable is not None
                and func_variable.is_locally_constant(dt_value)
            ):
                return False
        else:
            if not is_constant_over_dt(arg, variables, dt_value):
                return False
    return True
