def test01():
    __cdbkernel__ = create_scope()
    
    def post_process(ex):
        distribute(ex)
        eliminate_kronecker(ex)
        sort_product(ex)
        canonicalise(ex)
        collect_terms(ex)

    \delta{#}::KroneckerDelta;
    \GAMMA{#}::GammaMatrix;
    ex:=\GAMMA^{m} \GAMMA^{n p q} \GAMMA^{n p q m r} \delta^{a b} \delta^{c r};
    # display(ex)
    tst:= \GAMMA^{m} \GAMMA^{n p q} \GAMMA^{n p q m c} \delta^{a b} - @(ex);
    # display(tst)
    assert(tst==0)
    print("Test 01 passed")

test01()

def test02():
    def orderx(var, n):
        cn=Ex(n)
        drop_weight(var, $field=@(cn)$)
        return var

    {A,B,C}::Weight(label=field, value=1);
    ex:=A B C + A B + A A + A C + A B C D;
    orderx(ex, 2)
    tst:= A B C + A B C D - @(ex);
    assert(tst==0)
    print("Test 02 passed")

test02()

def test03():
    __cdbkernel__ = create_scope()
    ex:= Q Q Q Q Q Q;
    converge(ex):
        substitute(_, $Q->A+B, A B->3$, repeat=True)
        distribute(_)
    tst:= A A A A A A + 18 A A A A + 135 A A + 540 + 18 B B B B + 135 B B + B B B B B B - @(ex);
    assert(tst==0)
    print("Test 03 passed")

test03()


# 
# \partial{#}::PartialDerivative;        
# ex:=A_{m n} \partial_{r}{ B^{m n} } \partial^{r}{ Q } + \partial_{m}{ A^{m} } R;  
# num=1
# for partial in ex["\\partial"]:
#     partial.name="k"+str(num)
#     num+=1
# 
# ex;
#     
# # this one still not working?! Does not stop and restart at second product.
# \partial{#}::PartialDerivative;        
# ex:=A_{m n} \partial_{r}{ B^{m n} } \partial^{r}{ Q } + \partial_{m}{ A^{m} } R;  
# for prod in ex["\\prod"]:
#     num=1
#     for partial in prod["\\partial"]:
#         partial.name="k"+str(num)
#         num+=1
# 
# ex;
# print(tree(ex))
#     
# 
# \partial{#}::PartialDerivative;        
# ex:=A_{m n} \partial_{r}{ 3 B^{m n} } \partial^{r}{ 2 Q } + \partial_{m}{ 5 A^{m} } R;  
# for prod in ex["\\prod"]:
#     num=1
#     for partial in prod["\\partial"]:
#         for index in partial.indices():
#             print(index)
#             num+=1
# 
# ex;
# print(tree(ex))
#   

def test04():
    __cdbkernel__ = create_scope()
    \partial{#}::PartialDerivative;
    \dot{#}::Accent;
    NoMomList= [ Ex(r'A_{a? b?}'), Ex(r'T{#}'), Ex(r'S(t)'), Ex(r'a**2'), Ex(r"V'(\phi)"), Ex(r"\partial_{0}{a}") ]
    ex:=Q_{m n \dot{r} p q s_{2} k} + T_{m n \dot{r} p q s_{2} k} + a**2 \partial_{0}{a} A_{m n} \partial_{\dot{r}}{ B_{p q} } S(t) V'(\phi) \partial_{s_{2}}{ C } \partial_{0}{D_{k}} ;
    terms=ex["\\sum"].__next__().children()
    for term in terms:
        if term.name=="\\prod":
            num=1
            for factor in term.children():
                matches = map(lambda x: x.matches(factor), NoMomList)
                if not True in matches:
                    if factor.name=="\\partial":
                        for index in factor.own_indices():
                            k = Ex("k"+str(num))
                            factor.insert(k).append_child(index)
                            factor.insert(Ex("I"))
                        for arg in factor.args():
                            matches = map(lambda x: x.matches(arg), NoMomList)
                            if not True in matches:
                                 arg.append_child(Ex("k"+str(num)))
                                 factor.insert(arg)
                        factor.erase()
                        num+=1
                    else:
                        factor.append_child(Ex("k"+str(num)))
                        num+=1
        else:
            matches = map(lambda x: x.matches(term), NoMomList)
            if not True in matches:
                term.append_child(Ex("k1"))

    tst:= Q_{m n \dot{r} p q s_{2} k}(k1) + T_{m n \dot{r} p q s_{2} k} + a**2 \partial_{0}{a} A_{m n} k1_{\dot{r}} I B_{p q}(k1) S(t) V'(\phi) k2_{s_{2}} I C(k2) k3_{0} I D_{k}(k3) - @(ex);
    assert(tst==0)
    print("Test 04 passed")

test04()

def test05():
    ex:= A_{m n} B^{0 m} + C_{n};
    lst=[]
    for n in ex:
        lst.append(str(n))
    
    tst=['A_{m n} B^{0 m} + C_{n}', 'A_{m n} B^{0 m}', 'A_{m n}', 'm', 'n', 'B^{0 m}', '0', 'm', 'C_{n}', 'n']
    assert(tst==lst)
    print("Test 05 passed")

test05()

# ex:=A + B + C;
# lst=[]
# for sum in ex["\\sum"]:
#     for term in sum:
#         lst.append(term.name)
# 
# assert(lst==["A", "B", "C"])
# print("Test 05 passed")
        
 

# ex:= Q ( A_{m n} (Q+S) + B_{m n} ) + D_{m n};
# for term in ex.top().terms():
#     print("term:")
#     for factor in term.factors():
#         print("factor:")
#         print(factor)
# 
#
        
# \nabla{#}::Derivative;        
# ex:= \nabla_{q}{ T_{m n}^{p} };
# for term in ex.top().terms():
#     

def test_expand(ex):
   tst:= (A??)^{\dagger};
   for node in ex:
      if tst.matches( node ):
         rep=$P$
         lst=[]
         for prod in node["\\prod"]:
            for factor in prod.factors():
               lst.append($ @(factor) $)
         for factor in list(reversed(lst)):
            rep.top().append_child($ @(factor)^{\dagger} $)
         rep.top().name=r"\prod"
         node.replace(rep)
   return ex

def test06():
    __cdbkernel__ = create_scope()
    \dagger::Symbol;
    ex:= (A B C)^{\dagger} + Q + (D E)^{\dagger};
    test_expand(ex)
    tst:= C^{\dagger} B^{\dagger} A^{\dagger} + Q + E^{\dagger} D^{\dagger} - @(ex);
    assert(tst==0)
    print("Test 06 passed")

test06()

def test07():
    __cdbkernel__ = create_scope()
    {a,b,c,d#,z}::Indices.
    ex:= 1/2 D_{b} C^{a b c} A_{c};
    fi=ex.top().free_indices()
    assert(next(fi).name=='a')
    print("Test 07a passed")
    ex:= 1/2 C^{a b c} D_{b} A_{c};
    fi=ex.top().free_indices()
    assert(next(fi).name=='a')
    print("Test 07b passed")
    ex:= 1/2 D_{b} A_{c} C^{a b c};
    fi=ex.top().free_indices()
    assert(next(fi).name=='a')
    print("Test 07c passed")

test07()

def test08():
    __cdbkernel__ = create_scope()
    ex:= A C + B D + C ;
    for term in ex.top().terms():
        term.multiplier *= 3
    
    tst:= 3 A C + 3 B D + 3 C - @(ex);
    assert(tst==0)
    print("Test 08 passed")

test08()

def test09():
    __cdbkernel__ = create_scope()
    ex:= A + B(C+D) + Q(E);
    for node in ex:
       num=2
       # print("node", node)
       for term in node.terms():
          # print("term", term)
          term.multiplier *= num
          num+=1
    
    tst:= 2A + 3 B(2C+3D) + 4Q(2E) - @(ex);
    assert(tst==0)
    print("Test 09 passed")

test09()
     
def test10():     
    __cdbkernel__ = create_scope()
    A_{m n p}::TableauSymmetry(shape={1,1}, indices={1,2});
    p = TableauSymmetry.get($A_{m n p}$)
    p.attach($B_{m n p}$)
    ex:= B_{m n p} - B_{m p n};
    meld(ex)
    tst:= 2 B_{m n p} - @(ex);
    assert(tst==0)
    print("Test 10 passed")

test10()

def test11():
   __cdbkernel__ = create_scope()
   R_{m n p q}::RiemannTensor;
   p = RiemannTensor.get($R_{m n p q}$)
   try:
     p.attach($A_{m n}$)
     assert(1==0)
   except RuntimeError:
     print("Test 11 passed")

test11()

def test12():
   __cdbkernel__ = create_scope()
   sub:=3 q;
   ex:=A^{@(sub)};
   tst:= A^{3*q} - @(ex);
   assert(tst==0)
   print("Test 12 passed")

test12()
     
def test13():
    __cdbkernel__ = create_scope()
    {m,n,k,l}::Indices(flat, values={a,b});
    {\mu,\nu,\rho,\sigma}::Indices(curved, values={0,1,2,3});
    assert( $A_{m n}$.matches($A_{a b}$) == True )
    print("Test 13a passed")
    assert( $A_{m n}$.matches($A_{a a}$) == True )
    print("Test 13b passed")
    assert( $A_{m n}$.matches($A_{a c}$) == False )
    print("Test 13c passed")
    assert( $A_{m n}$.matches($A_{0 1}$) == False )
    print("Test 13d passed")
    assert( $A_{\mu\nu}$.matches($A_{0 1}$) == True )
    print("Test 13e passed")
    assert( $A_{\mu\nu}$.matches($A^{0 1}$) == True )     
    print("Test 13f passed")
    assert( $A_{\mu\nu}$.matches($A_{0 4}$) == False )
    print("Test 13g passed")
    assert( $A_{m n}$.matches($A_{k l}$) == True )
    print("Test 13h passed")
    assert( $A_{m n}$.matches($A_{m n}$) == True )    
    print("Test 13i passed")
    assert( $A_{m n}$.matches($A_{\mu \nu}$) == False )
    print("Test 13j passed")

test13()

def test14():
    __cdbkernel__ = create_scope()
    ex:={A,B} ~ {C,D};
    assert(ex==${A,B,C,D}$)
    print("Test 14a passed")
    ex1:= {A,B};
    ex2:= {C,D};
    ex3 = join(ex1, ex2)
    assert(ex3==${A,B,C,D}$)
    print("Test 14b passed")
    ex1:= A;
    ex2:= {C,D};
    ex3 = join(ex1, ex2)
    assert(ex3==${A,C,D}$)
    print("Test 14c passed")
    ex1:= {A,B};
    ex2:= C;
    ex3 = join(ex1, ex2)
    assert(ex3==${A,B,C}$)
    print("Test 14d passed")
    ex1:= A;
    ex2:= C;
    ex3 = join(ex1, ex2)
    assert(ex3==${A,C}$)
    print("Test 14e passed")
    ex1:= {A,B};
    ex2:= C;
    ex3:= {D, E};
    ex4= join(ex1, ex2, ex3)
    assert(ex4==${A,B,C,D,E}$)
    print("Test 14f passed")
    ex1:= A;
    ex2:= C;
    ex3:= {D, E};
    ex4= join(ex1, ex2, ex3)
    assert(ex4==${A,C,D,E}$)
    print("Test 14g passed")
    ex1:= A;
    ex4 = join(ex1, Ex(""))
    assert(ex4==$A$)
    print("Test 14h passed")
    ex1:= A;
    ex2:= \comma{};
    ex4= join(ex1, ex2)
    assert(ex4==$A$)
    print("Test 14i passed")

test14()

def test15():
    assert(Ex('X').matches(Ex(' '))==False)
    print("Test 15a passed")
    assert(Ex(' ').matches(Ex(' '))==True)
    print("Test 15b passed")
    assert(Ex(' ').matches(Ex('X'))==False)
    print("Test 15c passed")

test15()

def test16():
    # There are many ways in which a match can fail...
    __cdbkernel__ = create_scope()
    {a,b,c}::Indices(vector, position=independent).
    {i,j,k}::Indices(isospin, position=independent).
    assert($\Lambda_{a}$.matches($\Lambda_{i}$)==False)
    assert($\Lambda_{a}$.matches($\Lambda^{i}$)==False)
    print("Test 16 passed")
    
test16()

def test17():
    x::Weight(value=42, label=field);
    tst1 = Weight.get($x$, label="field").value("field")
    assert(tst1==42)
    print("Test 17a passed")
    ex:= 3 a;
    ex.top().multiplier = tst1
    tst2:= 42 a - @(ex);
    assert(tst2==0)
    print("Test 17b passed")

test17()


#  {i,j,k,l}::Indices(cartesian, values={x,y,z}).
#  {\mu,\nu,\rho,\sigma}::Indices(polar, values={r,\theta,\phi}).
#  ex:= x^{i} x^{i} + p_{\mu} p^{\mu};
#  for node in ex:
#       prop = Indices.get(node)
#       if prop is not None:
#           print(node)
#           print(node.name, prop.values)
#  
#  ind = Indices.get($i$);

def test18():
    ex:= A + B + C + D + E;
    s1 = ex[0:2]
    tst:= A + B - @(s1);
    assert(tst==0)
    print("Test 18a passed")
    s2 = ex[0:-2:2]
    tst:= A + C - @(s2);
    assert(tst==0)
    print("Test 18b passed")
    ex:= A + B + C + D + E + F;     
    s3 = ex[0:-2:2]
    tst:= A + C - @(s3);
    print("Test 18c passed")
    ex:= A + B + C + D + E + F;     
    s4 = ex[1:-2:2]
    tst:= B + D - @(s4);
    print("Test 18d passed")
    ex:= A + B + C + D + E;
    s5 = ex[2:3]
    tst:= C - @(s5);
    assert(tst==0)
    print("Test 18e passed")

test18()

# 
# import numpy as np
# a = np.array([[1,2],[3,4]])
# a[0,1]=3
# print(a)
# 
# ex:=[[1,2],[3,4]];
# ex[0][1]=3
# ex;
#