#!/usr/bin/python
# -*- coding:utf-8 -*-

import argparse,time,os

def GetTargetFile(targetfile):
	target_dict={}
	#{"EML4-ALK":["E12-A20","E8-A20",...]}
	with open(targetfile,'r') as tf:
		header = tf.readline()
		for line in tf:
			info =line.strip("\n").split("\t")
			gene5=info[0]
			gene3=info[2]
			detail_list=info[-1].split(",")
			detail_list_new=[]
			#ATIC_exon7-ALK_exon20
			for item in detail_list:
				item_info = item.split("-")
				item5_list=item_info[0].split("_")
				new_item5=item5_list[0][0]+item5_list[1].replace("exon","")
				item3_list=item_info[1].split("_")
				new_item3=item3_list[0][0]+item3_list[1].replace("exon","")
				detail_list_new.append(new_item5+"-"+new_item3)
			target_dict[gene5+"-"+gene3]=detail_list_new
	return target_dict

def GetBlackList(blacklist):
	blacklist_dict={}
	with open(blacklist,'r') as bl:
		for line in bl:
			info = line.strip("\n").split("\t")
			gene3=info[3]
			pos=info[4].split(":")[1]
			seq=info[5]
			if gene3 not in blacklist_dict.keys():
				blacklist_dict[gene3]=[]
			blacklist_dict[gene3].append([pos,seq])
	return blacklist_dict

def CheckMismatch(seq1_tmp,seq2_tmp,mismatch=2):

	seq1_tmp_len=len(seq1_tmp)
	count =0
	for index in range(seq1_tmp_len):
		if seq1_tmp[index] != seq2_tmp[index]:
			count +=1
		if count > mismatch:
			return False
	return True
	

def ISOverlap(seq1,seq2):

	#seq2---seq1
	for index in range(5):
		seq1_tmp = seq1[index:]
		seq1_len = len(seq1_tmp)
		seq2_tmp = seq2[:seq1_len]
		if CheckMismatch(seq1_tmp,seq2_tmp):
			return True
	
	#seq1--seq2
	for index in range(5):
		seq2_tmp = seq2[index:]
		seq2_len = len(seq2_tmp)
		seq1_tmp = seq1[:seq2_len]
		if CheckMismatch(seq1_tmp,seq2_tmp):
			return True
		
	return False

BASE_DICT={"A":"T","G":"C","T":"A","C":"G","N":"N"}

def CheckOverlap(seq_temp,item_seq):
	seq_temp_rev=[BASE_DICT[base] for base in seq_temp][::-1]
	temp_len=len(seq_temp_rev)
	item_len=len(item_seq)

	seq_temp_rev_f=seq_temp_rev[:int(temp_len/2)]
	seq_temp_rev_l=seq_temp_rev[-1*int(temp_len/2):]

	seq_temp_f=seq_temp[:int(temp_len/2)]
	seq_temp_l=seq_temp[-1*int(temp_len/2):]

	item_seq_f=item_seq[:int(item_len/2)]
	item_seq_l=item_seq[-1*int(item_len/2):]

	if ISOverlap(item_seq_f,seq_temp_rev_f) or ISOverlap(item_seq_l,seq_temp_rev_l):
		return True
	
	if ISOverlap(item_seq_f,seq_temp_f) or ISOverlap(item_seq_l,seq_temp_l):
		return True

	return False

def FilterBlacklist(blacklist_dict,FusionID,breakpoint,seq_temp):
	
	gene3=FusionID.replace("(+)","").replace("(-)","").split("-")[1]
	pos = int(breakpoint.split("-")[1].split(":")[1])
	if gene3 in blacklist_dict.keys():
		for item in blacklist_dict[gene3]:
			item_pos=int(item[0])
			item_seq=item[1]
			if abs(pos-item_pos)<5 and CheckOverlap(seq_temp,item_seq):
				return False
	return True

def FilterResults(infile,outf,target_dict,blacklist_dict,stat='d'):
	with open(infile,'r') as inputF:
		header = inputF.readline().strip("\n").split("\t")
		FusionID_index=header.index("FusionID")
		Exon_index=header.index("Exon")
		Transcript_index=header.index("Transcript")
		BreakPoints_index=header.index("BreakPoint")
		DoubleNum_index=header.index("DoubleNum")
		SingleNum_index=header.index("SingleNum")
		UMIkinds_index=header.index("UMIkind")
		Overlap_index=header.index("Overlap")
		FusionSeq_index=header.index("FusionSeq")
		PrimerPairs_index=header.index("PrimerPair")
		if stat == 'd':
			new_header=[header[index] for index in [FusionID_index,Exon_index,Transcript_index,BreakPoints_index,DoubleNum_index,SingleNum_index,UMIkinds_index,PrimerPairs_index]]
			new_header.append("TotalNum")
			outf.write("\t".join(new_header)+"\n")
		for line in inputF:
			info = line.strip("\n").split("\t")
			#是否在target 区域
			FusionID=info[0].replace("(+)","").replace("(-)","")
			ExonID=info[Exon_index]
			#in_target_filter = True if FusionID in target_dict.keys() and ExonID in target_dict[FusionID] else False
			in_target_filter = True if FusionID in target_dict.keys() else False

			#是否是背景库中点
			blackList_select = FilterBlacklist(blacklist_dict,info[FusionID_index],info[BreakPoints_index],info[FusionSeq_index])

			#single reads>25保留。double reads>=2就保留
			DoubleNum=int(info[DoubleNum_index])
			SingleNum=int(info[SingleNum_index])
			UMIkinds=int(info[UMIkinds_index])
			TotalNum=DoubleNum+SingleNum
			#num_select = True if ((DoubleNum >=2 and UMIkinds>=2)or (SingleNum>=25 and UMIkinds>=25)) else False
			num_select = True if DoubleNum >=2 or UMIkinds >=2 else False
			
			#Overlap
			Overlap_select=True if int(info[Overlap_index]) <=5 else False

			if in_target_filter and blackList_select and num_select and Overlap_select:
				new_info=[info[index] for index in [FusionID_index,Exon_index,Transcript_index,BreakPoints_index,DoubleNum_index,SingleNum_index,UMIkinds_index,PrimerPairs_index]]
				new_info.append(str(TotalNum))
				outf.write("\t".join(new_info)+"\n")


def main(indir,blacklist,targetfile):

	t1=time.time()
	target_dict=GetTargetFile(targetfile)
	blacklist_dict=GetBlackList(blacklist)

	outfile=indir+"/fusion_results.txt"
	outf=open(outfile,'w')

	double_results_file=indir+"/fusion_stat_final_anno.txt"
	
	if os.path.exists(double_results_file):
		FilterResults(double_results_file,outf,target_dict,blacklist_dict)

	outf.close()
	t2=time.time()
	print("Times: "+str(int(t2-t1))+"s")
	print("BreakPoints Filter Done!")


if __name__ == '__main__':
	parser = argparse.ArgumentParser(description='Filter FusionGene results,by target parterner gene,blacklist file,fusion reads num')
	parser.add_argument('-i', required=True, type=str, help="indir")
	parser.add_argument('-b', required=True, type=str, help="the blacklist file")
	parser.add_argument('-t', required=True, type=str, help="target parterner gene file")
	args = parser.parse_args()
	main(args.i,args.b,args.t)
