1from diofant import (ITE, And, Basic, Derivative, Eq, Equivalent, I, Implies,
2                     Integer, Integral, MatrixSymbol, Nand, Nor, Not, Or,
3                     Rational, Symbol, Tuple, Xor, cos, count_ops, exp, pi,
4                     sin, symbols)
5from diofant.abc import a, x, y, z
6
7
8__all__ = ()
9
10
11def test_count_ops_non_visual():
12    def count(val):
13        return count_ops(val, visual=False)
14    assert count(x) == 0
15    assert count(x) is not Integer(0)
16    assert count(x + y) == 1
17    assert count(x + y) is not Integer(1)
18    assert count(x + y*x + 2*y) == 4
19    assert count({x + y: x}) == 1
20    assert count({x + y: 2 + x}) is not Integer(1)
21    assert count(Or(x, y)) == 1
22    assert count(And(x, y)) == 1
23    assert count(Not(x)) == 1
24    assert count(Nor(x, y)) == 2
25    assert count(Nand(x, y)) == 2
26    assert count(Xor(x, y)) == 1
27    assert count(Implies(x, y)) == 1
28    assert count(Equivalent(x, y)) == 1
29    assert count(ITE(x, y, z)) == 1
30    assert count(ITE(True, x, y)) == 0
31
32
33def test_count_ops_visual():
34    ADD, MUL, POW, SIN, COS, AND, D, G = symbols(
35        'Add Mul Pow sin cos And Derivative Integral'.upper())
36    DIV, SUB, NEG = symbols('DIV SUB NEG')
37    NOT, OR, AND, XOR, IMPLIES, EQUIVALENT, _ITE, BASIC, TUPLE = symbols(
38        'Not Or And Xor Implies Equivalent ITE Basic Tuple'.upper())
39
40    def count(val):
41        return count_ops(val, visual=True)
42
43    assert count(7) is Integer(0)
44    assert count(-1) == NEG
45    assert count(-2) == NEG
46    assert count(Rational(2, 3)) == DIV
47    assert count(pi/3) == DIV
48    assert count(-pi/3) == DIV + NEG
49    assert count(I - 1) == SUB
50    assert count(1 - I) == SUB
51    assert count(1 - 2*I) == SUB + MUL
52
53    assert count(x) is Integer(0)
54    assert count(-x) == NEG
55    assert count(-2*x/3) == NEG + DIV + MUL
56    assert count(1/x) == DIV
57    assert count(1/(x*y)) == DIV + MUL
58    assert count(-1/x) == NEG + DIV
59    assert count(-2/x) == NEG + DIV
60    assert count(x/y) == DIV
61    assert count(-x/y) == NEG + DIV
62
63    assert count(x**2) == POW
64    assert count(-x**2) == POW + NEG
65    assert count(-2*x**2) == POW + MUL + NEG
66
67    assert count(x + pi/3) == ADD + DIV
68    assert count(x + Rational(1, 3)) == ADD + DIV
69    assert count(x + y) == ADD
70    assert count(x - y) == SUB
71    assert count(y - x) == SUB
72    assert count(-1/(x - y)) == DIV + NEG + SUB
73    assert count(-1/(y - x)) == DIV + NEG + SUB
74    assert count(1 + x**y) == ADD + POW
75    assert count(1 + x + y) == 2*ADD
76    assert count(1 + x + y + z) == 3*ADD
77    assert count(1 + x**y + 2*x*y + y**2) == 3*ADD + 2*POW + 2*MUL
78    assert count(2*z + y + x + 1) == 3*ADD + MUL
79    assert count(2*z + y**17 + x + 1) == 3*ADD + MUL + POW
80    assert count(2*z + y**17 + x + sin(x)) == 3*ADD + POW + MUL + SIN
81    assert count(2*z + y**17 + x + sin(x**2)) == 3*ADD + MUL + 2*POW + SIN
82    assert count(2*z + y**17 + x + sin(
83        x**2) + exp(cos(x))) == 4*ADD + MUL + 3*POW + COS + SIN
84
85    assert count(Derivative(x, x)) == D
86    assert count(Integral(x, x) + 2*x/(1 + x)) == G + DIV + MUL + 2*ADD
87    assert count(Basic()) is Integer(0)
88
89    assert count({x + 1: sin(x)}) == ADD + SIN
90    assert count([x + 1, sin(x) + y, None]) == ADD + SIN + ADD
91    assert count({x + 1: sin(x), y: cos(x) + 1}) == SIN + COS + 2*ADD
92    assert count({}) is Integer(0)
93    assert count([x + 1, sin(x)*y, None]) == SIN + ADD + MUL
94    assert count([]) is Integer(0)
95
96    assert count(Basic()) == 0
97    assert count(Basic(Basic(), Basic(x, x + y))) == ADD + 2*BASIC
98    assert count(Basic(x, x + y)) == ADD + BASIC
99    assert count(Or(x, y)) == OR
100    assert count(And(x, y)) == AND
101    assert count(And(x**y, z)) == AND + POW
102    assert count(Or(x, Or(y, And(z, a)))) == AND + OR
103    assert count(Nor(x, y)) == NOT + OR
104    assert count(Nand(x, y)) == NOT + AND
105    assert count(Xor(x, y)) == XOR
106    assert count(Implies(x, y)) == IMPLIES
107    assert count(Equivalent(x, y)) == EQUIVALENT
108    assert count(ITE(x, y, z)) == _ITE
109    assert count([Or(x, y), And(x, y), Basic(x + y)]) == ADD + AND + BASIC + OR
110
111    assert count(Basic(Tuple(x))) == BASIC + TUPLE
112    # It checks that TUPLE is counted as an operation.
113
114    assert count(Eq(x + y, 2)) == ADD
115
116
117def test_sympyissue_9324():
118    def count(val):
119        return count_ops(val, visual=False)
120
121    M = MatrixSymbol('M', 10, 10)
122    assert count(M[0, 0]) == 0
123    assert count(2 * M[0, 0] + M[5, 7]) == 2
124    P = MatrixSymbol('P', 3, 3)
125    Q = MatrixSymbol('Q', 3, 3)
126    assert count(P + Q) == 9
127    m = Symbol('m', integer=True)
128    n = Symbol('n', integer=True)
129    M = MatrixSymbol('M', m + n, m * m)
130    assert count(M[0, 1]) == 0
131