def update_cdga(gca, diff_dict={}):

    gens_dict = gca.gens_dict()
    differentials = {}
    if diff_dict:
        for gen_name in diff_dict.keys():
            gen = gens_dict[gen_name]
            product_gens = diff_dict[gen_name]
            gen0 = gens_dict[product_gens[0]]
            if type(product_gens[1]) == int: # is a power
                power = product_gens[1]
                mult = gen0
                for i in range(1, power):
                    mult *= gen0
                differentials.update({gen:mult})
            else:
                gen1 = gens_dict[product_gens[1]]
                differentials.update({gen:gen0*gen1})
#     print("differentials", differentials)
    return gca.cdg_algebra(differentials) 

    
def has_preimage(M, element):
    
    d = M.differential()
    deg = element.degree()-1
#     print("differential of ", element, "degree: ", deg)
    Md = d.differential_matrix(deg)
    gens = M.basis(deg)
    try:
        solution = Md.solve_left(vector(QQ,element.basis_coefficients()))
        primitive = sum(map(lambda coef, gen: coef*gen, solution, gens))
        print(element, " has preimage ", primitive)
        return True
    except ValueError:
        return False

def injectivity_step(M, phi, check_dim, names, 
                     degrees, differentials, index=0, depth=0, max_depth=5):
    
    print("Injectivity at dimension ", check_dim)
    if depth > max_depth:
        return [M, phi, names, degrees, differentials, True]
    deg = check_dim -1
    gens = [g for g in M.gens() if g.degree() < check_dim]
#     gens.sort(reverse=True)
    gen_products = []
    for i in range(len(gens)):
        e0 = gens[i]
        d = e0.degree()
        mcd = gcd(d, check_dim)
        if d%2 == 0 and mcd == d:
            power = True
            exp = check_dim / d
            mult = e0
            for j in range(1,exp-1):
                mult *= e0
            gen_products.append([e0, mult, power, exp])
        for j in range(i+1,len(gens)):
            power = False
            e1 = gens[j]
            if e0.degree() + e1.degree() == check_dim:
                gen_products.append([e0, e1, power])
#     print("products at dim ",check_dim,": ", gen_products)
    new_M = M
    new_phi = phi
    new_names = names
    new_degrees = degrees
    new_differentials = differentials
    for pair in gen_products:
        # Looking at the minimal model side
        gen0 = pair[0]
        deg_0 = gen0.degree()
        gen1 = pair[1]
        deg_1 = gen1.degree()
        M_product = gen0*gen1
        # Looking at the AplK side
        gen0_name = list(M.gens_dict().keys())[list(M.gens_dict().values()).index(gen0)]
        aplk_gen0 = phi[deg_0][gen0_name]
        is_power = pair[2]
        power = 0
        if is_power:
            power = pair[3]
            aplk_gen1 = aplk_gen0
            for i in range(1, power):
                aplk_gen1 *= aplk_gen0
        else:
            gen1_name = list(M.gens_dict().keys())[list(M.gens_dict().values()).index(gen1)]
            aplk_gen1 = phi[deg_1][gen1_name]
        aplk_product = aplk_gen0*aplk_gen1
        A = aplk_product.parent()
        # Check conditions
        if (M_product == M(0)) and (aplk_product == A.null_element()): # zero element
            pass
        elif (M_product == M(0)) and (aplk_product != A.null_element()): #check
            if aplk_product.is_closed():
                pass #trivial cohomology class (assuming aplk_product has preimage by differential)
            else:
                raise ValueError('zero associated to a non-closed element in AplK')
        elif (M_product.differential() != 0): # is not closed
            print("The element ",M_product, " is not closed")
            pass   
        elif (M_product.differential() == 0) and aplk_product.is_closed() and has_preimage(M, M_product):
            pass #trivial cohomology class (assuming aplk_product has preimage by differential)
        elif (M_product.differential() == 0) and aplk_product.is_closed(): #key step because M_product doesn't have preimage
            new_gen = 'y'+str(deg)+'_'+str(index)
            if aplk_product != A.null_element():
                try:
                    preimage = aplk_product.solve_primitive()
                except ValueError:
                    print(" The element ",M_product, " has non trivial cohomology class. No generator added.")
                    continue # we don't add a new generator because the element has non trivial cohomology class
                phi_new_gen = {new_gen: preimage}
            else: 
                phi_new_gen = {new_gen: A.null_element()}
            print(" Proceed to kill", M_product, " because has associated trivial cohomology class")
            print("Adding generator "+new_gen)
            new_names = names + [new_gen]
            new_degrees = tuple(list(degrees)+[deg])
            new_differentials = copy.deepcopy(differentials)
            if not is_power:
                new_differentials.update({new_gen: [gen0_name, gen1_name]})
            else:
                new_differentials.update({new_gen: [gen0_name, int(power)]})
            gca = GradedCommutativeAlgebra(QQ, names=new_names, degrees=new_degrees)
            new_M = update_cdga(gca, new_differentials)
            new_phi = copy.deepcopy(phi)
            d = new_phi[deg]
            d.update(phi_new_gen)
            return injectivity_step(new_M, new_phi, check_dim, new_names, 
                             new_degrees, new_differentials, index=index+1, depth=depth+1, max_depth=max_depth)
        else:
            pass

    return [new_M, new_phi, new_names, new_degrees, new_differentials, False]

def get_unmapped_generators(dim, coho_aplk, coho_gens, correspondence):

    dimensions = list(correspondence.keys())
    unmmaped_elements = copy.deepcopy(coho_aplk)
    for gen_key in coho_gens.keys():
        gen = coho_gens[gen_key]
        found = False
        for i in range(1,dim):
            if (not found) and (i in dimensions):
                for j in range(i,dim):
                    if (j in dimensions) and (i+j == dim):
                        gensi = correspondence[i]
                        gensj = correspondence[j]
                        for gi in gensi.values():
                            if not found:
                                for gj in gensj.values():
                                    product = gi.cup_product(gj) 
                                    if product == gen or -product == gen:
                                        found = True
                                        i_name = list(gensi.keys())[list(gensi.values()).index(gi)]
                                        j_name = list(gensj.keys())[list(gensj.values()).index(gj)]
                                        if gen_key in unmmaped_elements.keys():
                                            print("The element ",i_name+"*"+j_name," is already mapped")
                                            del unmmaped_elements[gen_key]
                                        break
    return unmmaped_elements

def surjective_step(A, M, phi, dim, names, degrees, differentials, cohomology_correspondence):
    
    print("Surjectivity at dimension ", dim)
    coho_aplk = A.get_cohomology_aplk(dim)
    cohomology_gens = A.cohomology_gens(dim)
    g = get_unmapped_generators(dim, coho_aplk, cohomology_gens, cohomology_correspondence)
    new_names = copy.deepcopy(names)
    new_degrees = copy.deepcopy(degrees)
    new_phi = copy.deepcopy(phi)
    new_correspondence = copy.deepcopy(cohomology_correspondence)
    phi_n_map = {}
    correspondence_n_map = {}
    for gen_key in g.keys():
        deg = gen_key[0]
        index = gen_key[1]
        new_var = 'x'+str(deg)+'_'+str(index)
        new_names += [new_var]
        new_degrees += [deg]
        phi_n_map[new_var] = g[gen_key]
        correspondence_n_map[new_var] = cohomology_gens[gen_key]
    gca = GradedCommutativeAlgebra(QQ, names=new_names, degrees=new_degrees)
    new_M = update_cdga(gca, differentials)
    new_phi[dim] = phi_n_map
    new_correspondence[dim] = correspondence_n_map
    return [new_M, new_phi, new_names, new_degrees, new_correspondence]
    
def minimal_model(S, max_depth=5, max_dim=5):
    
    fundamental_group = S.fundamental_group()
    if (fundamental_group.order() != 1):
        print("Warning: The space is not simply-connected")
    phi = {}
    cohomology_correspondence = {}
    A = AlgebraAplK(S)
    # First step of the algorithm
    first_dim = A.first_cohomology_group()
    phi_n_map = {}
    correspondence_n_map = {}
    if first_dim < 0:
        return [0,{},{}] # due to the function returns multiple objects
    elif first_dim > max_dim:
        return [0,{},{}] # due to the function returns multiple objects
    g = A.get_cohomology_aplk(first_dim)
    cohomology_gens = A.cohomology_gens(first_dim)
    names = []
    degrees = []
    for gen_key in g.keys():
        deg = gen_key[0]
        index = gen_key[1]
        new_var = 'x'+str(deg)+'_'+str(index)
        names += [new_var]
        degrees += [deg]
        phi_n_map[new_var] = g[gen_key]
        correspondence_n_map[new_var] = cohomology_gens[gen_key]
    differentials={}
    gca = GradedCommutativeAlgebra(QQ, names=names, degrees=degrees)
    M = gca.cdg_algebra(differentials)
    phi[first_dim] = phi_n_map
    cohomology_correspondence[first_dim] = correspondence_n_map
    # Now we have an isomorphism in the first non-null cohomology group
    # we check if phi is a monomorphism
    [M, phi, names, degrees, differentials, limit] = injectivity_step(M, phi, first_dim+1, names, 
                     degrees, differentials, index=0, depth=0, max_depth=5)
    if limit:
        print("The algorithm has reach the injectivity depth-limit at dimension ", first_dim)
        return [M, phi]
    else:
        dim = first_dim + 1
        for i in range(dim, max_dim):
            [M, phi, names, degrees, cohomology_correspondence] = surjective_step(A, M, phi, i, names, degrees, 
                                                       differentials, cohomology_correspondence)
            [M, phi, names, degrees, differentials, limit] = injectivity_step(M, phi, i+1, names, 
                     degrees, differentials, index=0, depth=0, max_depth=5)
            if limit:
                print("The algorithm has reach the injectivity depth-limit at dimension", dim)
                break            
        
    return [M, phi, cohomology_correspondence]

#######################################################
# TESTS
S1 = simplicial_complexes.Sphere(2)
S2 = simplicial_complexes.Sphere(2)
S = S1.wedge(S2, is_mutable=False)
#######################################################
#S = simplicial_complexes.ComplexProjectivePlane()
#######################################################
# S = simplicial_complexes.BrucknerGrunbaumSphere()
#######################################################
# S = simplicial_complexes.ProjectivePlane().wedge(simplicial_complexes.Sphere(2), is_mutable=False)
#######################################################
[sullivan_model, model_morphism, cohomology_correspondence] = minimal_model(S, max_depth=5, max_dim=10)
sullivan_model