"""
Function of this code: the function to merge the main steps.
Contact for help: gch_liu@163.com 
"""

import numpy as np
import pandas as pd
import joblib
import c1_data_preprocess
import seq2num
# import os
# import sklearn
# import xgboost



def y_get(filename):
    # _So = []  # list to sort

    data_path_06 = filename
    input_data_06, line_index = c1_data_preprocess.start(data_path_06, "seqs")
    # input_num_06 = len(input_data_06)
    input_data_06.iloc[:, 0].to_csv('example_test_name.txt', index=False, header=0)
    input_data_06.iloc[:, 1].to_csv('example_test_data.txt', index=False, header=0)

    seq_test_06 = []
    for line in open("example_test_data.txt"):  # here connect the output of c0_data_preprocess.py
        seq_test_06.append(line)
    num_seq_test_06 = seq2num.seq2num(seq_test_06)
    data_v06 = np.nan_to_num(num_seq_test_06)

    h_v06 = pd.DataFrame(data_v06)
    writer = pd.ExcelWriter('data_predict_example_test.xlsx')
    h_v06.to_excel(writer)
    writer.save()
    data2 = h_v06

    # merge all

    data2.reset_index(drop=True, inplace=True)
    data2.to_csv('datapredict_3labels_example_test.txt', index=False, header=None)

    datatrain1 = pd.read_table('datapredict_3labels_example_test.txt', sep=',', header=None)

    datatrain1[np.isinf(datatrain1)] = np.nan
    datatrain1.to_csv('datapredict_3labels_example_test_v2.txt', index=False, header=None)

    datatrain1 = np.array(datatrain1)
    '''Step2 model 1 (three class) prediction'''
    X_predict = datatrain1

    scaler = joblib.load('model/scaler_3labels_vNCBI_nonLTR_LTR_model_4_GPU.m')
    scaler.clip = False
    X_predict2 = scaler.transform(X_predict)

    model_XGBoost = joblib.load('model/model_xgboost_3labels_vNCBI_nonLTR_LTR_model_4_GPU.m')  # load model
    y_pred = model_XGBoost.predict(X_predict2)

    pred_result = pd.DataFrame({"predict value": np.transpose(y_pred)})
    # virus_filter1 = pred_result[pred_result['predict value'].isin([0])]
    virus_filter2 = datatrain1[pred_result['predict value'].isin([0])]
    line_index2 = np.array(line_index)[pred_result['predict value'].isin([0])]

    pred_result.replace([0, 1, 2], ['Infectious virus', 'Human ERV', 'Non-ERV human'], inplace=True)

 
    '''Step3 model 2 (six class) prediction'''
    X_predict3 = virus_filter2

    scaler = joblib.load('model/scaler_6labels_virus_model_2_GPU.m')
    try:
        scaler.clip = False
        X_predict4 = scaler.transform(X_predict3)

        model_XGBoost2 = joblib.load('model/model_xgboost_6labels_virus_model_2_GPU.m')  # load model

        y_pred2 = model_XGBoost2.predict(X_predict4)

        pred_result2 = pd.DataFrame({"predict value": np.transpose(y_pred2)})

        pred_result2.replace([0, 1, 2, 3, 4, 5],
                             ['double strand DNA (dsDNA)', 'single strand DNA (ssDNA)', 'retro-transcribing',
                              'negative-sense single strand RNA (ssRNA (-))',
                              'positive-sense single strand RNA (ssRNA (+))',
                              'double strand RNA (dsRNA)'], inplace=True)
    # writer=pd.ExcelWriter("result_example_test_final.xlsx")
    # pred_result2.to_excel(writer,sheet_name="predict")
    except:
        y_pred2 = []
    # writer.save()
#    import work

    print('------------the end----------')
    return y_pred, y_pred2, line_index, line_index2


if __name__ == "__main__":
    # display all data
    np.set_printoptions(threshold=np.inf)
    pd.set_option('display.max_columns', None)
    pd.set_option('display.max_rows', None)

    '''Step1 input fasta data'''