import sys
import argparse
import logging
import time

def pase_info(seq):
    info = {'SVLEN': 0, 'END': 0, "SVTYPE": '', "RE": 0, "CHR2": ''}
    for i in seq.split(';'):
        if i.split('=')[0] in ["SVLEN", "END", "RE"]:
            try:
                info[i.split('=')[0]] = abs(int(float(i.split('=')[1])))
            except:
                pass
        if i.split('=')[0] in ["CHR2"]:
            info[i.split('=')[0]] = i.split('=')[1]
        if i.split('=')[0] in ["SVTYPE"]:
            info[i.split('=')[0]] = i.split('=')[1][0:3]

    return info

def phase_GT(seq):
    i = seq.split(':')
    if i[0] in ['0/1', '1/0']:
        return 'het'
    elif i[0] == '1/1':
        return 'hom'
    else:
        return 'unknown'

def load_callset(path):
    callset = dict()
    file = open(path, 'r')
    for line in file:
        seq = line.strip('\n').split('\t')
        if seq[0][0] == '#':
            continue

        chr = seq[0]
        pos = int(seq[1])
        info = pase_info(seq[7])
        if info['SVTYPE'] == "TRA":
            info['SVTYPE'] = "BND"

        if info['SVTYPE'] == "BND":
            if seq[4][0] == ']':
                form = ']]N'
                chr2 = seq[4].split(':')[0][1:]
                pos2 = int(seq[4].split(':')[1][:-2])
            elif seq[4][0] == '[':
                form = '[[N'
                chr2 = seq[4].split(':')[0][1:]
                pos2 = int(seq[4].split(':')[1][:-2])
            else:
                if seq[4][1] == ']':
                    form = 'N]]'
                    chr2 = seq[4].split(':')[0][2:]
                    pos2 = int(seq[4].split(':')[1][:-1])
                else:
                    form = 'N[['
                    chr2 = seq[4].split(':')[0][2:]
                    pos2 = int(seq[4].split(':')[1][:-1])
            if info['SVTYPE'] not in callset:
                callset[info['SVTYPE']] = list()
            if info['END'] == 0:
                info['CHR2'] = chr2
                info['END'] = pos2
            try:
                if int(chr) <= int(info['CHR2']):
                    if form == 'N[[':
                        form = ']]N'
                    if form == ']]N':
                        form = 'N[['
                    callset[info['SVTYPE']].append([chr, pos, info['CHR2'], info['END'], form, phase_GT(seq[9]), 0])
                else:
                    callset[info['SVTYPE']].append([info['CHR2'], info['END'], chr, pos, form, phase_GT(seq[9]), 0])
            except:
                callset[info['SVTYPE']].append([chr, pos, info['CHR2'], info['END'], form, phase_GT(seq[9]), 0])

    file.close()
    return callset

def eval(call, ans, offect):
    tpcall = 0
    for i in call["BND"]:
        for j in ans["BND"]:
            if i[0] == j[0] and i[2] == j[2] and abs(i[1]-j[1]) <= offect and abs(i[3]-j[3]) <= offect:
                tpcall += 1
                break
    fp = len(call["BND"]) - tpcall
    return fp, len(call["BND"])

def main_ctrl(args):
    # load ground truth set
    base_list = load_callset(args.base)
    comp_list = load_callset(args.comp)
    
    fp, total = eval(comp_list, base_list, args.offect)
    logging.info('False positive in BND: %d'%(fp))
    logging.info('Total amount of BND: %d'%(total))

def main(argv):
	args = parseArgs(argv)
	setupLogging(False)
	# print args
	starttime = time.time()
	main_ctrl(args)
	logging.info("Finished in %0.2f seconds."%(time.time() - starttime))

USAGE="""\
	Evaluate SV callset generated by simulations.
	Author: Shuqi Cao
	Email: sqcao@stu.hit.edu.cn
"""

def parseArgs(argv):
	parser = argparse.ArgumentParser(prog="evaluation on BND", description=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter)
	parser.add_argument("base", type=str, help="Ground truth of BNDs.")
	parser.add_argument("comp", type=str, help="BND callsets to be benched.")
	parser.add_argument('-o', '--offect', help = "Offect of translocation overlaping.[%(default)s]", default = 1000, type = int)
	args = parser.parse_args(argv)
	return args

def setupLogging(debug=False):
	logLevel = logging.DEBUG if debug else logging.INFO
	logFormat = "%(asctime)s [%(levelname)s] %(message)s"
	logging.basicConfig( stream=sys.stderr, level=logLevel, format=logFormat )
	logging.info("Running %s" % " ".join(sys.argv))

if __name__ == '__main__':
	main(sys.argv[1:])