def phi_polinomial(pol, model_morphism, A):
    
    pol_keys = pol.dict().keys()
    gens = pol.parent().gens()
    m_gens = [g for g,e in zip(gens,sum([vector(k) for k in pol_keys])) if e]
    aplkGens = [model_morphism[g.degree()][str(g)] for g in m_gens]
    aplkElement = A.null_element()
    for exp in pol.dict().keys():
        coef = pol.dict()[exp]
        # To-Do: check if exp and aplkGens are sorted in the same way
        term = coef*prod([gen**e for gen,e in zip(aplkGens, exp)])
        aplkElement += term
    return aplkElement
    
    
def get_ker_elements(sullivan_model, model_morphism, A, dim):
    # Get the basis of Hm
    hm_basis = sullivan_model.cohomology(dim).basis()
    if hm_basis:
        Hc = A.cohomology_group(dim)
        apl_correspondence = {}
        cochain_correspondence = {}
        Hc_correspondence = {}
        M = zero_matrix(QQ,len(Hc.basis()),len(hm_basis))
        i = 0
        hm_basis_keys = [g.representative() for g in hm_basis.keys()]
        for gen_repr in hm_basis_keys:
            # get an element of the basis of Hm
            # get a representative of a cohomology class in M
            # get the apl representative for each generator of M^p
            apl_correspondence[gen_repr] = phi_polinomial(gen_repr, model_morphism, A)
            cochain = apl_correspondence[gen_repr].integral()
            cochain_correspondence[gen_repr] = cochain
            vcochain = A.cochain_to_vector(cochain,dim)
            Hc_correspondence[gen_repr] = Hc(vcochain)
            M[:,i] = Hc_correspondence[gen_repr]
            i+=1
        kernelGens = M.right_kernel().basis()
        m_kernel = [sum([a*b for a,b in zip(hm_basis_keys,g)]) for g in kernelGens]
        return m_kernel
    else:
        return []

def convert_to(element, gca):
    d = element.dict()
    new_e = {tuple(list(k)+[0]): d[k] for k in d.keys()}
    return gca(new_e)
    
def injectivity_step(M, phi, check_dim, names, 
                     degrees, differentials, A, index=0, depth=0, max_depth=5):
    
    print("Injectivity at dimension ", check_dim)
    if depth > max_depth:
        return [M, phi, names, degrees, differentials, True]
    new_M = M
    new_phi = phi
    new_names = names
    new_degrees = degrees
    new_differentials = differentials
    m_kernel = get_ker_elements(M, phi, A, check_dim)
    if len(m_kernel)>0:
        phi_new_gen = {}
        for kerElement in m_kernel:
            new_gen = 'y'+str(check_dim-1)+'_'+str(index)
            print(" Proceed to kill", kerElement, " because has associated trivial cohomology class")
            print("Adding generator "+new_gen)
            new_names = names + [new_gen]
            new_degrees = tuple(list(degrees)+[check_dim-1])
            gca = GradedCommutativeAlgebra(QQ, names=new_names, degrees=new_degrees)
            gen = gca.gens_dict()[new_gen]
            aplkElement = phi_polinomial(kerElement, phi, A)
            print(aplkElement)
            print(aplkElement.is_zero())
            print(aplkElement == A.null_element())
            if aplkElement == A.null_element():
                preimage = A.null_element()
            else:
                preimage = aplkElement.solve_primitive()
            phi_new_gen = {str(gen): preimage}
            new_differentials = {convert_to(k,gca):convert_to(differentials[k],gca) for k in differentials.keys()}
            new_differentials.update({gen: convert_to(kerElement, gca)})
            new_M = gca.cdg_algebra(new_differentials) 
            new_phi = {k: {s:phi[k][s] for s in phi[k].keys()} for k in phi.keys()}
            d = new_phi[check_dim-1]
            d.update(phi_new_gen)
            return injectivity_step(new_M, new_phi, check_dim, new_names, 
                             new_degrees, new_differentials, A, index=index+1, depth=depth+1, max_depth=max_depth)
    else:
        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 = 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 = deepcopy(names)
    new_degrees = deepcopy(degrees)
    new_phi = {k: {s:phi[k][s] for s in phi[k].keys()} for k in phi.keys()}
    new_correspondence = 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, A, 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, A, 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