def test_rewrite():
    x, y, z = symbols('x y z')
    f1 = sin(x) + cos(x)
    assert f1.rewrite(cos,exp) == exp(I*x)/2 + sin(x) + exp(-I*x)/2
    assert f1.rewrite([cos],sin) == sin(x) + sin(x + pi/2, evaluate=False)
    f2 = sin(x) + cos(y)/gamma(z)
    assert f2.rewrite(sin,exp) == -I*(exp(I*x) - exp(-I*x))/2 + cos(y)/gamma(z)

def test_issue_3109():
    from sympy import root, Rational
    I = S.ImaginaryUnit
    assert sqrt(33**(9*I/10)) == -33**(9*I/20)
    assert root((6*I)**(2*I), 3).as_base_exp()[1] == Rational(1, 3) # != 2*I/3
    assert root((6*I)**(I/3), 3).as_base_exp()[1] == I/9
    assert sqrt(exp(3*I)) == exp(3*I/2)
    assert sqrt(-sqrt(3)*(1 + 2*I)) == sqrt(sqrt(3))*sqrt(-1 - 2*I)

    def eval(cls, n, m, theta, phi):
        n, m, theta, phi = [sympify(x) for x in (n, m, theta, phi)]

        # Handle negative index m and arguments theta, phi
        if m.could_extract_minus_sign():
            m = -m
            return S.NegativeOne**m * exp(-2*I*m*phi) * Ynm(n, m, theta, phi)
        if theta.could_extract_minus_sign():
            theta = -theta
            return Ynm(n, m, theta, phi)
        if phi.could_extract_minus_sign():
            phi = -phi
            return exp(-2*I*m*phi) * Ynm(n, m, theta, phi)

 def _eval_evalf(self, prec):
     from sympy import exp, pi, I
     z, period = self.args
     p = periodic_argument(z, period)._eval_evalf(prec)
     if abs(p) > pi or p == -pi:
         return self  # Cannot evalf for this argument.
     return (abs(z)*exp(I*p))._eval_evalf(prec)

def test_conditional():
    X = Geometric('X', S(2)/3)
    Y = Poisson('Y', 3)
    assert P(X > 2, X > 3) == 1
    assert P(X > 3, X > 2) == S(1)/3
    assert P(Y > 2, Y < 2) == 0
    assert P(Eq(Y, 3), Y >= 0) == 9*exp(-3)/2

def test_rewrite():
    x, y, z = symbols('x y z')
    a, b = symbols('a b')
    f1 = sin(x) + cos(x)
    assert f1.rewrite(cos,exp) == exp(I*x)/2 + sin(x) + exp(-I*x)/2
    assert f1.rewrite([cos],sin) == sin(x) + sin(x + pi/2, evaluate=False)
    f2 = sin(x) + cos(y)/gamma(z)
    assert f2.rewrite(sin,exp) == -I*(exp(I*x) - exp(-I*x))/2 + cos(y)/gamma(z)
    assert Max(a, b).rewrite(Piecewise) == Piecewise((a, a >= b), (b, True))
    assert Max(x, y, z).rewrite(Piecewise) == Piecewise((x, (x >= y) & (x >= z)), (y, y >= z), (z, True))
    assert Max(x, y, a, b).rewrite(Piecewise) == Piecewise((a, (a >= b) & (a >= x) & (a >= y)),
        (b, (b >= x) & (b >= y)), (x, x >= y), (y, True))
    assert Min(a, b).rewrite(Piecewise) == Piecewise((a, a <= b), (b, True))
    assert Min(x, y, z).rewrite(Piecewise) == Piecewise((x, (x <= y) & (x <= z)), (y, y <= z), (z, True))
    assert Min(x,  y, a, b).rewrite(Piecewise) ==  Piecewise((a, (a <= b) & (a <= x) & (a <= y)),
        (b, (b <= x) & (b <= y)), (x, x <= y), (y, True))

def _unpolarify(eq, exponents_only, pause=False):
    if not isinstance(eq, Basic) or eq.is_Atom:
        return eq

    if not pause:
        if isinstance(eq, exp_polar):
            return exp(_unpolarify(eq.exp, exponents_only))
        if isinstance(eq, principal_branch) and eq.args[1] == 2*pi:
            return _unpolarify(eq.args[0], exponents_only)
        if (
            eq.is_Add or eq.is_Mul or eq.is_Boolean or
            eq.is_Relational and (
                eq.rel_op in ('==', '!=') and 0 in eq.args or
                eq.rel_op not in ('==', '!='))
        ):
            return eq.func(*[_unpolarify(x, exponents_only) for x in eq.args])
        if isinstance(eq, polar_lift):
            return _unpolarify(eq.args[0], exponents_only)

    if eq.is_Pow:
        expo = _unpolarify(eq.exp, exponents_only)
        base = _unpolarify(eq.base, exponents_only,
            not (expo.is_integer and not pause))
        return base**expo

    if eq.is_Function and getattr(eq.func, 'unbranched', False):
        return eq.func(*[_unpolarify(x, exponents_only, exponents_only)
            for x in eq.args])

    return eq.func(*[_unpolarify(x, exponents_only, True) for x in eq.args])

def test_custom_codegen():
    from sympy.printing.ccode import C99CodePrinter
    from sympy.functions.elementary.exponential import exp

    printer = C99CodePrinter(settings={'user_functions': {'exp': 'fastexp'}})
    gen = C99CodeGen(printer=printer)
    gen.preprocessor_statements.append('#include "fastexp.h"')

    x, y = symbols('x y')
    expr = exp(x + y)

    expected = (
        '#include "expr.h"\n'
        '#include <math.h>\n'
        '#include "fastexp.h"\n'
        'double expr(double x, double y) {\n'
        '   double expr_result;\n'
        '   expr_result = fastexp(x + y);\n'
        '   return expr_result;\n'
        '}\n'
    )

    result = codegen(('expr', expr), header=False, empty=False, code_gen=gen)
    source = result[0][1]
    assert source == expected

def test_twave():
    A1, phi1, A2, phi2, f = symbols('A1, phi1, A2, phi2, f')
    n = Symbol('n')  # Refractive index
    t = Symbol('t')  # Time
    x = Symbol('x')  # Spatial varaible
    k = Symbol('k')  # Wave number
    E = Function('E')
    w1 = TWave(A1, f, phi1)
    w2 = TWave(A2, f, phi2)
    assert w1.amplitude == A1
    assert w1.frequency == f
    assert w1.phase == phi1
    assert w1.wavelength == c/(f*n)
    assert w1.time_period == 1/f
    w3 = w1 + w2
    assert w3.amplitude == sqrt(A1**2 + 2*A1*A2*cos(phi1 - phi2) + A2**2)
    assert w3.frequency == f
    assert w3.wavelength == c/(f*n)
    assert w3.time_period == 1/f
    assert w3.angular_velocity == 2*pi*f
    assert w3.wavenumber == 2*pi*f*n/c
    assert simplify(w3.rewrite('sin') - sqrt(A1**2 + 2*A1*A2*cos(phi1 - phi2)
    + A2**2)*sin(pi*f*n*x*s/(149896229*m) - 2*pi*f*t + atan2(A1*cos(phi1)
    + A2*cos(phi2), A1*sin(phi1) + A2*sin(phi2)) + pi/2)) == 0
    assert w3.rewrite('pde') == epsilon*mu*Derivative(E(x, t), t, t) + Derivative(E(x, t), x, x)
    assert w3.rewrite(cos) == sqrt(A1**2 + 2*A1*A2*cos(phi1 - phi2)
    + A2**2)*cos(pi*f*n*x*s/(149896229*m) - 2*pi*f*t + atan2(A1*cos(phi1)
    + A2*cos(phi2), A1*sin(phi1) + A2*sin(phi2)))
    assert w3.rewrite('exp') == sqrt(A1**2 + 2*A1*A2*cos(phi1 - phi2)
    + A2**2)*exp(I*(pi*f*n*x*s/(149896229*m) - 2*pi*f*t
    + atan2(A1*cos(phi1) + A2*cos(phi2), A1*sin(phi1) + A2*sin(phi2))))

def _lambert(eq, x):
    """
    Given an expression assumed to be in the form
        ``F(X, a..f) = a*log(b*X + c) + d*X + f = 0``
    where X = g(x) and x = g^-1(X), return the Lambert solution if possible:
        ``x = g^-1(-c/b + (a/d)*W(d/(a*b)*exp(c*d/a/b)*exp(-f/a)))``.
    """
    eq = _mexpand(expand_log(eq))
    mainlog = _mostfunc(eq, log, x)
    if not mainlog:
        return []  # violated assumptions
    other = eq.subs(mainlog, 0)
    if (-other).func is log:
        eq = (eq - other).subs(mainlog, mainlog.args[0])
        mainlog = mainlog.args[0]
        if mainlog.func is not log:
            return []  # violated assumptions
        other = -(-other).args[0]
        eq += other
    if not x in other.free_symbols:
        return [] # violated assumptions
    d, f, X2 = _linab(other, x)
    logterm = collect(eq - other, mainlog)
    a = logterm.as_coefficient(mainlog)
    if a is None or x in a.free_symbols:
        return []  # violated assumptions
    logarg = mainlog.args[0]
    b, c, X1 = _linab(logarg, x)
    if X1 != X2:
        return []  # violated assumptions

    u = Dummy('rhs')
    sol = []
    # check only real solutions:
    for k in [-1, 0]:
        l = LambertW(d/(a*b)*exp(c*d/a/b)*exp(-f/a), k)
        # if W's arg is between -1/e and 0 there is
        # a -1 branch real solution, too.
        if k and not l.is_real:
            continue
        rhs = -c/b + (a/d)*l

        solns = solve(X1 - u, x)
        for i, tmp in enumerate(solns):
            solns[i] = tmp.subs(u, rhs)
            sol.append(solns[i])
    return sol

    def eval(cls, a, x):
        # For lack of a better place, we use this one to extract branching
        # information. The following can be
        # found in the literature (c/f references given above), albeit scattered:
        # 1) For fixed x != 0, lowergamma(s, x) is an entire function of s
        # 2) For fixed positive integers s, lowergamma(s, x) is an entire
        #    function of x.
        # 3) For fixed non-positive integers s,
        #    lowergamma(s, exp(I*2*pi*n)*x) =
        #              2*pi*I*n*(-1)**(-s)/factorial(-s) + lowergamma(s, x)
        #    (this follows from lowergamma(s, x).diff(x) = x**(s-1)*exp(-x)).
        # 4) For fixed non-integral s,
        #    lowergamma(s, x) = x**s*gamma(s)*lowergamma_unbranched(s, x),
        #    where lowergamma_unbranched(s, x) is an entire function (in fact
        #    of both s and x), i.e.
        #    lowergamma(s, exp(2*I*pi*n)*x) = exp(2*pi*I*n*a)*lowergamma(a, x)
        from sympy import unpolarify, I
        if x == 0:
            return S.Zero
        nx, n = x.extract_branch_factor()
        if a.is_integer and a.is_positive:
            nx = unpolarify(x)
            if nx != x:
                return lowergamma(a, nx)
        elif a.is_integer and a.is_nonpositive:
            if n != 0:
                return 2*pi*I*n*(-1)**(-a)/factorial(-a) + lowergamma(a, nx)
        elif n != 0:
            return exp(2*pi*I*n*a)*lowergamma(a, nx)

        # Special values.
        if a.is_Number:
            if a is S.One:
                return S.One - exp(-x)
            elif a is S.Half:
                return sqrt(pi)*erf(sqrt(x))
            elif a.is_Integer or (2*a).is_Integer:
                b = a - 1
                if b.is_positive:
                    if a.is_integer:
                        return factorial(b) - exp(-x) * factorial(b) * Add(*[x ** k / factorial(k) for k in range(a)])
                    else:
                        return gamma(a) * (lowergamma(S.Half, x)/sqrt(pi) - exp(-x) * Add(*[x**(k-S.Half) / gamma(S.Half+k) for k in range(1, a+S.Half)]))

                if not a.is_Integer:
                    return (-1)**(S.Half - a) * pi*erf(sqrt(x)) / gamma(1-a) + exp(-x) * Add(*[x**(k+a-1)*gamma(a) / gamma(a+k) for k in range(1, S(3)/2-a)])

 def fdiff(self, argindex=1):
     """
     Returns the first derivative of this function.
     """
     if argindex == 1:
         return exp(*self.args)
     else:
         raise ArgumentIndexError(self, argindex)

def test_discrete_probability():
    X = Geometric('X', S(1)/5)
    Y = Poisson('Y', 4)
    assert P(Eq(X, 3)) == S(16)/125
    assert P(X < 3) == S(9)/25
    assert P(X > 3) == S(64)/125
    assert P(X >= 3) == S(16)/25
    assert P(X <= 3) == S(61)/125
    assert P(Ne(X, 3)) == S(109)/125
    assert P(Eq(Y, 3)) == 32*exp(-4)/3
    assert P(Y < 3) == 13*exp(-4)
    assert P(Y > 3).equals(32*(-S(71)/32 + 3*exp(4)/32)*exp(-4)/3)
    assert P(Y >= 3).equals(32*(-39/32 + 3*exp(4)/32)*exp(-4)/3)
    assert P(Y <= 3) == 71*exp(-4)/3
    assert P(Ne(Y, 3)).equals(
        13*exp(-4) + 32*(-71/32 + 3*exp(4)/32)*exp(-4)/3)
    assert P(X < S.Infinity) is S.One
    assert P(X > S.Infinity) is S.Zero

    def eval(cls, a, x):
        # For lack of a better place, we use this one to extract branching
        # information. The following can be
        # found in the literature (c/f references given above), albeit scattered:
        # 1) For fixed x != 0, lowergamma(s, x) is an entire function of s
        # 2) For fixed positive integers s, lowergamma(s, x) is an entire
        #    function of x.
        # 3) For fixed non-positive integers s,
        #    lowergamma(s, exp(I*2*pi*n)*x) =
        #              2*pi*I*n*(-1)**(-s)/factorial(-s) + lowergamma(s, x)
        #    (this follows from lowergamma(s, x).diff(x) = x**(s-1)*exp(-x)).
        # 4) For fixed non-integral s,
        #    lowergamma(s, x) = x**s*gamma(s)*lowergamma_unbranched(s, x),
        #    where lowergamma_unbranched(s, x) is an entire function (in fact
        #    of both s and x), i.e.
        #    lowergamma(s, exp(2*I*pi*n)*x) = exp(2*pi*I*n*a)*lowergamma(a, x)
        from sympy import unpolarify, I
        if x == 0:
            return S.Zero
        nx, n = x.extract_branch_factor()
        if a.is_integer and a.is_positive:
            nx = unpolarify(x)
            if nx != x:
                return lowergamma(a, nx)
        elif a.is_integer and a.is_nonpositive:
            if n != 0:
                return 2*pi*I*n*(-1)**(-a)/factorial(-a) + lowergamma(a, nx)
        elif n != 0:
            return exp(2*pi*I*n*a)*lowergamma(a, nx)

        # Special values.
        if a.is_Number:
            # TODO this should be non-recursive
            if a is S.One:
                return S.One - exp(-x)
            elif a is S.Half:
                return sqrt(pi)*erf(sqrt(x))
            elif a.is_Integer or (2*a).is_Integer:
                b = a - 1
                if b.is_positive:
                    return b*cls(b, x) - x**b * exp(-x)

                if not a.is_Integer:
                    return (cls(a + 1, x) + x**a * exp(-x))/a

 def fdiff(self, argindex=2):
     from sympy import meijerg, unpolarify
     if argindex == 2:
         a, z = self.args
         return -exp(-unpolarify(z))*z**(a - 1)
     elif argindex == 1:
         a, z = self.args
         return uppergamma(a, z)*log(z) + meijerg([], [1, 1], [0, 0, a], [], z)
     else:
         raise ArgumentIndexError(self, argindex)

def test_expand():
    m0 = OperationsOnlyMatrix([[x*(x + y), 2], [((x + y)*y)*x, x*(y + x*(x + y))]])
    # Test if expand() returns a matrix
    m1 = m0.expand()
    assert m1 == Matrix(
        [[x*y + x**2, 2], [x*y**2 + y*x**2, x*y + y*x**2 + x**3]])

    a = Symbol('a', real=True)

    assert OperationsOnlyMatrix(1, 1, [exp(I*a)]).expand(complex=True) == \
           Matrix([cos(a) + I*sin(a)])

    def eval(cls, a, z):
        from sympy import unpolarify, I, expint
        if z.is_Number:
            if z is S.NaN:
                return S.NaN
            elif z is S.Infinity:
                return S.Zero
            elif z is S.Zero:
                # TODO: Holds only for Re(a) > 0:
                return gamma(a)

        # We extract branching information here. C/f lowergamma.
        nx, n = z.extract_branch_factor()
        if a.is_integer and (a > 0) == True:
            nx = unpolarify(z)
            if z != nx:
                return uppergamma(a, nx)
        elif a.is_integer and (a <= 0) == True:
            if n != 0:
                return -2*pi*I*n*(-1)**(-a)/factorial(-a) + uppergamma(a, nx)
        elif n != 0:
            return gamma(a)*(1 - exp(2*pi*I*n*a)) + exp(2*pi*I*n*a)*uppergamma(a, nx)

        # Special values.
        if a.is_Number:
            # TODO this should be non-recursive
            if a is S.One:
                return exp(-z)
            elif a is S.Half:
                return sqrt(pi)*(1 - erf(sqrt(z)))  # TODO could use erfc...
            elif a.is_Integer or (2*a).is_Integer:
                b = a - 1
                if b.is_positive:
                    return b*cls(b, z) + z**b * exp(-z)
                elif b.is_Integer:
                    return expint(-b, z)*unpolarify(z)**(b + 1)

                if not a.is_Integer:
                    return (cls(a + 1, z) - z**a * exp(-z))/a

 def fdiff(self, argindex=4):
     if argindex == 1:
         # Diff wrt n
         raise ArgumentIndexError(self, argindex)
     elif argindex == 2:
         # Diff wrt m
         raise ArgumentIndexError(self, argindex)
     elif argindex == 3:
         # Diff wrt theta
         n, m, theta, phi = self.args
         return (m * cot(theta) * Ynm(n, m, theta, phi) +
                 sqrt((n - m)*(n + m + 1)) * exp(-I*phi) * Ynm(n, m + 1, theta, phi))
     elif argindex == 4:
         # Diff wrt phi
         n, m, theta, phi = self.args
         return I * m * Ynm(n, m, theta, phi)
     else:
         raise ArgumentIndexError(self, argindex)

    def test_cf(dist, support_lower_limit, support_upper_limit):
        pdf = density(dist)
        t = S('t')
        x = S('x')

        # first function is the hardcoded CF of the distribution
        cf1 = lambdify([t], characteristic_function(dist)(t), 'mpmath')

        # second function is the Fourier transform of the density function
        f = lambdify([x, t], pdf(x)*exp(I*x*t), 'mpmath')
        cf2 = lambda t: mpmath.nsum(lambda x: f(x, t), [support_lower_limit, support_upper_limit], maxdegree=10)

        # compare the two functions at various points
        for test_point in [2, 5, 8, 11]:
            n1 = cf1(test_point)
            n2 = cf2(test_point)

            assert abs(re(n1) - re(n2)) < 1e-12
            assert abs(im(n1) - im(n2)) < 1e-12

    def _shift(func, s):
        z = func.args[-1]
        d = z.collect(x, evaluate=False)
        b = list(d)[0]
        a = d[b]

        if isinstance(a, exp_polar):
            a = exp(a.as_base_exp()[1])
            z = a * b

        t = b.as_base_exp()
        b = t[1] if t[0] is x else S(0)
        r = s / b
        an = (i + r for i in func.args[0][0])
        ap = (i + r for i in func.args[0][1])
        bm = (i + r for i in func.args[1][0])
        bq = (i + r for i in func.args[1][1])

        return a**-r, meijerg((an, ap), (bm, bq), z)