r"""
<Short one-line summary that ends with no period>

<Paragraph description>

EXAMPLES::

<Lots and lots of examples>

AUTHORS:

- Carlos Alquézar Baeta (2023-02-07): initial version
- Miguel Ángel Marco (2023-02-07): initial version
- Jorge Martín Morales (2023-02-07): initial version

"""

# ****************************************************************************
#       Copyright (C) 2023 Carlos Alquézar Baeta <alquezar@unizar.es>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#                  https://www.gnu.org/licenses/
# ****************************************************************************

def phi_polinomial(pol, model_morphism, A):
    #pol_keys = pol.dict().keys()
    m_gens = pol.parent().gens()
    #m_gens = [g for g,e in zip(gens,sum([vector(k) for k in pol_keys])) if e]
    aplkGens = {g: model_morphism[g.degree()][str(g)] for g in m_gens}
    aplkElement = A.null_element(A)
    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*A.one()
        for i in range(len(exp)):
            e = exp[i]
            if e:
                aGen = aplkGens[m_gens[i]]
                term *= aGen**e
        #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(A,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:
            print("get_ker_element "+str(i),flush=True)
            # 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
        """
        H = A.K.cohomology_ring(base_ring=QQ)
        PiMatrix = H.contraction().pi().to_matrix(dim) # AT_Model
        apl_correspondence = {}
        cochain_correspondence = {}
        Hc_correspondence = {}
        M = zero_matrix(QQ, len(H.basis(dim)), len(hm_basis))
        i = 0
        hm_basis_keys = [g.representative() for g in hm_basis.keys()]
        for gen_repr in hm_basis_keys:
            print("get_ker_element "+str(i),flush=True)
            # 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] = PiMatrix*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)


@parallel(ncpus=Integer(24))
def find_preimage(A, phi, kerElement):
    print("finding preimage for ker_element ", kerElement, flush=True)
    aplkElement = phi_polinomial(kerElement, phi, A)
    preimage = {}
    if aplkElement == A.null_element(A):
        preimage = A.null_element(A)
        #print("Associating Apl(K)'s null element to the new generator", flush=True)
    else:
        preimage = aplkElement.solve_primitive()
    return [kerElement, preimage]

def zip_with_constants(a,b,l):
    return [(a, b, i) for i in l]
    
def injectivity_step(data, check_dim, max_depth=5):
    """"
    The method checks if the phi map induces injectivity at cohomology. If it is 
    not the case, the method adds generators to the GCDA M in order to kill the
    the proper cohomology classes to make the map injective.

    INPUT::
        - data:  A data structure containing the following parameters:
            - sma: The minimal Sullivan algebra under construction
            - morphism: The map between the current state of the Sullivan minimal model and the Apl(K)
            - names: names of the current generators of the sma
            - degrees: the degrees of the generators of the sma
            - differencials: the differential of the generators of the sma
            - AplK: The algebra Apl(K)
            - y_current_indexes: the current number of generators added to each degree in order to build injectivity
            
        - check_dim: The dimension the method is going to check the injectivity
        - max_depth: maximum number of generators to be added for dimension 'check_dim'

    OUTPUT::
        - data: The updated data structure
        - limit_reached: boolean. If 'True', the method has reached the maximum number of generators to be added for dimension 'check_dim'. Otherwise, 'False'.
    """

    print("Injectivity at dimension ", check_dim, flush=True)
    limit_reached = False
    current_index = data['y_current_indexes'][check_dim-1]
    if current_index >= max_depth:
        limit_reached = True
        return [data, limit_reached]

    M = data['msa']
    phi = data['morphism']
    A = data['AplK']
    names = data['names']
    degrees = data['degrees']
    differentials = data['differentials']
    m_kernel = get_ker_elements(M, phi, A, check_dim)
    
    if len(m_kernel) > 0:
        new_M = data['msa']
        new_phi = data['morphism']
        new_names = data['names']
        new_degrees = data['degrees']
        new_differentials = data['differentials']
        if len(m_kernel) > max_depth - current_index:
            limit_reached = True
            m_kernel = m_kernel[:max_depth-current_index]
        
        phi_new_gen = {}
        #preimages = {}
        
        #parallelization
        n = len(m_kernel)
        print("Zip parameters before parallelization", flush=True)
        #parameters = list(zip(n*[A],n*[phi],m_kernel))
        parameters = zip_with_constants(A, phi ,m_kernel)
        print("Finish zipping", flush=True)
        preimages = {i[1][0]:i[1][1] for i in find_preimage(parameters)}
       # for kerElement in m_kernel:
       #    preimages[kerElement] = find_preimage(A, phi, kerElement)
           #print(preimages, flush=True)
           
        kerGens = []
        for kerElement in preimages.keys():
            new_gen = 'y'+str(check_dim-1)+'_'+str(current_index)
            current_index += 1
            print(" Proceed to kill", kerElement,
                  " because has associated trivial cohomology class", flush=True)
            print("Adding generator "+new_gen, flush=True)
            new_names += [new_gen]
            new_degrees = tuple(list(new_degrees)+[check_dim-1])
            kerGens += [(kerElement, new_gen)]
        
        gca = GradedCommutativeAlgebra(
                QQ, names=new_names, degrees=new_degrees)
              
        new_differentials = {convert_to(k, gca): convert_to(
                differentials[k], gca) for k in differentials.keys()}  
        new_phi = {k: {s: phi[k][s]
                for s in phi[k].keys()} for k in phi.keys()}
        for g in kerGens:
            kerElement = g[0]
            new_gen = g[1]
            gen = gca.gens_dict()[new_gen]
            preimage = preimages[kerElement]
            new_differentials.update({gen: convert_to(kerElement, gca)})
            phi_new_gen = {str(gen): preimage}
            d = new_phi[check_dim-1]
            d.update(phi_new_gen)
        new_M = gca.cdg_algebra(new_differentials)
        ## update data
        data['msa'] = new_M
        data['morphism'] = new_phi
        data['names'] = new_names
        data['degrees'] = new_degrees
        data['differentials'] = new_differentials
        data['y_current_indexes'][check_dim-1] = current_index
            
        return injectivity_step(data, check_dim, max_depth=max_depth)
    else:
        return [data, limit_reached]
    #return [new_M, new_phi, new_names, new_degrees, new_differentials, depth_limit]


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", flush=True)
                                            del unmmaped_elements[gen_key]
                                        break
    return unmmaped_elements


def surjective_step(data, dim):

    print("Surjectivity at dimension ", dim, flush=True)
    A = data['AplK']
    cohomology_correspondence = data['cohomology_correspondence']
    names = data['names']
    degrees = data['degrees']
    phi = data['morphism']
    differentials = data['differentials']
    coho_aplk = A.get_cohomology_aplk(A,dim)
    cohomology_gens = A.cohomology_gens(A,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 = gca.cdg_algebra(differentials)
    new_phi[dim] = phi_n_map
    new_correspondence[dim] = correspondence_n_map
    ## update data
    data['msa'] = new_M
    data['morphism'] = new_phi
    data['cohomology_correspondence'] = new_correspondence
    data['names'] = new_names
    data['degrees'] = new_degrees
    data['y_current_indexes'] = {dim:0}
    return data


def build_initial_data(A, dim):
    gens = A.get_cohomology_aplk(A,dim)
    cohomology_gens = A.cohomology_gens(A,dim)
    phi_n_map = {}
    correspondence_n_map = {}
    names = []
    degrees = []
    for gen_key in gens.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] = gens[gen_key]
        correspondence_n_map[new_var] = cohomology_gens[gen_key]
    differentials = {}
    gca = GradedCommutativeAlgebra(QQ, names=names, degrees=degrees)
    M = gca.cdg_algebra(differentials)
    data = {'AplK':A}
    data['msa'] = M # minimal Sullivan Algebra
    data['morphism'] = {dim:phi_n_map}
    data['names'] = names
    data['y_current_indexes'] = {dim:0}
    data['degrees'] = degrees
    data['differentials'] = differentials
    data['cohomology_correspondence'] = {dim:correspondence_n_map}
    return data

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", flush=True)
    A = AlgebraAplK(S)
    # First step of the algorithm
    first_dim = A.first_cohomology_group()
    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
    data = build_initial_data(A, first_dim)    
    # Now we have an isomorphism in the first non-null cohomology group
    # we check if the morphism is a monomorphism at degree +1
    [data, limit] = injectivity_step(data, first_dim+1, max_depth=max_depth)
    if limit:
        print(
            "The algorithm has reach the injectivity depth-limit at dimension ", first_dim, flush=True)
        return [data['msa'], data['morphism']]
    else:
        dim = first_dim + 1
        for i in range(dim, max_dim+1):
            print("Model morphism", flush=True)
            print(data['morphism'], flush=True)
            data = surjective_step(data, i) # SURJECTIVITY
            [data, limit] = injectivity_step(data, i+1, max_depth=max_depth) # INJECTIVITY
            if limit:
                print(
                    "The algorithm has reach the injectivity depth-limit at dimension", i+1, flush=True)
                break

    return [data['msa'], data['morphism'], data['cohomology_correspondence']]