def test_collect_5():
    """Collect with respect to a tuple"""
    a, x, y, z, n = symbols('axyzn')
    assert collect(x**2*y**4 + z*(x*y**2)**2 + z + a*z, [x*y**2, z]) in [
                z*(1 + a + x**2*y**4) + x**2*y**4,
                z*(1 + a) + x**2*y**4*(1 + z) ]
    assert collect((1+ (x+y) + (x+y)**2).expand(), [x,y]) == 1 + y + x*(1 + 2*y) + x**2  + y**2

def test_collect_func():
    f = ((x + a + 1)**3).expand()

    assert collect(f, x) == a**3 + 3*a**2 + 3*a + x**3 + x**2*(3*a + 3) + x*(3*a**2 + 6*a + 3) + 1
    assert collect(f, x, factor) == x**3 + 3*x**2*(a + 1) + 3*x*(a + 1)**2 + (a + 1)**3

    assert collect(f, x, evaluate=False) == {S.One: a**3 + 3*a**2 + 3*a + 1, x: 3*a**2 + 6*a + 3, x**2: 3*a + 3, x**3: 1}

def test_collect_4():
    """Collect with respect to a power"""
    a, b, c, x = symbols('a,b,c,x')

    assert collect(a*x**c + b*x**c, x**c) == x**c*(a + b)
    # issue 6096: 2 stays with c (unless c is integer or x is positive0
    assert collect(a*x**(2*c) + b*x**(2*c), x**c) == x**(2*c)*(a + b)

 def __init__(self, m, **kwargs):
     (a, b), (c, d) = m
     a = cancel(a)
     b = collect(b, ohm)
     c = collect(c, 1/ohm)
     d = cancel(d)
     super(Transmission, self).__init__([[a, b], [c, d]], **kwargs)
     if self.shape != (2, 2):
         raise ValueError("Transmission Matrix must be 2x2")

def test_collect_D():
    D = Derivative
    f = Function('f')
    x,a,b = symbols('xab')
    fx  = D(f(x), x)
    fxx = D(f(x), x,x)

    assert collect(a*fx + b*fx, fx) == (a + b)*fx
    assert collect(a*D(fx,x) + b*D(fx,x), fx)   == (a + b)*D(fx, x)
    assert collect(a*fxx     + b*fxx    , fx)   == (a + b)*D(fx, x)

def calcTF(A, B, C, n):

    s = sympy.Symbol('s')
    I = sympy.eye(n) * s
    symA = makeSymMat(A)


    #sysMat = I - symA
    #sysMat = sysMat.inv()
    #print 'sysMat'
    #print sysMat
    charEqn = (I - symA).det()
    #print 'determinant'
    #print sympy.collect(sympy.expand(sympy.simplify(charEqn)), s, evaluate=False)
    myCharEqn = charEqn
    myTFNum = C * (I - symA).adjugate() * B

    #print 'My characteristic eqn'
    #print myCharEqn
    #print 'numerator'
    #print myTFNum

    alphaDicts = []
    betaDicts = []
    #print 'rows'
    for elem in myTFNum:
        #print elem

        #newEqn = sympy.simplify(sympy.expand(elem))
        #newEqn = sympy.ratsimp(newEqn)
        #tfFrac = sympy.fraction(newEqn)
        elem = sympy.expand(sympy.simplify(elem))

        #alphas are the denominator
        currDictAlpha = sympy.collect(myCharEqn, s, evaluate=False)
        if len(currDictAlpha) > 0:
            alphaDicts.append(currDictAlpha)

        #betas are the numerator
        currDictBeta = sympy.collect(elem, s, evaluate=False)
        if len(currDictBeta) > 0:
            betaDicts.append(currDictBeta)


    #print  'Transfer Function'
    #print myTF
    #print  'betas'
    #print betaDicts
    myBetas = betaDicts
    myAlphas = alphaDicts
    #print 'alphas'
    #print alphaDicts
    #print 'done'

    return(myTFNum, myCharEqn, myAlphas, myBetas)

def test_issue_13143():
    fx = f(x).diff(x)
    e = f(x) + fx + f(x)*fx
    # collect function before derivative
    assert collect(e, Wild('w')) == f(x)*(fx + 1) + fx
    e = f(x) + f(x)*fx + x*fx*f(x)
    assert collect(e, fx) == (x*f(x) + f(x))*fx + f(x)
    assert collect(e, f(x)) == (x*fx + fx + 1)*f(x)
    e = f(x) + fx + f(x)*fx
    assert collect(e, [f(x), fx]) == f(x)*(1 + fx) + fx
    assert collect(e, [fx, f(x)]) == fx*(1 + f(x)) + f(x)

def test_collect_D_0():
    D = Derivative
    f = Function('f')
    x, a, b = symbols('x,a,b')
    fxx = D(f(x), x, x)

    assert collect(a*fxx + b*fxx, fxx) == (a + b)*fxx

    def run(expr):
        # Return semantic (rebuilt expression, factorization candidates)

        if expr.is_Number or expr.is_Symbol:
            return expr, [expr]
        elif expr.is_Indexed or expr.is_Atom:
            return expr, []
        elif expr.is_Add:
            rebuilt, candidates = zip(*[run(arg) for arg in expr.args])

            w_numbers = [i for i in rebuilt if any(j.is_Number for j in i.args)]
            wo_numbers = [i for i in rebuilt if i not in w_numbers]

            w_numbers = collect_const(expr.func(*w_numbers))
            wo_numbers = expr.func(*wo_numbers)

            if aggressive is True and wo_numbers:
                for i in flatten(candidates):
                    wo_numbers = collect(wo_numbers, i)

            rebuilt = expr.func(w_numbers, wo_numbers)
            return rebuilt, []
        elif expr.is_Mul:
            rebuilt, candidates = zip(*[run(arg) for arg in expr.args])
            rebuilt = collect_const(expr.func(*rebuilt))
            return rebuilt, flatten(candidates)
        elif expr.is_Equality:
            rebuilt, candidates = zip(*[run(expr.lhs), run(expr.rhs)])
            return expr.func(*rebuilt, evaluate=False), flatten(candidates)
        else:
            rebuilt, candidates = zip(*[run(arg) for arg in expr.args])
            return expr.func(*rebuilt), flatten(candidates)

    def op(self, op_name, *args):
        print "RUNNING %s"% op_name, len(self.terms)
        new_terms = dict()
        
        for key in self.terms:
            X = self.terms[key]

            r_val = getattr(X, op_name)(*args)
           
            new_key = str(X)
            new_terms[new_key] = X.copy()

            old_val = sympy.Symbol(key)
            new_val = sympy.Symbol(new_key)
            
            if(op_name=="delete" and r_val):
                new_coeff = sympy.Symbol(r_val) 
                self.val = self.val.subs(old_val,old_val*new_coeff)

            self.val = self.val.subs(old_val, new_val)

        self.terms = new_terms

        # Collect like terms
        for X in self.terms:
            self.val = sympy.collect(self.val, sympy.Symbol(X))

def test_gcd_terms():
    f = 2*(x + 1)*(x + 4)/(5*x**2 + 5) + (2*x + 2)*(x + 5)/(x**2 + 1)/5 + (2*x + 2)*(x + 6)/(5*x**2 + 5)

    assert _gcd_terms(f) == ((S(6)/5)*((1 + x)/(1 + x**2)), 5 + x, 1)
    assert _gcd_terms(Add.make_args(f)) == ((S(6)/5)*((1 + x)/(1 + x**2)), 5 + x, 1)

    assert gcd_terms(f) == (S(6)/5)*((1 + x)*(5 + x)/(1 + x**2))
    assert gcd_terms(Add.make_args(f)) == (S(6)/5)*((1 + x)*(5 + x)/(1 + x**2))

    assert gcd_terms((2*x + 2)**3 + (2*x + 2)**2) == 4*(x + 1)**2*(2*x + 3)

    assert gcd_terms(0) == 0
    assert gcd_terms(1) == 1
    assert gcd_terms(x) == x
    assert gcd_terms(2 + 2*x) == Mul(2, 1 + x, evaluate=False)
    arg = x*(2*x + 4*y)
    garg = 2*x*(x + 2*y)
    assert gcd_terms(arg) == garg
    assert gcd_terms(sin(arg)) == sin(garg)

    # issue 3040-like
    alpha, alpha1, alpha2, alpha3 = symbols('alpha:4')
    a = alpha**2 - alpha*x**2 + alpha + x**3 - x*(alpha + 1)
    rep = (alpha, (1 + sqrt(5))/2 + alpha1*x + alpha2*x**2 + alpha3*x**3)
    s = (a/(x - alpha)).subs(*rep).series(x, 0, 1)
    assert simplify(collect(s, x)) == -sqrt(5)/2 - S(3)/2 + O(x)

    # issue 2818
    assert _gcd_terms([S.Zero, S.Zero]) == (0, 0, 1)
    assert _gcd_terms([2*x + 4]) == (2, x + 2, 1)

def residue(expr, x, x0):
    """
    Finds the residue of ``expr`` at the point x=x0.

    The residue is defined as the coefficient of 1/(x-x0) in the power series
    expansion about x=x0.

    Examples:

    >>> from sympy import Symbol, residue, sin
    >>> x = Symbol("x")
    >>> residue(1/x, x, 0)
    1
    >>> residue(1/x**2, x, 0)
    0
    >>> residue(2/sin(x), x, 0)
    2

    This function is essential for the Residue Theorem [1].

    The current implementation uses series expansion to calculate it. A more
    general implementation is explained in the section 5.6 of the Bronstein's
    book [2]. For purely rational functions, the algorithm is much easier. See
    sections 2.4, 2.5, and 2.7 (this section actually gives an algorithm for
    computing any Laurent series coefficient for a rational function). The
    theory in section 2.4 will help to understand why the resultant works in
    the general algorithm. For the definition of a resultant, see section 1.4
    (and any previous sections for more review).

    [1] http://en.wikipedia.org/wiki/Residue_theorem
    [2] M. Bronstein: Symbolic Integration I, Springer Verlag (2005)

    """
    from sympy import collect

    expr = sympify(expr)
    if x0 != 0:
        expr = expr.subs(x, x + x0)
    s = expr.series(x, 0, 0).removeO()
    # TODO: this sometimes helps, but the series expansion should rather be
    # fixed, see #1627:
    if s == 0:
        s = expr.series(x, 0, 1).removeO()
    if s == 0:
        s = expr.series(x, 0, 6).removeO()
    s = collect(s, x)
    if x0 != 0:
        s = s.subs(x, x - x0)
    a = Wild("r", exclude=[x])
    c = Wild("c", exclude=[x])
    r = s.match(a / (x - x0) + c)
    if r:
        return r[a]
    elif isinstance(s, Add):
        # TODO: this is to overcome a bug in match (#1626)
        for t in s.args:
            r = t.match(a / (x - x0) + c)
            if r:
                return r[a]
    return Integer(0)

def coeff(expr, term):
    expr = sp.collect(expr, term)
    symbols = list(term.atoms(sp.Symbol))
    w = sp.Wild("coeff", exclude=symbols)
    m = expr.match(w * term + sp.Wild("rest"))
    if m:
        return m[w]

 def __init__(self):
     h, g, h0, g0, a, d = sympy.symbols('h g h0 g0 a d')
     #g1 = sympy.Max(-d, g0 + h0 - 1/(4*a))
     g1 = g0 + h0 - 1/(4*a)
     h1 = h0 - 1/(2*a)
     parabola = d - a*h*h       # =g on boundary
     slope = g1 + h - h1        # =g on line with slope=1 through z1
     h2 = sympy.solve(parabola-slope,h)[0] # First solution is above z1
     g2 = h2 - h1 + g1          # Line has slope of 1
     g3 = g1
     h3 = sympy.solve((parabola-g).subs(g, g1),h)[1] # Second is on right
     r_a = sympy.Rational(1,24) # a=1/24 always
     self.h = tuple(x.subs(a,r_a) for x in (h0, h1, h2, h3))
     self.g = tuple(x.subs(a,r_a) for x in (g0, g1, g2, g3))
     def integrate(f):
         ia = sympy.integrate(
             sympy.integrate(f,(g, g1, g1+h-h1)),
             (h, h1, h2)) # Integral of f over right triangle
         ib = sympy.integrate(
             sympy.integrate(f,(g, g1, parabola)),
             (h, h2, h3)) # Integral of f over region against parabola
         return (ia+ib).subs(a, r_a)
     i0 = integrate(1)               # Area = integral of pie slice
     E = lambda f:(integrate(f)/i0)  # Expected value wrt Lebesgue measure
     sigma = lambda f,g:E(f*g) - E(f)*E(g)
     self.d = d
     self.Eh = collect(E(h),sympy.sqrt(d-g0-h0+6))
     self.Eg = E(g)
     self.Sigmahh = sigma(h,h)
     self.Sigmahg = sigma(h,g)
     self.Sigmagg = sigma(g,g)
     return

def collect_into_dict_include_zero_and_constant_terms(expr, syms):
    
    expr_terms_dict = sympy.collect(expr,syms,exact=True,evaluate=False)
    for sym in syms:
        if sym not in expr_terms_dict.keys(): expr_terms_dict[sym] = 0
    if 1 not in expr_terms_dict: expr_terms_dict[1] = 0
    return expr_terms_dict

 def exponential_euler_code(self):
     '''
     Generates Python code for an exponential Euler step.
     Not efficient for the moment!
     '''
     all_variables = self._eq_names + self._diffeq_names + self._alias.keys() + ['t']
     vars_tmp = [name + '__tmp' for name in self._diffeq_names]
     lines = ','.join(self._diffeq_names) + '=P._S\n'
     lines += ','.join(vars_tmp) + '=P._dS\n'
     for name in self._diffeq_names:
         # Freeze
         namespace = self._namespace[name]
         expr = optimiser.freeze(self._string[name], all_variables, namespace)
         # Find a and b in dx/dt=a*x+b
         sym_expr = symbolic_eval(expr)
         if isinstance(sym_expr, float):
             lines += name + '__tmp[:]=' + name + '+(' + str(expr) + ')*dt\n'
         else:
             sym_expr = sym_expr.expand()
             sname = sympy.Symbol(name)
             terms = sympy.collect(sym_expr, name, evaluate=False)
             if sname ** 0 in terms:
                 b = terms[sname ** 0]
             else:
                 b = 0
             if sname in terms:
                 a = terms[sname]
             else:
                 a = 0
             lines += name + '__tmp[:]=' + str(-b / a + (sname + b / a) * sympy.exp(a * sympy.Symbol('dt'))) + '\n'
     lines += 'P._S[:]=P._dS'
     #print lines
     return compile(lines, 'Exponential Euler update code', 'exec')

def sympy_polynomial_to_function(expr, symbs,
                                 func_name='sympy_polynomial',
                                 return_string=False):
    poly = sp.poly(sp.collect(sp.expand(expr), symbs), *symbs)
    input_names = map(attrgetter('name'), poly.gens)
    num_names = len(input_names)

    def term_string((powers, coeff)):
        indices = filter(lambda i: powers[i] != 0,
                         range(num_names))
        if indices:
            terms = map(lambda i: '(%s ** %s)' % (input_names[i],
                                                  powers[i]),
                        indices)
            return '%s * %s' % (float(coeff), ' * '.join(terms))
        else:
            return '%s' % float(coeff)

    poly_str = ' + '.join(map(term_string, poly.terms()))
    function_str = 'def %s(%s):\n    return %s' % (
        func_name, ', '.join(input_names), poly_str)

    if return_string:
        return function_str

    exec_env = {}
    exec function_str in exec_env
    return exec_env[func_name]

def get_conditionally_linear_system(eqs):
    '''
    Convert equations into a linear system using sympy.
    
    Parameters
    ----------
    eqs : `Equations`
        The model equations.
    
    Returns
    -------
    coefficients : dict of (sympy expression, sympy expression) tuples
        For every variable x, a tuple (M, B) containing the coefficients M and
        B (as sympy expressions) for M * x + B
    
    Raises
    ------
    ValueError
        If one of the equations cannot be converted into a M * x + B form.

    Examples
    --------
    >>> from brian2 import Equations
    >>> eqs = Equations("""
    ... dv/dt = (-v + w**2) / tau : 1
    ... dw/dt = -w / tau : 1
    ... """)
    >>> system = get_conditionally_linear_system(eqs)
    >>> print(system['v'])
    (-1/tau, w**2.0/tau)
    >>> print(system['w'])
    (-1/tau, 0)

    '''
    diff_eqs = eqs.substituted_expressions
    
    coefficients = {}
    
    for name, expr in diff_eqs:
        var = sp.Symbol(name, real=True)        
        
        # Coefficients
        wildcard = sp.Wild('_A', exclude=[var])
        #Additive constant
        constant_wildcard = sp.Wild('_B', exclude=[var])    
        pattern = wildcard*var + constant_wildcard
    
        # Factor out the variable
        s_expr = sp.collect(expr.sympy_expr.expand(), var)
        matches = s_expr.match(pattern)
        
        if matches is None:
            raise ValueError(('The expression "%s", defining the variable %s, '
                             'could not be separated into linear components') %
                             (s_expr, name))
        coefficients[name] = (matches[wildcard].simplify(),
                              matches[constant_wildcard].simplify())
    
    return coefficients

def list_collect_v1(row, listin):
    rowout = copy.copy(row)
    for c, item in enumerate(rowout):
        temp = item
        for var in listin:
            temp = sympy.collect(temp, var)
        rowout[c] = temp
    return rowout

def test_collect_Wild():
    """Collect with respect to functions with Wild argument"""
    a, b, x, y = symbols('a b x y')
    f = Function('f')
    w1 = Wild('.1')
    w2 = Wild('.2')
    assert collect(f(x) + a*f(x), f(w1)) == (1 + a)*f(x)
    assert collect(f(x, y) + a*f(x, y), f(w1)) == f(x, y) + a*f(x, y)
    assert collect(f(x, y) + a*f(x, y), f(w1, w2)) == (1 + a)*f(x, y)
    assert collect(f(x, y) + a*f(x, y), f(w1, w1)) == f(x, y) + a*f(x, y)
    assert collect(f(x, x) + a*f(x, x), f(w1, w1)) == (1 + a)*f(x, x)
    assert collect(a*(x + 1)**y + (x + 1)**y, w1**y) == (1 + a)*(x + 1)**y
    assert collect(a*(x + 1)**y + (x + 1)**y, w1**b) == a*(x + 1)**y + (x + 1)**y
    assert collect(a*(x + 1)**y + (x + 1)**y, (x + 1)**w2) == (1 + a)*(x + 1)**y
    assert collect(a*(x + 1)**y + (x + 1)**y, w1**w2) == (1 + a)*(x + 1)**y