import subprocess
import os
import pandas as pd
import traceback
import argparse
from Alignment_Part import alignment
def runCmd(cmd) :
  res = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  sout ,serr = res.communicate()
  return res.returncode, sout, serr
def dir_build(dir_path):
    cmd = f'mkdir -p {dir_path}'
    returncode, sout, serr = runCmd(cmd)
    print(cmd)
    if returncode == 0:
        final_result = "Build Dir."
    else:
        final_result = serr.decode("gbk")
    return returncode, final_result
def soft_build(input_file,dir_path):
    input_file = os.path.abspath(input_file)
    output_file = dir_path + os.path.basename(input_file)
    print(os.path.splitext(input_file)[-1])
    if os.path.splitext(input_file)[-1] != ".sorted":
        output_file = dir_path + os.path.basename(input_file) + ".sorted"
    cmd = f'ln -s {input_file} {output_file}'
    returncode, sout, serr = runCmd(cmd)
    print(cmd)
    if returncode == 0:
        final_result = "Soft Link Build."
    else:
        final_result = serr.decode("gbk")
    return returncode, final_result
def miRNAminerCLI(genome_file,minMirLen,maxMirLen,dir_path,thread,sRNAminer):
    cmd = f'ls {dir_path}*sorted | parallel -j {thread} "java -cp {sRNAminer} biocjava.sRNA.Miner.miRNA.miRNAminerCLI ' \
          f'--inAln {{}} --outputPrefix {{}} --inGenome {genome_file} --minMiRLen {minMirLen} --maxMiRLen {maxMirLen} 2>{{}}.sRNAminer.log"'

    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = "miRNAminer Step 1: Mining from Lib [Finish]"
    else:
        final_result = serr.decode('gbk')
    return returncode, final_result


def DetectKnownMIR(sRNAminer,dir_path,thread,sRNAminer_database):
    cmd = f'ls {dir_path}*tab | parallel -j {thread} "java -cp {sRNAminer} ' \
          f'biocjava.sRNA.Miner.miRNA.DetectKnownMiR --inMiRList {{}} --miRBaseDB {sRNAminer_database} --outList {{}}.known' \
          f' --matureColIndex 6 --starColIndex 7 --starAbundanceIndex 11"'
    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = "miRNAminer Step 2: Dececting Known miRNA [Finish]"
    else:
        final_result = serr.decode('gbk')
    return returncode, final_result

def filter(input_file):
    returncode = 0
    try:
        output_file = input_file + ".filter"
        miRNA_filter = pd.read_table(input_file,encoding='utf-8')
        miRNA_filter = miRNA_filter[miRNA_filter["Flag"]!="Unlikely"]
        miRNA_filter.to_csv(output_file, sep="\t",index=None)
        success = "miRNAminer Step 3: Filtering [Finish]"
        return returncode,success
    except Exception:
        returncode += 1
        error_str = traceback.format_exc()
        return returncode,error_str

def make_list(dir_path,file_list):
    try:
        file_list = [os.path.basename(i) + ".tab.known.filter" for i in file_list]

        output_file = open(dir_path + "all.miRNA.list", "w+")
        for root, dirs, files in os.walk(dir_path):
            for file in files:
                if file in file_list:
                    output_file.write(root + file + "\n")
        output_file.close()
    except Exception:
        error_str = traceback.format_exc()
        print(error_str)


def MIRmerger(dir_path,sRNAminer):
    cmd = "java -cp " + sRNAminer + " biocjava.sRNA.Miner.miRNA.MIRmerger --inFileList " +  dir_path  + "all_miRNA.list --outFile " + dir_path + "merged.miRNA.list.xls"

    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = "miRNAminer Step 4: Merging [Finish]"
    else:
        final_result = serr.decode('gbk')
    return returncode, final_result

def DetectKnownMIR_2(dir_path,sRNAminer,sRNAminer_database):
    cmd = "java -cp " + sRNAminer + " biocjava.sRNA.Miner.miRNA.DetectKnownMiR --inMiRList " + dir_path + "merged.miRNA.list.xls" \
          " --miRBaseDB " + sRNAminer_database +" --outList "+ dir_path + "merged.miRNA.list.anno.xls"

    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = "miRNAminer Step 5: Detecting Known miRNA--Merge [Finish]"
    else:
        final_result = serr.decode('gbk')
    return returncode, final_result

def MIRnamer(dir_path,sRNAminer):
    cmd = "java -cp " + sRNAminer + " biocjava.sRNA.Miner.miRNA.MIRnamer --inFile " + dir_path +\
          "merged.miRNA.list.anno.xls --outFile " + dir_path + "merged.miRNA.list.anno.rename.xls"

    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = "miRNAminer Step 6: Renaming miRNA [Finish]"
    else:
        final_result = serr.decode('gbk')
    return returncode, final_result
def for_miRNAinfo(dir_path):
    returncode = 0
    try:
        input_file_1 = dir_path + "merged.miRNA.list.anno.rename.xls"
        output_file = dir_path + "miRNAinfo.tab"
        miRNAInfo = pd.read_table(input_file_1,encoding='utf-8')
        miRNAInfo["miRNA_ID-Flag"] = miRNAInfo["miRNA_ID"] + "-" +miRNAInfo["Flag"]
        miRNAInfo = miRNAInfo[["Chr_ID","Precursor_Start","Precursor_End","miRNA_ID-Flag","Strand","Mature_Seq","Star_Seq"]]
        miRNAInfo.to_csv(output_file, sep="\t",index=None)
        success = "miRNAminer Step 7: Establishing miRNAinfo.tab [Finish]"
        return returncode,success
    except Exception:
        returncode += 1
        error_str = traceback.format_exc()
        return returncode,error_str
def MIRinfo2GFF(dir_path,genome_file,sRNAminer):
    cmd = "java -cp " + sRNAminer + " biocjava.sRNA.Miner.miRNA.MIRinfo2GFF --inGenomeFa " + genome_file +\
          " --inMIRinfo " + dir_path + "miRNAinfo.tab --outGff3 " +  dir_path +"miRNA.gff3"

    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = "miRNAminer Step 8: Establishing a GFF file [Finish]"
    else:
        final_result = serr.decode('gbk')
    return returncode, final_result

def miRNA_identification_merger_pipeline(genome_file,input_file_1,minMirLen,maxMirLen,Batch_mode,dir_path,thread,sRNAminer,sRNAminer_database):
    print("Start...")
    n = 0

    if not dir_path.endswith("/") or dir_path.endswith("\\"):
        dir_path = dir_path + "/"
    final_dir_path = dir_path + "miRNA_identification/"
    while os.path.exists(final_dir_path):
        n += 1
        final_dir_path = dir_path + "miRNA_identification_" + str(n) + "/"
    dir_path = final_dir_path
    returncode, final_result = dir_build(dir_path)
    print(final_result)
    if returncode == 0:
        log_file = dir_path + "miRNA_identification.log"
        log_file_content = open(log_file,"w+")
        returncode,final_result = alignment.build_index_TBtools(genome_file, sRNAminer)
        print(final_result, file=log_file_content)
        if returncode == 0:
            print(final_result)
            if Batch_mode:
                input_file_content = open(input_file_1,"r")
                file_list = [i.strip() for i in input_file_content.readlines()]
                new_file_list = []
                for input_file in file_list:
                    returncode, final_result = soft_build(input_file, dir_path)
                if returncode == 0:
                    returncode, final_result = miRNAminerCLI(genome_file,minMirLen,maxMirLen,dir_path,thread,sRNAminer)
                    print(final_result, file=log_file_content)
                if returncode == 0:
                    print(final_result)
                    returncode, final_result = DetectKnownMIR(sRNAminer,dir_path,thread,sRNAminer_database)
                    print(final_result, file=log_file_content)
                if returncode == 0:
                    print(final_result)
                    for input_file in file_list:
                        input_file = dir_path + os.path.basename(input_file) + ".tab.known"
                        returncode, final_result = filter(input_file)
                        input_file = dir_path + os.path.basename(input_file) + ".filter"
                        new_file_list.append(input_file)
                        if returncode != 0:
                            print(final_result, file=log_file_content)
                            break
                    if returncode == 0:
                        print(final_result)
                        print(final_result, file=log_file_content)
                        check_num_miRNA_signal = 0
                        all_miRNA_list_tmp = open(dir_path + "all_miRNA.list","w+")
                        for i in new_file_list:
                            print(i, file=all_miRNA_list_tmp)
                            with open(i, "r") as f:
                                file_content = f.readlines()
                            if len(file_content) == 1:
                                check_num_miRNA_signal += 1
                        all_miRNA_list_tmp.close()
                        if check_num_miRNA_signal == len(new_file_list):
                            returncode += 1
                            final_result = "No miRNA is found in all Libs."
                            print(final_result, file=log_file_content)
                        if returncode == 0:
                            returncode, final_result = MIRmerger(dir_path, sRNAminer)
                            print(final_result, file=log_file_content)
                        if returncode == 0:
                            print(final_result)
                            returncode, final_result = DetectKnownMIR_2(dir_path, sRNAminer,sRNAminer_database)
                            print(final_result, file=log_file_content)
                            if returncode == 0:
                                print(final_result)
                                returncode, final_result = MIRnamer(dir_path, sRNAminer)
                                print(final_result, file=log_file_content)
                                if returncode == 0:
                                    print(final_result)
                                    returncode, final_result = for_miRNAinfo(dir_path)
                                    print(final_result, file=log_file_content)
                                    if returncode == 0:
                                        print(final_result)
                                        returncode, final_result = MIRinfo2GFF(dir_path,genome_file, sRNAminer)
                                        print(final_result, file=log_file_content)
                print(final_result)
            else:
                file_list = str(input_file_1).split(",")
                new_file_list = []
                for input_file in file_list:
                    returncode, final_result = soft_build(input_file, dir_path)
                if returncode == 0:
                    returncode, final_result = miRNAminerCLI(genome_file, minMirLen, maxMirLen, dir_path, thread,
                                                             sRNAminer)
                    print(final_result, file=log_file_content)
                if returncode == 0:
                    print(final_result)
                    returncode, final_result = DetectKnownMIR(sRNAminer, dir_path, thread, sRNAminer_database)
                    print(final_result, file=log_file_content)
                if returncode == 0:
                    print(final_result)
                    for input_file in file_list:
                        input_file = dir_path + os.path.basename(input_file) + ".tab.known"
                        returncode, final_result = filter(input_file)
                        input_file = dir_path + os.path.basename(input_file) + ".filter"
                        new_file_list.append(input_file)
                        if returncode != 0:
                            print(final_result, file=log_file_content)
                            break
                    if returncode == 0:
                        print(final_result)
                        print(final_result, file=log_file_content)
                        check_num_miRNA_signal = 0
                        all_miRNA_list_tmp = open(dir_path + "all_miRNA.list", "w+")
                        for i in new_file_list:
                            print(i, file=all_miRNA_list_tmp)
                            with open(i, "r") as f:
                                file_content = f.readlines()
                            if len(file_content) == 1:
                                check_num_miRNA_signal += 1
                        all_miRNA_list_tmp.close()
                        if check_num_miRNA_signal == len(new_file_list):
                            returncode += 1
                            final_result = "No miRNA is found in all Libs."
                            print(final_result, file=log_file_content)
                        if returncode == 0:
                            returncode, final_result = MIRmerger(dir_path, sRNAminer)
                            print(final_result, file=log_file_content)
                        if returncode == 0:
                            print(final_result)
                            returncode, final_result = DetectKnownMIR_2(dir_path, sRNAminer, sRNAminer_database)
                            print(final_result, file=log_file_content)
                            if returncode == 0:
                                print(final_result)
                                returncode, final_result = MIRnamer(dir_path, sRNAminer)
                                print(final_result, file=log_file_content)
                                if returncode == 0:
                                    print(final_result)
                                    returncode, final_result = for_miRNAinfo(dir_path)
                                    print(final_result, file=log_file_content)
                                    if returncode == 0:
                                        print(final_result)
                                        returncode, final_result = MIRinfo2GFF(dir_path, genome_file, sRNAminer)
                                        print(final_result, file=log_file_content)
                print(final_result)
        else:
            print(final_result)
        log_file_content.close()
    if returncode == 0:
        print("Finish!")





if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", "-i", required=False, help="use_for_rfam")
    # parser.add_argument("--genome_file", "-g", required=False, help="use_for_rfam")
    args = parser.parse_args()
    input_file_arg = args.input_file
