import pytest

from diofant import (Derivative, Expr, Function, I, Integer, Rational, Symbol,
                     cos, cot, exp, log, pi, sin, symbols, tan)


__all__ = ()


def test_diff():
    x, y = symbols('x, y')
    assert Rational(1, 3).diff(x) is Integer(0)
    assert I.diff(x) is Integer(0)
    assert pi.diff(x) is Integer(0)
    assert x.diff(x, 0) == x
    assert (x**2).diff(x, 2, x) == 0
    assert (x**2).diff(x, y, 0) == 2*x
    assert (x**2).diff(x, y) == 0
    pytest.raises(ValueError, lambda: x.diff(1, x))

    a = Symbol('a')
    b = Symbol('b')
    c = Symbol('c')
    p = Integer(5)
    e = a*b + b**p
    assert e.diff(a) == b
    assert e.diff(b) == a + 5*b**4
    assert e.diff(b).diff(a) == 1
    e = a*(b + c)
    assert e.diff(a) == b + c
    assert e.diff(b) == a
    assert e.diff(b).diff(a) == 1
    e = c**p
    assert e.diff(c, 6) == 0
    assert e.diff(c, 5) == 120
    e = c**2
    assert e.diff(c) == 2*c
    e = a*b*c
    assert e.diff(c) == a*b

    f = Function('f')
    assert f(x).diff(x, 2).diff(f(x).diff(x, 1)) == 0


def test_diff2():
    n3 = Integer(3)
    n2 = Integer(2)
    n6 = Integer(6)
    x, c = map(Symbol, 'xc')

    e = n3*(-n2 + x**n2)*cos(x) + x*(-n6 + x**n2)*sin(x)
    assert e == 3*(-2 + x**2)*cos(x) + x*(-6 + x**2)*sin(x)
    assert e.diff(x).expand() == x**3*cos(x)

    e = (x + 1)**3
    assert e.diff(x) == 3*(x + 1)**2
    e = x*(x + 1)**3
    assert e.diff(x) == (x + 1)**3 + 3*x*(x + 1)**2
    e = 2*exp(x*x)*x
    assert e.diff(x) == 2*exp(x**2) + 4*x**2*exp(x**2)


def test_diff3():
    a, b, c = map(Symbol, 'abc')
    p = Integer(5)
    e = a*b + sin(b**p)
    assert e == a*b + sin(b**5)
    assert e.diff(a) == b
    assert e.diff(b) == a + 5*b**4*cos(b**5)
    e = tan(c)
    assert e == tan(c)
    assert e.diff(c) in [cos(c)**(-2), 1 + sin(c)**2/cos(c)**2, 1 + tan(c)**2]
    e = c*log(c) - c
    assert e == -c + c*log(c)
    assert e.diff(c) == log(c)
    e = log(sin(c))
    assert e == log(sin(c))
    assert e.diff(c) in [sin(c)**(-1)*cos(c), cot(c)]
    e = 2**a/log(2)
    assert e.diff(a) == 2**a


def test_diff_no_eval_derivative():
    class My(Expr):
        def __new__(cls, x):
            return Expr.__new__(cls, x)

    x, y = symbols('x y')
    # My doesn't have its own _eval_derivative method
    assert isinstance(My(x).diff(x), Derivative)
    # it doesn't have y so it shouldn't need a method for this case
    assert My(x).diff(y) == 0


def test_speed():
    # this should return in 0.0s. If it takes forever, it's wrong.
    x = Symbol('x')
    assert x.diff(x, 10**8) == 0


def test_deriv_noncommutative():
    A = Symbol('A', commutative=False)
    f = Function('f')
    x = Symbol('x')
    assert A*f(x)*A == f(x)*A**2
    assert A*f(x).diff(x)*A == f(x).diff(x) * A**2
