#!/usr/bin/env python
#==============================================================================
# annotate-juncs.py
#
# Shawn Driscoll
# 20160819
#
# Gene Expression Laboratory, Pfaff
# Salk Institute for Biological Studies
#
# Annotate a junctions file such as those generated by 'bam2juncs' which is 
# kind of like a psi/theta quantification or the juncdb quantification
#==============================================================================

import sys, argparse, re
import subprocess as sp
from os import unlink
from os.path import isfile, expanduser

# from subprocess import Popen
# from random import gauss, random, sample
# from scipy.stats import norm
# import numpy as np
# import numpy.random as npr

# R support
# import rpy2.robjects as robjects
# r = robjects.r


HOME = expanduser("~")

#==============================================================================
# main
#==============================================================================

def main(args):

	# variables
	juncdb = {}
	j5p = {}
	j3p = {}
	tid = ""
	gid = ""
	gname = ""
	ltid = ""
	laln = []
	jid = ""
	jstats = dict(num=0, num_annot=0, num_unannot_anchored=0)

	#
	# Start parsing the GTF
	#

	sys.stderr.write("Parsing {}\n".format(args.gtf))
	try:
		fin = open(args.gtf, "r")
	except:
		sys.stderr.write("Failed to open GTF ({})\n".format(args.gtf))
		return 1

	sys.stderr.write("Parsing GTF for junctions...\n")

	for szl in fin:
		aln = szl.strip().split("\t")

		# skip any non exon lines
		if aln[2] != "exon":
			continue

		# find the transcript id
		rres = re.search("transcript_id \"([^\"]+)\"", aln[8])
		if rres:
			tid = rres.group(1)
		else:
			sys.stderr.write("[main] failed to parse transcript id")
			fin.close()
			return(1)

		# check if current transcript id matches the last one. if so then we have a junction
		# because we have gone from one exon to the next
		if tid == ltid:

			# junction is from laln[4] to aln[3]
			jid = "{}:{}-{}".format(aln[0], str(int(laln[4])+1), str(int(aln[3])-1))

			# build 3p and 5p ids
			jid5p = "{}:{}".format(aln[0], str(int(laln[4])+1))
			jid3p = "{}:{}".format(aln[0], str(int(aln[3])-1))


			# is this junction in the database already?
			if jid not in juncdb:
				juncdb[jid] = dict(ref=aln[0], strand=aln[6], gnames={}, gids={}, tids={})

			# is this 5p end in the db already?
			if jid5p not in j5p:
				j5p[jid5p] = dict(gnames={}, gids={}, tids={}, strand=aln[6], targets=[])

			if jid3p not in j3p:
				j3p[jid3p] = dict(gnames={}, gids={}, tids={}, strand=aln[6], targets=[])

			# update targets for the 3p and 5p ends
			j5p[jid5p]['targets'].append(jid3p)
			j3p[jid3p]['targets'].append(jid5p)

			# add transcript id to the sets
			juncdb[jid]["tids"][tid] = 0
			j5p[jid5p]["tids"][tid] = 0
			j3p[jid3p]["tids"][tid] = 0

			# add gene id to the sets
			rres = re.search("gene_id \"([^\"]+)\"", aln[8])
			gid = ""
			if rres:
				gid = rres.group(1)
				juncdb[jid]['gids'][rres.group(1)] = 0
				j5p[jid5p]["gids"][rres.group(1)] = 0
				j3p[jid3p]["gids"][rres.group(1)] = 0

			# add gene name to the set
			rres = re.search("gene_name \"([^\"]+)\"", aln[8])
			gname = ""
			if rres:
				gname = rres.group(1)
				juncdb[jid]['gnames'][rres.group(1)] = 0
				j5p[jid5p]["gnames"][rres.group(1)] = 0
				j3p[jid3p]["gnames"][rres.group(1)] = 0

		laln = list(aln)
		ltid = tid

	fin.close()

	# last line
	if tid == ltid:

		# junction is from laln[4] to aln[3]
		jid = "{}:{}-{}".format(aln[0], str(int(laln[4])+1), str(int(aln[3])-1))

		# build 3p and 5p ids
		jid5p = "{}:{}".format(aln[0], str(int(laln[4])+1))
		jid3p = "{}:{}".format(aln[0], str(int(aln[3])-1))

		# is this junction in the database already?
		if jid not in juncdb:
			juncdb[jid] = dict(ref=aln[0], strand=aln[6], gnames={}, gids={}, tids={})

		# is this 5p end in the db already?
		if jid5p not in j5p:
			j5p[jid5p] = dict(gnames={}, gids={}, tids={}, strand=aln[6], targets=[])

		if jid3p not in j3p:
			j3p[jid3p] = dict(gnames={}, gids={}, tids={}, strand=aln[6], targets=[])

		# update targets for the 3p and 5p ends
		j5p[jid5p]['targets'].append(jid3p)
		j3p[jid3p]['targets'].append(jid5p)

		# add transcript id to the sets
		juncdb[jid]["tids"][tid] = 0
		j5p[jid5p]["tids"][tid] = 0
		j3p[jid3p]["tids"][tid] = 0

		# add gene id to the sets
		rres = re.search("gene_id \"([^\"]+)\"", aln[8])
		gid = ""
		if rres:
			gid = rres.group(1)
			juncdb[jid]['gids'][rres.group(1)] = 0
			j5p[jid5p]["gids"][rres.group(1)] = 0
			j3p[jid3p]["gids"][rres.group(1)] = 0

		# add gene name to the set
		rres = re.search("gene_name \"([^\"]+)\"", aln[8])
		gname = ""
		if rres:
			gname = rres.group(1)
			juncdb[jid]['gnames'][rres.group(1)] = 0
			j5p[jid5p]["gnames"][rres.group(1)] = 0
			j3p[jid3p]["gnames"][rres.group(1)] = 0

	sys.stderr.write("done.\n")

	#
	# now read the junctions file and annotate the junctions. we need to keep some 
	# statistics as well.
	#

	try:
		fin = open(args.juncs, "r")
	except:
		sys.stderr.write("Failed to open junctions file ({})\n".format(args.juncs))
		return 1

	sys.stderr.write("Parsing junctions file and annotating junctions\n")
	if args.r:
		sys.stderr.write("Ouptut will contain only annotated junctions!\n")

	if args.d > 0:
		lheader = fin.readline().strip().split("\t")
		if lheader[4] == "name":
			lheader[4] = "id"
		lheader += ["flag", "gene_name", "gene_id", "transcript_id", "strand"]
		print "\t".join(lheader)

	for szl in fin:
		aln = szl.strip().split("\t")

		# position is the first part

		# TODO: add in that bit of code from asi-count-....py that figures out the intron
		# ends from the extra overlap dimensions in a typical junctions output from an aligner like STAR

		jid = "{}:{}-{}".format(aln[0], str(aln[1]), str(aln[2]))
		flag = int(0)

		# flags...
		# 0x1 = annotated
		# 0x2 = 5p annotated
		# 0x4 = 3p annotated
		# 0x8 = novel, LR annotated, common genes

		dannot = dict(flag=int(0), gene_id="none", transcript_id="none", gene_name="none", strand="u")
		if jid in juncdb:
			dannot['flag'] |= 0x1
			dannot['gene_id'] = ",".join(sorted(juncdb[jid]['gids'].keys()))
			dannot['gene_name'] = ",".join(sorted(juncdb[jid]['gnames'].keys()))
			dannot['transcript_id'] = ",".join(sorted(juncdb[jid]['tids'].keys()))
			dannot['strand'] = juncdb[jid]['strand']
		else:
			jid5p = "{}:{}".format(aln[0], str(aln[1]))
			jid3p = "{}:{}".format(aln[0], str(aln[2]))

			if jid5p in j5p:
				dannot['flag'] |= 0x2

			if jid3p in j3p:
				dannot['flag'] |= 0x4

			if (jid5p in j5p) and (jid3p in j3p):

				# both ends are annotated - lets see if we can make sense of this
				gid5p = j5p[jid5p]['gids'].keys()
				gn5p = j5p[jid5p]['gnames'].keys()
				gid3p = j3p[jid3p]['gids'].keys()
				gn3p = j3p[jid3p]['gnames'].keys()

				# intersect these
				gid_common = set(gid5p).intersection(gid3p)
				gn_common = set(gn5p).intersection(gn3p)

				# if there is an intersection..
				if len(gid_common) > 0:
					# we have a common gene id
					dannot["gene_id"] = ",".join(sorted(list(gid_common)))
					dannot["flag"] |= 0x8
				else:
					# no common gene ids...
					dannot["gene_id"] = ",".join(sorted(gid5p)) + "|" + ",".join(sorted(gid3p))

				if len(gn_common) > 0:
					# we have a common gene name
					dannot["gene_name"] = ",".join(sorted(list(gn_common)))
				else:
					dannot["gene_name"] = ",".join(sorted(gn5p)) + "|" + ",".join(sorted(gn3p))

			else:

				if jid5p in j5p:
					dannot["gene_id"] = ",".join(sorted(j5p[jid5p]['gids'].keys()))
					dannot["gene_name"] = ",".join(sorted(j5p[jid5p]['gnames'].keys()))

				elif jid3p in j3p:
					dannot["gene_id"] = ",".join(sorted(j3p[jid3p]['gids'].keys()))
					dannot["gene_name"] = ",".join(sorted(j3p[jid3p]['gnames'].keys()))

		lout = aln + ["{:d}".format(dannot['flag']), dannot['gene_name'], dannot['gene_id'], dannot['transcript_id'], dannot['strand']]

		if (not args.r) | (args.r and (dannot['flag'] & 0x1)):
			print "\t".join(lout)

	fin.close()

	return 0


# --
# runcmd
# run a system level command in subprocess. optionally you can return the process.
# if the process isn't returned then the function waits for the process to finish
def runcmd(cmd, returnProcess=False):
	sys.stderr.write("CMD: {}\n".format(cmd))
	p1 = sp.Popen(cmd.split())

	if returnProcess==True:
		return(p1)

	p1.wait()
	return(0)

#==============================================================================
# entry point
#==============================================================================


parser = argparse.ArgumentParser(description="...")
parser.add_argument('gtf', type=str, help="GTF to extract junctions from")
parser.add_argument('juncs', type=str, 
	help="Junctions file. Expects to find chrom, start and end as the first three columns.")
parser.add_argument('-d', action="store_const", default=False, const=True,
	help="Input junctions includes a header row")
parser.add_argument('-f', action="store_const", const=True, default=False, 
	help="Reformat to match JUNCDB output")
parser.add_argument('-r', action="store_const", const=True, default=False, 
	help="Remove unannotated junctions")
parser.add_argument('-o', type=str, action="store", default=None, 
	help="File name for annotation report")

args = parser.parse_args()

if __name__ == "__main__":

	try:
		sys.exit(main(args))
	except KeyboardInterrupt:
		sys.stderr.write("\nkilled it\n")

