""" Assert SymPy Expression Equality """
# Author: Ken Sible
# Email:  ksible *at* outlook *dot* com

from mpmath import mp, mpf, mpc, fabs, log10
import sympy as sp, sys, random

precision = 30

def expand_vardict(vardict):
    if all(not isinstance(vardict[var], list) for var in vardict):
        return vardict
    for var in vardict:
        if not isinstance(vardict[var], list):
            vardict[var] = [vardict[var]]
    return expand_vardict({(var + '[' + str(i) + ']') if len(vardict[var]) > 1 else var: vardict[var][i]
            for var in vardict for i in range(len(vardict[var]))})

def compute_value(symdict, replaced, reduced, factor):
    # determine the precision for floating point arithmetic
    mp.dps = factor * precision
    # substitute unique random number for each free symbol in symdict
    # into every common subexpression from CSE and into main expression
    for var, subexpr in replaced:
        for symbol in subexpr.free_symbols:
            subexpr = subexpr.subs(symbol, symdict[symbol])
        if subexpr.atoms(sp.pi):
            subexpr = subexpr.subs(sp.pi, mp.pi)
        symdict[var] = subexpr
    reduced = reduced[0]
    for symbol in reduced.free_symbols:
        reduced = reduced.subs(symbol, symdict[symbol])
        if reduced.atoms(sp.pi):
            reduced = reduced.subs(sp.pi, mp.pi)
    reduced = reduced.subs(sp.Function('NRPyAbs'), sp.Abs)
    value = mpc(sp.N(reduced)) if isinstance(reduced, complex) else mpf(reduced)
    mp.dps = precision
    return value

def update_vardict(vardict):
    # expand vardict into component mapping
    vardict = expand_vardict(vardict)
    # extract every free symbol present in vardict
    symdict = {symbol: None for var in vardict if isinstance(vardict[var], sp.Basic)
                   for symbol in vardict[var].free_symbols}
    for free_symbol in symdict:
        # seed random number generator with free_symbol hash value
        random.seed(hash(free_symbol))
        # update symdict with mapping: free_symbol -> unique random number
        symdict[free_symbol] = mpf(random.random())
    for var in vardict:
        # apply CSE to every expression in vardict
        replaced, reduced = sp.cse(vardict[var], order='none')
        # calculate value after substituting the unique random number
        # from each free symbol in symdict into every expression in vardict
        value = compute_value(symdict, replaced, reduced, factor=1)
        # double the precision (factor = 2) whenever value within range of zero
        if fabs(value) != mpf('0.0') and fabs(value) < 10 ** ((-2.0/3) * precision):
            _value = compute_value(symdict, replaced, reduced, factor=2)
            if fabs(_value) < 10 ** (-(4.0/3) * precision):
                value = mpf('0.0')
        # update vardict with mapping: variable -> (pseudo) unique number
        vardict[var] = value
    return vardict

def assert_equal(vardict_1, vardict_2, suppress_message=False):
    """ Assert SymPy Expression Equality

        >>> from sympy import sin, cos
        >>> from sympy.abc import x

        >>> assert_equal(sin(2*x), 2*sin(x)*cos(x))
        Assertion Passed!
        >>> assert_equal(cos(2*x), cos(x)**2 - sin(x)**2)
        Assertion Passed!

        >>> assert_equal(cos(2*x), 1 - 2*sin(x)**2)
        Assertion Passed!
        >>> assert_equal(cos(2*x), 1 + 2*sin(x)**2)
        Traceback (most recent call last):
        ...
        AssertionError

        >>> vardict_1 = {'A': sin(2*x), 'B': cos(2*x)}
        >>> vardict_2 = {'A': 2*sin(x)*cos(x), 'B': cos(x)**2 - sin(x)**2}
        >>> assert_equal(vardict_1, vardict_2)
        Assertion Passed!

        >>> assert_equal('(a^2 - b^2) - (a + b)*(a - b)', 0)
        Assertion Passed!
    """
    if not isinstance(vardict_1, dict):
        vardict_1 = {'': vardict_1}
    if not isinstance(vardict_2, dict):
        vardict_2 = {'': vardict_2}
    for var_1, var_2 in zip(vardict_1, vardict_2):
        if not isinstance(vardict_1[var_1], sp.Basic):
            vardict_1[var_1] = sp.sympify(vardict_1[var_1])
        if not isinstance(vardict_2[var_2], sp.Basic):
            vardict_2[var_2] = sp.sympify(vardict_2[var_2])
    # update each vardict with mapping: variable -> (pseudo) unique number
    vardict_1, vardict_2 = update_vardict(vardict_1), update_vardict(vardict_2)
    # assert whether SDA >= (2/3) * precision, implying expression equality
    for var_1, var_2 in zip(vardict_1, vardict_2):
        n_1, n_2 = vardict_1[var_1], vardict_2[var_2]
        if n_1 == n_2: continue
        E_rel = 2 * fabs(n_1 - n_2)/(fabs(n_1) + fabs(n_2))
        assert -log10(E_rel) + 1 >= (2.0/3) * precision
    if not suppress_message:
        print('Assertion Passed!')

if __name__ == "__main__":
    import doctest
    sys.exit(doctest.testmod()[0])
