import argparse
from Bio.Seq import reverse_complement
from sRNA_abundance import calculate_alignments_read
import traceback
import os
def detInfo(mappingLines):
    lines = mappingLines.split('\n')
    posDetInfo = {}
    for line in lines:
        if line == "": continue
        (readid, seq, chrom, start, end, abund, strand, hits) = line.split('\t')
        start = int(start) + 1
        end = int(end) + 1
        if len(seq) >=20 and len(seq) <= 24:
            if strand == "+":
                if readid not in posDetInfo:
                    posDetInfo[readid] = {}
                    num=int(abund)/int(hits)
                    posDetInfo[readid]["seq"] = seq
                    posDetInfo[readid]["num"] = num
                    posDetInfo[readid]["site_info"] = chrom + ":" + str(start)  + "-" + str(end) + "(" + strand + ")"
            elif strand == "-":
                if readid not in posDetInfo:
                    posDetInfo[readid] = {}
                    num=int(abund)/int(hits)
                    posDetInfo[readid]["seq"] = reverse_complement(seq)
                    posDetInfo[readid]["num"] = num
                    posDetInfo[readid]["site_info"] = chrom + ":" + str(start)  + "-" + str(end) + "(" + strand + ")"

    return (posDetInfo)



def ReadmappingFile(mappingFile):
    FILE = open(mappingFile, "r")
    pos2abund = {}
    while True:  # group together the mapping lines with the same chromosomal position
        line = FILE.readline().rstrip()
        if not line: break
        eles = line.split('\t')
        start_pos = str(int(eles[3]) + 1)
        if eles[2] not in pos2abund:
            pos2abund[eles[2]] = {}
            if start_pos not in pos2abund[eles[2]]:
                pos2abund[eles[2]][start_pos] = line + "\n"
            else:
                pos2abund[eles[2]][start_pos] += line + "\n"
        else:
            if start_pos not in pos2abund[eles[2]]:
                pos2abund[eles[2]][start_pos] = line + "\n"
            else:
                pos2abund[eles[2]][start_pos] += line + "\n"

    FILE.close()
    return pos2abund

def calSRNAabun_each(sRNAloci_file,mappingFile,dir_path):
    pos2abund = ReadmappingFile(mappingFile)
    normNum = calculate_alignments_read.cal_abundance(mappingFile)

    out_seq_file_name = dir_path + os.path.basename(mappingFile) + ".sRNALociSeq.abun.xls"
    out_seq_file = open(out_seq_file_name, "w+")
    print("sRNALoci_ID\tCHR_ID\tStart_Pos\tEnd_Pos\tRead_ID\tSeq\tSeq Length\tAbundance\tRead Pos",file=out_seq_file)
    with open(sRNAloci_file,"r") as f:
        for sRNA_line in f:
            if sRNA_line.startswith("sRNAloci_ID\tChr_ID\tStart_Pos"):
                continue
            else:
                print_Line_dict = {}
                sRNA_line_eles = sRNA_line.split('\t')
                sRNAloci_ID = sRNA_line_eles[0]
                Chr_ID = sRNA_line_eles[1]
                Start_Pos = sRNA_line_eles[2]
                End_Pos = sRNA_line_eles[3]
                for pos in range(int(Start_Pos), int(End_Pos)):
                    if Chr_ID in pos2abund:
                        if str(pos) in pos2abund[Chr_ID]:
                            posInfo = detInfo(pos2abund[Chr_ID][str(pos)])
                            if len(posInfo) != 0:
                                for key in posInfo.keys():
                                    if key not in print_Line_dict.keys():
                                        print_Line_dict[key] = posInfo[key]
                                    elif key in print_Line_dict.keys():
                                        print_Line_dict[key]["site_info"] += "," + posInfo[key]["site_info"]
                                        print_Line_dict[key]["num"] += posInfo[key]["num"]
                if len(print_Line_dict) != 0:
                    for read_id in print_Line_dict:
                        normalize_num = print_Line_dict[read_id]["num"] * 10000000 / int(normNum)
                        print(sRNAloci_ID,Chr_ID,Start_Pos,End_Pos,read_id,print_Line_dict[read_id]["seq"],len(print_Line_dict[read_id]["seq"]),round(normalize_num,2),
                              print_Line_dict[read_id]["site_info"],sep="\t",file=out_seq_file)
                else:
                    print(sRNAloci_ID + "\tNo reads found",file=out_seq_file)
        out_seq_file.close()

def sRNA_extract_merge_pipeline(sRNAloci_file,mappingFile,dir_path,Batch_mode):
    print("Start...")
    try:
        if not dir_path.endswith("/") or dir_path.endswith("\\"):
            dir_path = dir_path + "/"
        if Batch_mode:
            input_file_content = open(mappingFile, "r")
            file_list = [i.strip() for i in input_file_content.readlines()]
            for input_file in file_list:
                calSRNAabun_each(sRNAloci_file,input_file,dir_path)
            input_file_content.close()
        else:
            file_list = str(mappingFile).split(",")
            for input_file in file_list:
                calSRNAabun_each(sRNAloci_file,input_file,dir_path)
        print("Finish")
    except Exception as e:
        print(e)

