import numpy as np
import json
import os
import sys
import pickle
from select_data import read_db_homology
from threads import Procerssrunner
from read_get_gene_seq import read_gene_sequences
from access_data_rest import update_rest, update_rest_protein
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import IUPAC


def write_fasta(sequences, name):
    with open(name+".fa", "w") as file:
        for seq in sequences:
            if sequences[seq] == "":
                continue
            record = SeqRecord(Seq(sequences[seq], IUPAC.protein), id=seq)
            SeqIO.write(record, file, "fasta")


def read_data_synteny(nop, name):
    smg = []
    sml = []
    indexes = []
    for i in range(nop):
        try:
            with open("temp_"+name+"/thread_"+str(i+1) +
                      "_smg.temp", "rb") as file:
                smg = smg+pickle.load(file)
            with open("temp_"+name+"/thread_"+str(i+1) +
                      "_sml.temp", "rb") as file:
                sml = sml+pickle.load(file)
            with open("temp_"+name+"/thread_"+str(i+1) +
                      "_indexes.temp", "rb") as file:
                indexes = indexes+pickle.load(file)
        except Exception as e:
            print("Problem with thread", i+1, "detected for", name, e)
            continue
    print(len(indexes))
    return smg, sml, indexes


def load_neighbor_genes():
    with open("processed/neighbor_genes.json", "r") as file:
        lsy = dict(json.load(file))
    print(len(lsy))
    print("Neighbor Genes Loaded")
    return lsy


def read_data_homology(dirname):
    lf = os.listdir(dirname)
    if len(lf) == 0:
        print("No Files in the Directory!!!!!!!")
        sys.exit(1)
    a_h = []
    d_h = []
    for x in lf:
        df, n = read_db_homology(dirname, x)
        n = n.split()[0]
        try:
            indexes = np.load("processed/"+n+"_selected_indexes.npy")
        except Exception:
            print("Incomplete data for:", n)
        df = df.loc[indexes]
        a_h.append(df)
        d_h.append(n)
    return a_h, d_h


def main():
    arg = sys.argv
    nop = int(arg[-1])
    n = 3
    a_h, d_h = read_data_homology("data_homology")
    print("Data Read")
    lsy = load_neighbor_genes()
    gene_sequences = read_gene_sequences(
        a_h, lsy, "geneseq", "gene_seq_positive")
    gene_sequences = update_rest(gene_sequences, "gene_seq_positive")
    print("Gene Sequences Loaded.")
    if not os.path.isdir("processed/synteny_matrices"):
        os.mkdir("processed/synteny_matrices")
    ndir = "processed/synteny_matrices/"
    nf1 = "synteny_matrices_global"
    nf2 = "synteny_matrices_local"
    nf3 = "indexes"
    for i in range(len(a_h)):
        df = a_h[i]
        part = len(df)//nop
        pr = Procerssrunner()
        pr.start_processes(nop, df, gene_sequences, lsy, part, n, d_h[i])
        smg, sml, indexes = read_data_synteny(nop, d_h[i])
        print(len(indexes))
        np.save(ndir+str(d_h[i])+"_"+nf1, smg)
        np.save(ndir+str(d_h[i])+"_"+nf2, sml)
        np.save(ndir+str(d_h[i])+"_"+nf3, indexes)
        a_h[i] = df.loc[indexes]
    print("Synteny Matrices Created Successfully :)")
    protein_sequences = read_gene_sequences(
        a_h, lsy, "pro_seq", "pro_seq_positive")
    protein_sequences = update_rest_protein(
        protein_sequences, "pro_seq_positive")
    write_fasta(protein_sequences, "protein_seq_positive")
    print("Protein Sequences Loaded.")


if __name__ == "__main__":
    main()