import subprocess
import argparse
import os
import pandas as pd
from itertools import islice
import traceback
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 feature_region(input_file,dir_path):
    returncode = 0
    final_result = ""
    try:
        output_file = open(dir_path + "miRNA.feature.region", "w+")
        with open(input_file, 'r') as f: # 这里放的是gff3
            for line in islice(f, 1, None):
                line_list = line.split("\t")
                col_3 = line_list[2]
                if col_3 == "miRNA":
                    line_list[8] = line_list[8].split(";")[2].replace("ID=", "") # sRNAminer version
                    str_line = "\t".join([line_list[8], line_list[0], line_list[3], line_list[4], line_list[6]])
                    print(str_line, file=output_file)
        return returncode, final_result
    except Exception:
        returncode += 1
        final_result = traceback.format_exc()
        return returncode, final_result


def sRNAQuantify(input_file,dir_path,sRNAminer):
    output_file = dir_path + os.path.basename(input_file) + ".miRNA.feature.region.exp"
    cmd = "java -cp " + sRNAminer +" biocjava.sRNA.Miner.sRNAQuantify --inAln " + input_file +\
          " --insRNAregion " + dir_path + "miRNA.feature.region --outsRNAExp " + output_file

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

def Merge_Exp(dir_path,file_list):
    returncode = 0
    final_result = ""
    try:
        output_file = dir_path + "merged.miRNA.exp.xls"
        EXP_DICT = {}
        file_list = [os.path.basename(i) + ".miRNA.feature.region.exp" for i in file_list]
        for root, dirs, files in os.walk(dir_path):
            for file in files:
                if file in file_list:
                    EXP_DICT[file] = {}
                    sample_name = os.path.basename(file)
                    open_file = open(dir_path + file, "r")
                    while True:
                        line = open_file.readline().rstrip()
                        if not line: break
                        line_list = line.split("\t")
                        EXP_DICT[sample_name][line_list[0]] = line_list[5]

        data = pd.DataFrame(EXP_DICT)
        data = data.reset_index()
        data = data.rename(columns={"index": "GeneID"})
        data.to_csv(output_file,sep="\t",index=False)
        final_result = "Calculating Finish!"
        return returncode, final_result
    except Exception:
        returncode += 1
        final_result = traceback.format_exc()
        return returncode, final_result

def miRNA_expression_merger_pipeline(gff_file,input_file_1,dir_path,Batch_mode,sRNAminer):
    print("Start...")
    if not dir_path.endswith("/") or dir_path.endswith("\\"):
        dir_path = dir_path + "/"
    log_file = dir_path + "miRNA_expression.log"
    log_file_content = open(log_file, "w+")
    returncode, final_result = feature_region(gff_file,dir_path)
    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()]
            for input_file in file_list:
                returncode,final_result = sRNAQuantify(input_file,dir_path,sRNAminer)
                print(final_result, file=log_file_content)
                if returncode != 0:
                    break
            if returncode == 0:
                returncode, final_result = Merge_Exp(dir_path,file_list)
                print(final_result, file=log_file_content)
            input_file_content.close()
            print(final_result)

        else:
            file_list = str(input_file_1).split(",")
            for input_file in file_list:
                returncode, final_result = sRNAQuantify(input_file, dir_path, sRNAminer)
                print(final_result, file=log_file_content)
                if returncode != 0:
                    break
            if returncode == 0:
                returncode, final_result = Merge_Exp(dir_path, file_list)
                print(final_result, file=log_file_content)
            print(final_result)
    else:
        print(final_result)
    log_file_content.close()
    if returncode == 0:
        print("Finish!")