import numpy as np
import pytest
import sympy
from numpy.testing import assert_equal

import brian2
from brian2 import (
    DimensionMismatchError,
    Expression,
    Hz,
    Statements,
    get_dimensions,
    ms,
    mV,
    second,
    volt,
)
from brian2.core.preferences import prefs
from brian2.utils.logger import catch_logs


def sympy_equals(expr1, expr2):
    """
    Test that whether two string expressions are equal using sympy, allowing
    e.g. for ``sympy_equals("x * x", "x ** 2") == True``.
    """
    s_expr1 = sympy.nsimplify(sympy.sympify(expr1).expand())
    s_expr2 = sympy.nsimplify(sympy.sympify(expr2).expand())
    return s_expr1 == s_expr2


@pytest.mark.codegen_independent
def test_expr_creation():
    """
    Test creating expressions.
    """
    expr = Expression("v > 5 * mV")
    assert expr.code == "v > 5 * mV"
    assert (
        "v" in expr.identifiers
        and "mV" in expr.identifiers
        and not "V" in expr.identifiers
    )
    with pytest.raises(SyntaxError):
        Expression("v 5 * mV")


@pytest.mark.codegen_independent
def test_split_stochastic():
    tau = 5 * ms
    expr = Expression("(-v + I) / tau")
    # No stochastic part
    assert expr.split_stochastic() == (expr, None)

    # No non-stochastic part -- note that it should return 0 and not None
    expr = Expression("sigma*xi/tau**.5")
    non_stochastic, stochastic = expr.split_stochastic()
    assert sympy_equals(non_stochastic.code, 0)
    assert "xi" in stochastic
    assert len(stochastic) == 1
    assert sympy_equals(stochastic["xi"].code, "sigma/tau**.5")

    expr = Expression("(-v + I) / tau + sigma*xi/tau**.5")
    non_stochastic, stochastic = expr.split_stochastic()
    assert "xi" in stochastic
    assert len(stochastic) == 1
    assert sympy_equals(non_stochastic.code, "(-v + I) / tau")
    assert sympy_equals(stochastic["xi"].code, "sigma/tau**.5")

    expr = Expression("(-v + I) / tau + sigma*xi_1/tau**.5 + xi_2*sigma2/sqrt(tau_2)")
    non_stochastic, stochastic = expr.split_stochastic()
    assert set(stochastic.keys()) == {"xi_1", "xi_2"}
    assert sympy_equals(non_stochastic.code, "(-v + I) / tau")
    assert sympy_equals(stochastic["xi_1"].code, "sigma/tau**.5")
    assert sympy_equals(stochastic["xi_2"].code, "sigma2/tau_2**.5")

    expr = Expression("-v / tau + 1 / xi")
    with pytest.raises(ValueError):
        expr.split_stochastic()


@pytest.mark.codegen_independent
def test_str_repr():
    """
    Test the string representation of expressions and statements. Assumes that
    __str__ returns the complete expression/statement string and __repr__ a
    string of the form "Expression(...)" or "Statements(...)" that can be
    evaluated.
    """
    expr_string = "(v - I)/ tau"
    expr = Expression(expr_string)

    # use sympy to check for equivalence of expressions (terms may have be
    # re-arranged by sympy)
    assert sympy_equals(expr_string, str(expr))
    assert sympy_equals(expr_string, eval(repr(expr)).code)

    # Use exact string equivalence for statements
    statement_string = "v += w"
    statement = Statements(statement_string)

    assert str(statement) == "v += w"
    assert repr(statement) == "Statements('v += w')"


if __name__ == "__main__":
    test_expr_creation()
    test_split_stochastic()
    test_str_repr()
