#Marta Pla Diaz, May 26 (2021), for python2.7


#With this script you obtain the gene names and its reading frame

def abrir_fichero(nomfichero_genes_pangenome):
    dicc_genes_pangenome = {}
    fich = open(nomfichero_genes_pangenome, 'r')
    for line in fich:
        if line.endswith('\n'):  
            line = line[:len(line) - 1]
        if line.endswith('\r'):  
            line = line[:len(line) - 1]
        if line.strip():
            line_div = line.split("_")
            if str(len(line_div)) == "2":
                gene_name_final = line_div[0]+"_"+line_div[1]
                dicc_genes_pangenome[gene_name_final] = line
            if str(len(line_div)) == "3":
                gene_name_final = line_div[1] + "_" + line_div[2]
                dicc_genes_pangenome[gene_name_final] = line
            if str(len(line_div)) == "4":
                gene_name_final = line_div[2] + "_" + line_div[3]
                dicc_genes_pangenome[gene_name_final] = line
            if str(len(line_div)) == "5":
                gene_name_final = line_div[0] + "_" + line_div[4]
                dicc_genes_pangenome[gene_name_final] = line
            if str(len(line_div)) == "6":
                gene_name_final = line_div[4] + "_" + line_div[5]
                dicc_genes_pangenome[gene_name_final] = line
            if str(len(line_div)) == "7":
                gene_name_final = line_div[5] + "_" + line_div[6]
                dicc_genes_pangenome[gene_name_final] = line
            if str(len(line_div)) == "8":
                gene_name_final = line_div[6] + "_" + line_div[7]
                dicc_genes_pangenome[gene_name_final] = line
    return dicc_genes_pangenome

def abrir_ficheros_sense_refs(nomfichero_sense):
    dicc_pauta_genes = {}
    fich = open(nomfichero_sense, 'r')
    for line in fich:
        if line.endswith('\n'):  
            line = line[:len(line) - 1]
        if line.endswith('\r'): 
            line = line[:len(line) - 1]
        line_div = line.split(" ")
        nombre_gen = line_div[0]
        dicc_pauta_genes[nombre_gen] = line_div[1]
    return dicc_pauta_genes


def check_pauta_lectura_genes(dicc_genes_pangenome,dicc_sense_genes_BosniaA,dicc_sense_genes_CDC2,
                              dicc_sense_genes_SS14,dicc_sense_genes_Nichols):
    dicc_sense_pangenome = {}
    for gen in dicc_genes_pangenome:
        if "TENDBA" in gen:
            gen_mod = gen[:len(gen) - 4]
            for i in dicc_sense_genes_BosniaA:
                if gen_mod == i:
                    pauta = dicc_sense_genes_BosniaA[i]
                    name_gen_final = dicc_genes_pangenome[gen]
                    dicc_sense_pangenome[name_gen_final] = pauta
        if "TPECDC2" in gen:
            gen_mod = gen[:len(gen) - 4]
            for i in dicc_sense_genes_CDC2:
                if gen_mod == i:
                    pauta = dicc_sense_genes_CDC2[i]
                    name_gen_final = dicc_genes_pangenome[gen]
                    dicc_sense_pangenome[name_gen_final] = pauta
        if "TPASS" in gen:
            gen_mod = gen[:len(gen) - 4]
            if not "TPASS_2" in gen_mod:
                gen_mod = gen_mod.replace("TPASS_", "TPASS_2")
            for i in dicc_sense_genes_SS14:
                if gen_mod == i:
                    pauta = dicc_sense_genes_SS14[i]
                    name_gen_final = dicc_genes_pangenome[gen]
                    dicc_sense_pangenome[name_gen_final] = pauta
        if "TPANIC" in gen:
            gen_mod = gen[:len(gen) - 4]
            for i in dicc_sense_genes_Nichols:
                if gen_mod == i:
                    pauta = dicc_sense_genes_Nichols[i]
                    name_gen_final = dicc_genes_pangenome[gen]
                    dicc_sense_pangenome[name_gen_final] = pauta
    return dicc_sense_pangenome

def generar_output(dicc_sense_pangenome, nomfichero_genes_pangenome):
    print "\n---> Generando el fichero con las los genes del pangenoma y su sentido de lectura...\n"
    nomfichero_output = nomfichero_genes_pangenome[:len(nomfichero_genes_pangenome) - 4]
    nomfichero_output = nomfichero_output + "_genes_sense.txt"
    fich = open(nomfichero_output, "w")
    for gen in dicc_sense_pangenome:
        fich.write(str(str(gen) + " " + dicc_sense_pangenome[gen])+ "\n")
    print "\n-----------------------------------------------  FIN  -----------------------------------------------------"


def main():
    nomfichero_genes_pangenome = "genes_pangenome_25_05_21.txt"
    nom_fichero_sense_Nichols = "Nichols_NCBI_genes_sense.txt"
    nom_fichero_sense_SS14 = "SS14_NCBI_genes_sense.txt"
    nom_fichero_sense_CDC2 = "CDC2_NCBI_genes_sense.txt"
    nom_fichero_sense_BosniaA = "BosniaA_NCBI_genes_sense.txt"
    dicc_sense_genes_Nichols = abrir_ficheros_sense_refs(nom_fichero_sense_Nichols)
    dicc_sense_genes_SS14 = abrir_ficheros_sense_refs(nom_fichero_sense_SS14)
    dicc_sense_genes_CDC2 =  abrir_ficheros_sense_refs(nom_fichero_sense_CDC2)
    dicc_sense_genes_BosniaA = abrir_ficheros_sense_refs(nom_fichero_sense_BosniaA)
    dicc_genes_pangenome =abrir_fichero(nomfichero_genes_pangenome)
    dicc_sense_pangenome =check_pauta_lectura_genes(dicc_genes_pangenome,dicc_sense_genes_BosniaA,dicc_sense_genes_CDC2,
                                                    dicc_sense_genes_SS14,dicc_sense_genes_Nichols)
    generar_output(dicc_sense_pangenome, nomfichero_genes_pangenome)

if __name__ == '__main__':
    main()