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 hcsiRNAMinerCLI(minRegionGap, minRepeatRatio,dir_path,thread,sRNAminer):
    cmd = f'ls {dir_path}*.hcsiRNA.bowtie.aln.sorted| parallel -j {thread} "java -cp {sRNAminer} ' \
          f'biocjava.sRNA.Miner.hcsiRNA.hcsiRNAMinerCLI ' \
          f'--inSortedAln {{}} --outRegionInfo {{}}.hcsiRNA --minRegionGap {minRegionGap} --minRepeatRatio {minRepeatRatio}"'
    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = "hcsiRNAminer Step 1: Mining from Lib [Finish]"
    else:
        final_result = serr.decode('gbk')
    return returncode, final_result

def HcsiRNAmerger_1(dir_path,file_list):
    error_signal = 0
    error_str = ""
    file_list = [os.path.basename(i)+".hcsiRNA.bowtie.aln.sorted.hcsiRNA" for i in file_list]
    try:
        output_file = open(dir_path + "hcsiRNA.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()

        return error_signal,error_str
    except Exception:
        error_signal += 1
        error_str = traceback.format_exc()
        return error_signal,error_str
def HcsiRNAmerger_2(dir_path,sRNAminer):
    input_hcsiRNA_list = dir_path + "hcsiRNA.list"
    output_merge_list = dir_path + "merged.hcsiRNA.info"
    cmd = "java -cp " + sRNAminer +" biocjava.sRNA.Miner.hcsiRNA.HcsiRNAmerger --inFileList " + input_hcsiRNA_list + \
          " --outFile " + output_merge_list

    returncode, sout, serr = runCmd(cmd)
    if returncode == 0:
        final_result = "hcsiRNAminer Step 2: Merging hc-siRNA file [Finish]"
    else:
        final_result = serr.decode('gbk')
    return returncode, final_result

def hcsiRNAToGff(input_file):
    returncode = 0
    final_result = ""
    try:
        dir_path = os.path.dirname(input_file) + "/"
        output_file = dir_path + "merged.hcsiRNA.gff3"
        input_file_content = open(input_file,"r")
        output_file_content = open(output_file,"w+")
        print("##gff-version 3",file=output_file_content)
        first_line_signal = 1
        n = 1
        while True:
            line = input_file_content.readline().rstrip()
            if not line:break
            if first_line_signal == 1:
                first_line_signal += 1
                continue
            else:
                (Repeat_ID, Chr_ID, Start_Pos, End_Pos, maxRepeatSpeciesRatio, SpeciesNum, LibNum) = line.split("\t")
                print(Chr_ID, "hcsiRNAminer", "hc-siRNA", Start_Pos, End_Pos, ".", ".", ".","ID=hcsiRNA" + str(n) +
                      ";maxRepeatSpeciesRatio=" + str(maxRepeatSpeciesRatio) + ";SpeciesNum=" + str(SpeciesNum) +
                      ";LibNum=" + str(LibNum),sep="\t",file=output_file_content)
                n += 1
        input_file_content.close()
        output_file_content.close()
        final_result = "hcsiRNAminer Step 4: Establishing a GFF file [Finish]"
        return returncode, final_result
    except Exception:
        returncode += 1
        final_result = traceback.format_exc()
        return returncode, final_result




def filter_24_PHAS(input_hc_file,input_24_PHAS_file):
    returncode = 0
    final_result = ""
    try:
        input_hcsiRNA = pd.read_table(input_hc_file)
        PHAS_24_PHAS_file = open(input_24_PHAS_file,"r")
        dir_path = os.path.dirname(input_hc_file) + "/"
        output_file = dir_path + "final.hcsiRNA.loci.xls"
        first_line_signal = 1
        remove_list = []
        while True:
            line = PHAS_24_PHAS_file.readline().rstrip()
            if not line:break
            if first_line_signal == 1:
                first_line_signal += 1
                continue
            else:
                CHR = line.split("\t")[1]
                STRAT_POS = int(line.split("\t")[2])
                END_POS =int(line.split("\t")[3])
                tmp_pd_1 = input_hcsiRNA[(input_hcsiRNA["Chr_ID"] == CHR) & (input_hcsiRNA["Start_Pos"] >= STRAT_POS) &
                                       (input_hcsiRNA["Start_Pos"] < END_POS) & (input_hcsiRNA["End_Pos"] >= END_POS)]
                tmp_pd_2 = input_hcsiRNA[(input_hcsiRNA["Chr_ID"] == CHR) & (input_hcsiRNA["Start_Pos"] <= STRAT_POS) &
                                         (input_hcsiRNA["End_Pos"] > STRAT_POS) & (input_hcsiRNA["End_Pos"] <= END_POS)]
                tmp_pd_3 = input_hcsiRNA[(input_hcsiRNA["Chr_ID"] == CHR) & (input_hcsiRNA["Start_Pos"] >= STRAT_POS) &
                                         (input_hcsiRNA["End_Pos"] <= END_POS)]
                tmp_pd_4 = input_hcsiRNA[(input_hcsiRNA["Chr_ID"] == CHR) & (input_hcsiRNA["Start_Pos"] < STRAT_POS) &
                                         (input_hcsiRNA["End_Pos"] > END_POS)]

                remove_list.extend(list(tmp_pd_1["Repeat_ID"]))
                remove_list.extend(list(tmp_pd_2["Repeat_ID"]))
                remove_list.extend(list(tmp_pd_3["Repeat_ID"]))
                remove_list.extend(list(tmp_pd_4["Repeat_ID"]))

        remove_list = list(set(remove_list))
        input_hcsiRNA = input_hcsiRNA.set_index('Repeat_ID')
        input_hcsiRNA = input_hcsiRNA.drop(remove_list).reset_index()
        input_hcsiRNA.to_csv(output_file, sep="\t",index=None)
        final_result = "hcsiRNAminer Step 3: Filter [Finish]"
        return returncode, final_result
    except Exception:
        returncode += 1
        final_result = traceback.format_exc()
        return returncode, final_result

def align(genome_fasta,input_fasta,thread,mis_num,dir_path):
    returncode = 0
    if os.path.exists(input_fasta):
        output_file = dir_path + os.path.basename(input_fasta) + ".hcsiRNA.bowtie"
        cmd = "bowtie -a --best --strata -p " + thread + " -v " + mis_num + " -f " + genome_fasta + " " + 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

    else:
        returncode += 1
        final_result = f'File not Found:{input_fasta}'
        return returncode, final_result


def bowtie_format_change(dir_path,thread,sRNAminer):
    cmd = f'ls {dir_path}*hcsiRNA.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}*hcsiRNA.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 hc_siRNA_identification_merger_pipeline(genome_file,index,input_file,mismatch,thread,minRegionGap,minRepeatRatio,PHAS24_file,Batch_mode,dir_path,sRNAminer):
    print("Start...")
    n = 0

    if not dir_path.endswith("/") or dir_path.endswith("\\"):
        dir_path = dir_path + "/"
    final_dir_path = dir_path + "hcsiRNA_identification/"
    while os.path.exists(final_dir_path):
        n += 1
        final_dir_path = dir_path + "hcsiRNA_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 + "hc-siRNA_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,"r")
                file_list = [i.strip() for i in input_file_content.readlines()]
                new_file_list = []
                input_file_content.close()
                for input_file in file_list:
                    returncode, final_result = align(index,input_file,thread,mismatch,dir_path) #.bowtie
                    print(final_result, file=log_file_content)
                    if returncode == 0:
                        print(final_result)
                if returncode == 0:
                    returncode, final_result = bowtie_format_change(dir_path, thread, sRNAminer)  # .aln
                    print(final_result, file=log_file_content)
                if returncode == 0:
                    print(final_result)
                    returncode, final_result = sort_bowtie(dir_path, thread, sRNAminer)  # .aln
                    print(final_result, file=log_file_content)
                if returncode == 0:
                    print(final_result)
                    print("Preparing for Bowtie Files(sRNAminer Format) [Finish]")
                for input_file in file_list:
                    input_file = dir_path + os.path.basename(input_file) + ".hcsiRNA.bowtie.aln.sorted"
                    new_file_list.append(
                        dir_path + os.path.basename(input_file) + ".hcsiRNA")
                if returncode == 0:
                    returncode, final_result = hcsiRNAMinerCLI(minRegionGap,minRepeatRatio, dir_path, thread,sRNAminer)  # .hcsiRNA
                    print(final_result, file=log_file_content)

                if returncode == 0:
                    print(final_result)
                    check_num_hcsiRNA_signal = 0
                    all_hcsiRNA_list_tmp = open(dir_path + "hcsiRNA.list", "w+")
                    for i in new_file_list:
                        print(i, file=all_hcsiRNA_list_tmp)
                        with open(i, "r") as f:
                            file_content = f.readlines()
                        if len(file_content) == 0:
                            check_num_hcsiRNA_signal += 1
                    all_hcsiRNA_list_tmp.close()
                    if check_num_hcsiRNA_signal == len(new_file_list):
                        returncode += 1
                        final_result = "No hc-siRNA is found in all Libs."
                        print(final_result, file=log_file_content)
                    if returncode == 0:
                        returncode, final_result = HcsiRNAmerger_2(dir_path, sRNAminer)

                        print(final_result, file=log_file_content)
                    if returncode == 0:
                        print(final_result)
                        input_hc_file = dir_path + "merged.hcsiRNA.info"
                        returncode, final_result = filter_24_PHAS(input_hc_file,PHAS24_file)
                        print(final_result, file=log_file_content)
                        if returncode == 0:
                            print(final_result)
                            input_for_gff3 = dir_path + "final.hcsiRNA.loci.xls"
                            returncode, final_result = hcsiRNAToGff(input_for_gff3)
                            print(final_result, file=log_file_content)

                print(final_result)
            else:
                file_list = str(input_file).split(",")
                new_file_list = []
                for input_file in file_list:
                    returncode, final_result = align(index, input_file, thread, mismatch, dir_path)  # .bowtie
                    print(final_result, file=log_file_content)
                    if returncode == 0:
                        print(final_result)
                if returncode == 0:
                    returncode, final_result = bowtie_format_change(dir_path, thread, sRNAminer)  # .aln
                    print(final_result, file=log_file_content)
                if returncode == 0:
                    print(final_result)
                    returncode, final_result = sort_bowtie(dir_path, thread, sRNAminer)  # .aln
                    print(final_result, file=log_file_content)
                if returncode == 0:
                    print(final_result)
                    print("Preparing for Bowtie Files(sRNAminer Format) [Finish]")
                for input_file in file_list:
                    input_file = dir_path + os.path.basename(input_file) + ".hcsiRNA.bowtie.aln.sorted"
                    new_file_list.append(
                        dir_path + os.path.basename(input_file) + ".hcsiRNA")
                if returncode == 0:
                    returncode, final_result = hcsiRNAMinerCLI(minRegionGap, minRepeatRatio, dir_path, thread,
                                                               sRNAminer)  # .hcsiRNA
                    print(final_result, file=log_file_content)


                if returncode == 0:
                    print(final_result)
                    check_num_hcsiRNA_signal = 0
                    all_hcsiRNA_list_tmp = open(dir_path + "hcsiRNA.list", "w+")
                    for i in new_file_list:
                        print(i, file=all_hcsiRNA_list_tmp)
                        with open(i, "r") as f:
                            file_content = f.readlines()
                        if len(file_content) == 0:
                            check_num_hcsiRNA_signal += 1
                    all_hcsiRNA_list_tmp.close()
                    if check_num_hcsiRNA_signal == len(new_file_list):
                        returncode += 1
                        final_result = "No hc-siRNA is found in all Libs."
                        print(final_result, file=log_file_content)
                    if returncode == 0:
                        returncode, final_result = HcsiRNAmerger_2(dir_path, sRNAminer)

                        print(final_result, file=log_file_content)
                    if returncode == 0:
                        print(final_result)
                        input_hc_file = dir_path + "merged.hcsiRNA.info"
                        returncode, final_result = filter_24_PHAS(input_hc_file, PHAS24_file)

                        print(final_result, file=log_file_content)
                        if returncode == 0:
                            print(final_result)
                            input_for_gff3 = dir_path + "final.hcsiRNA.loci.xls"
                            returncode, final_result = hcsiRNAToGff(input_for_gff3)
                            print(final_result, file=log_file_content)
                print(final_result)
        log_file_content.close()
    if returncode == 0:
        print("Finish!")