"""
Function of this code: run this function to do the whole methods.
Contact for help: gch_liu@163.com 
"""

import os

from main import y_get
import re
from Bio import SeqIO
import warnings

warnings.filterwarnings('ignore')


def Isother(string):
    pattern = re.compile(r'[^A|C|G|T|a|c|g|t]')
    string = string.replace('\n', '')
    string = string.replace('\t', '')
    ish = re.search(pattern, string)
    # print(ish)
    return ish


def get_namefrc():
    names = []
    with open(filename) as f:
        names = f.readlines()
    check = {}
    for name in names:
        if name == '\n':
            continue
        elif name[0] != '>':
            if Isother(name.strip('\n'))==None and check[name_check]==0:
                check[name_check] = 0
            else:
                check[name_check] = 1
        else:
            name_check = name
            check[name_check] = 0
    idx = 1
    for name in names:
        if name[0] == '>':
            frc_name[name] = idx
            if check[name] == 0:
                reasonable_name.append(name)
            idx = idx + 1
    # print(len(reasonable_name))
    # for name in reasonable_name:
    #    print(name)
    return


def get_class1():
    # print(type(y_pred))
    idx = 0
    y_pred,_,y_index,_ = y_get(filename)
    for i in range(len(y_pred)):
        key = next((k for k, v in frc_name.items() if v == (y_index[i]+1)), None)
        if y_pred[i] == 0:
            frc_type1[key] = "Infectious virus"
            reasonable_name_type.append(reasonable_name[idx])
        elif y_pred[i] == 1:
            frc_type1[key] = "Non-ERV human"
        elif y_pred[i] == 2:

            frc_type1[key] = "ERV"
        idx = idx + 1
    # print(len(reasonable_name))
    # print(len(reasonable_name_type))


def getclass_2():
    # print(len(y_pred2))
    idx = 0
    _, y_pred1,_, y_index1 = y_get(filename)
    for i in range(len(y_pred1)):
        key = next((k for k, v in frc_name.items() if v == (y_index1[i] + 1)), None)
        if y_pred1[i] == 0:
            frc_type2[key] = "Double strand DNA (dsDNA)"
        elif y_pred1[i] == 1:
            frc_type2[key] = "Single strand DNA (ssDNA)"
        elif y_pred1[i] == 2:
            frc_type2[key] = "Retro-transcribing"
        elif y_pred1[i] == 3:
            frc_type2[key] = "Negative-sense single strand RNA (ssRNA (-))"
        elif y_pred1[i] == 4:
            frc_type2[key] = "Positive-sense single strand RNA (ssRNA (+))"
        elif y_pred1[i] == 5:
            frc_type2[key] = "Double strand RNA (dsRNA)"
        idx = idx + 1
    return


def final_work():
    with open(filename) as f:
        names = f.readlines()
    for name in names:
        # as name
        tmp = []
        if name[0] == '>':
            # id
            tmp.append(frc_name[name])
            # name
            tmp.append(name.strip('\n'))
            # type_1
            if name in frc_type1:
                tmp.append(frc_type1[name])
            else:
                tmp.append("Containing non-A/C/G/T/a/c/g/t symbol")
            # type_2
            if name in frc_type2:
                tmp.append(frc_type2[name])
            else:
                tmp.append("/")
            final.append(tmp)


# id name type_1 type_2


def _Write_to_excel():
    import pandas as pd
    for i in range(0, 4):
        result.append([])
    for i in range(0, 4):
        for item in final:
            result[i].append(item[i])
    print(result[0])
    dic1 = {'ID': result[0], 'Name': result[1], 'Three classes': result[2],
            'Six virus subgroups': result[3]}
    df = pd.DataFrame(dic1)
    print(1)
    df.to_excel("output_files/{}_output_results.xlsx".format(file[:-6]), index=False)


if __name__ == "__main__":
    dirlist = os.listdir("input_files")
    is_remove = 0
    for filename in os.listdir("input_files"):
        file = filename
        filename = "input_files/" + filename
        if file[-5:] == "fastq":
            savename = filename[:-1] + "a"
            SeqIO.convert(filename, "fastq", savename, "fasta")
            filename = savename
        is_remove = 1

        frc_name = {}
        check = []
        frc_type1 = {}
        frc_type2 = {}
        reasonable_name = []
        reasonable_name_type = []
        final = []
        result = []

        # y_pred,_,line_index = y_get(filename)
        get_namefrc()
        get_class1()
        getclass_2()
        final_work()
        _Write_to_excel()
        os.remove("datapredict_3labels_example_test_v2.txt")
        os.remove("datapredict_3labels_example_test.txt")
        os.remove("example_test_data.txt")
        os.remove("data_predict_example_test.xlsx")
        os.remove("example_test_name.txt")
        os.remove("seqspreprossed_seq_name.txt")
        os.remove("seqspreprossed_seq_data.txt")
#        if is_remove:
#            os.remove(savename)
