import pandas as pd
import sys
import progressbar
import os
from neighbor_genes import read_genome_maps
from process_data import create_data_homology_ls
from read_get_gene_seq import read_gene_sequences
from access_data_rest import update_rest_protein
from prepare_synteny_matrix import write_fasta


def read_database(fname, dirname):
    df = pd.read_csv(dirname+"/"+fname, sep="\t", header=None)
    label_dict = dict(ortholog_one2one=1,
                      other_paralog=0,
                      non_homolog=2,
                      ortholog_one2many=1,
                      ortholog_many2many=1,
                      within_species_paralog=0,
                      gene_split=4)
    label = []
    for _, row in df.iterrows():
        label.append(label_dict[row[7]])
    df = df.assign(label=label)
    df = df.drop(7, axis=1)
    df = df.drop(0, axis=1)
    df.columns = ["gene_stable_id", "species", "homology_gene_stable_id",
                  "homology_species", "goc", "wga", "label"]
    return df


def read_prediction_file_folder(dir_name):
    lf = os.listdir(dir_name)
    a_h = []
    d_h = []
    for x in progressbar.progressbar(lf):
        df = read_database(x, dir_name)
        a_h.append(df)
        d_h.append(x.split(".")[0])
    return a_h, d_h


def create_synteny_features(a_h, d_h, n, a, d, ld, ldg, cmap, cimap, name):
    lsy = create_data_homology_ls(a_h, d_h, n, a, d, ld, ldg, cmap, cimap, 0)
    protein_sequences = read_gene_sequences(
        a_h, lsy, "pro_seq", "prediction_"+name)
    protein_sequences = update_rest_protein(
        protein_sequences, "prediction_"+name)
    write_fasta(protein_sequences, "prediction_"+name)
    print("Protein Sequences Loaded")


def main():
    arg = sys.argv
    dirname = arg[-1]
    a_h, d_h = read_prediction_file_folder(dirname)
    n = 3

    a, d, ld, ldg, cmap, cimap = read_genome_maps()  # read the genome maps
    print("Genome Maps Loaded.")

    create_synteny_features(a_h, d_h, n, a, d, ld, ldg, cmap, cimap, dirname)


if __name__ == "__main__":
    main()