from sage.all import *
import cProfile, pstats
import itertools
load('aplK.py')

def test():
    gens = load("K3_aplKgens")
    S = simplicial_complexes.K3Surface()
    S.set_immutable()
    #S = SimplicialComplex([[0,1,3,4,5],[0,2,6,7,8],[1,2,9,10,11,12],[11,13],[12,13]],is_mutable=False)
    A = AlgebraAplK(S)
    H4 = A.cohomology_group(4)
    #gens_dim2 = gens[2]
    """
    dim = 2
    g2_0 = gens[dim][(dim,0)]
    g2_1 = gens[dim][(dim,1)]
    g2_2 = gens[dim][(dim,2)]
    """
    g4_0 = gens[4][(4,0)]
    products = {}
    elements = {}
    for i in range(len(gens[2].keys())):
        for j in range(i,len(gens[2].keys())):
            g1 = gens[2][(2,i)]
            g2 = gens[2][(2,j)]
            g = g1*g2
            c = g.integral()
            v = A.cochain_to_vector(c,4)
            products[(i,j)] = H4(v)
            elements[(i,j)] = g-H4(v)[0]*g4_0

    save(elements,"K3_gens_prods_deg4")
    primitives = {}

    for key in elements.keys():
        print("Solving primitive of product " + str(key), flush=True)
        g = elements[key]
        f = g.solve_primitive()
        primitives[key] = f
        save(primitives,"K3_primitives_prods_deg4")
        
    for key in elements.keys():
        print("Checking differential of element " + str(key),flush=True)
        g = elements[key]
        f = primitives[key]
        print("f.differential() == g", flush=True)
        print(f.differential() == g, flush=True)

    print("Check finished", flush=True)
    
#############################################
cProfile.runctx("test()", globals(), locals(), "Profile.prof")
s = pstats.Stats("Profile.prof")
s.strip_dirs().sort_stats("cumtime").print_stats()