import subprocess
import argparse
import traceback
import os
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 build_index_bowtie(input_genome):
    cmd = f'bowtie-build {input_genome} {input_genome}'
    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = sout.decode('gbk')
    else:
        final_result = serr.decode('gbk')
    return returncode, final_result

def build_index_TBtools(input_fasta,sRNAminer):
    output_fasta = input_fasta + ".TBtools.fa"
    cmd = "java -cp " + sRNAminer + " biocjava.bioIO.FastX.FastaIndex.MakeFastaIndex --inFa " + input_fasta + " --outFa " + output_fasta

    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = "Build Index of sRNAminer [Finish]"
    else:
        final_result = serr.decode('gbk')
    return returncode,final_result

def align(index,input_fasta,thread,locate,mis_num,dir_path):
    output_file = dir_path + os.path.basename(input_fasta) + ".bowtie"
    cmd = "bowtie -a -m "+ str(locate) + " --best --strata -p " + str(thread) + " -v " + str(mis_num) + " -f " + index + " " + input_fasta + " " + output_file

    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = sout.decode('gbk') + "\n" + serr.decode('gbk')
    else:
        final_result = serr.decode('gbk')
    return returncode,final_result


def bowtie_format_change(dir_path,thread,sRNAminer):
    cmd = f'ls {dir_path}*.bowtie| parallel -j {thread} "java -cp {sRNAminer} biocjava.sRNA.Tools.sRNAseqAlignmentFormater --inFile {{}} --outFile {{}}.aln"'
    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = "Changing Bowtie format [Finish]"
    else:
        final_result = serr.decode('gbk')
    return returncode,final_result

def sort_bowtie(dir_path,thread,sRNAminer):
    cmd = f'ls {dir_path}*.bowtie.aln|parallel -j {thread} "java -cp {sRNAminer} biocjava.sRNA.Tools.sRnaAlnIndexBuilder --inAln {{}}"'

    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = "Sort Bowtie File [Finish]"
    else:
        final_result = serr.decode('gbk')
    return returncode,final_result

def alignment_pipeline(index,input_file_1,locate,mismatch,thread,dir_path,Batch_mode,sRNAminer):
    print("Start...")
    returncode = 0
    n = 0

    if not dir_path.endswith("/") or dir_path.endswith("\\"):
        dir_path = dir_path + "/"
    final_dir_path = dir_path + "Alignment/"
    while os.path.exists(final_dir_path):
        n += 1
        final_dir_path = dir_path + "Alignment_" + str(n) + "/"
    dir_path = final_dir_path
    returncode, final_result = dir_build(dir_path)
    print(final_result)
    if returncode == 0:
        log_file = dir_path + "alignment.log"
        log_file_content = open(log_file, "w+")
        final_result = ""
        if Batch_mode:
            input_file_content = open(input_file_1, "r")
            file_list = [i.strip() for i in input_file_content.readlines()]
            for input_file in file_list:
                returncode, final_result = align(index, input_file, thread, locate, mismatch,dir_path)  # .bowtie
                print(final_result, file=log_file_content)
                print(final_result)
            if returncode == 0:
                returncode, final_result = bowtie_format_change(dir_path, thread, sRNAminer)  # .aln
                print(final_result, file=log_file_content)
                print(final_result)
            if returncode == 0:
                returncode, final_result = sort_bowtie(dir_path,thread, sRNAminer)  # .aln
                print(final_result, file=log_file_content)
            print(final_result)
            input_file_content.close()
        else:
            file_list = str(input_file_1).split(",")
            for input_file in file_list:
                returncode, final_result = align(index, input_file, thread, locate, mismatch, dir_path)  # .bowtie
                print(final_result, file=log_file_content)
                print(final_result)
            if returncode == 0:
                returncode, final_result = bowtie_format_change(dir_path, thread, sRNAminer)  # .aln
                print(final_result, file=log_file_content)
                print(final_result)
            if returncode == 0:
                returncode, final_result = sort_bowtie(dir_path, thread, sRNAminer)  # .aln
                print(final_result, file=log_file_content)
            print(final_result)
        log_file_content.close()
    if returncode == 0:
        print("Finish!")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", "-i",required=True,help="use_for_rfam")
    parser.add_argument("--genome_fasta", "-g",required=True,help="use_for_rfam")
    parser.add_argument("--thread", "-t",required=True,help="use_for_rfam")
    parser.add_argument("--mis_num", "-m",required=True,help="use_for_rfam")
    parser.add_argument("--dir_path", "-o",required=True,help="use_for_rfam")
    args = parser.parse_args()
    input_file_arg = args.input_file
    genome_fasta_arg = args.genome_fasta
    thread_arg = args.thread
    mis_num_arg = args.mis_num
    dir_path_arg = args.dir_path
    print("Start analyze")
    alignment_pipeline(genome_fasta_arg,input_file_arg,thread_arg,mis_num_arg,dir_path_arg)
    print("Analyze finish!")