intron_gtf

####get introns and intron/exon joint coordinates form a gtf file
####Input GTF file
####13.03.2017
####Cheked for Coordinate pairs that the latter coordinate is always larger 

#!/usr/bin/python
import os
import sys
import argparse
from itertools import islice, combinations

#### passing arguments

parser=argparse.ArgumentParser()
parser.add_argument('input', help='Input GTF file')
parser.add_argument('-output', help='output GTF file')
parser.add_argument('-junction',  action='store_true', help='Will the coordinates of intron/exon junction sites be saved as another file')
args=parser.parse_args()

if not os.path.isfile(args.input): 
	sys.exit('ERROR: File '+args.input+' seems not to be exist')

#### define functions


def overlap(coor1, coor2):
####For sorted coordinates
	if coor1[0]<=coor2[0]:
		if coor1[1]>=coor2[0]:
			return True
		else:
			return False
	elif coor1[0]<=coor2[1]:
		return True
	else: 
		return False

def merge_olp(num_list):
	out=[];start=num_list[0][0];end=num_list[0][1]
	for i in range(len(num_list)-1):
		if end < num_list[i+1][0]:
			out.append([start,end]) 
			start, end = num_list[i+1]
		else:
			end=max(end,num_list[i+1][1])
	else:
		out.append([start, end])
	return out	

def get_intron(gene, exons):
	intron=[]
	if gene[0]<exons[0][0]:
		intron.append([gene[0],exons[0][0]-1])
	for i in range(len(exons)-1):
		if exons[i+1][0]-1>exons[i][1]:
			intron.append([exons[i][1]+1, exons[i+1][0]-1])
	else:
		if gene[1]-1>exons[-1][1]:
			intron.append([exons[-1][1]+1, gene[1]])
	return intron


def get_junctions(gene, introns):
	junctions=[item for sublist in introns for item in sublist]
	if not introns==[]:
		if junctions[0]==gene[0]:
			junctions.pop(0)
		if junctions[-1]==gene[1]:
			junctions.pop()
		return(junctions)
	else:
		return([])
	
def remove_dups(num_list):
####remove duplicates for a list contains lists with same length
	num_list.sort()
	index=[]
	for i in range(len(num_list)-1):
		if num_list[i]==num_list[i+1]:
	       		index.append(i) 
	return [v for i, v in enumerate(num_list) if i not in index]

	
def Main(infile):
	template=[]
	geneID=[]
	gene=[]
	exons=[]
	number_introns=0
	inf=open(infile, 'r')
	if args.junction:
		out_junction=open(os.path.dirname(os.path.abspath(args.input))+'/junction_'+os.path.basename(args.input), 'w')

	if not args.output:
		filename=os.path.dirname(os.path.abspath(args.input))+'/intron_'+os.path.basename(args.input)
		ouf=open(filename, 'w')
	else:
		ouf=open(args.output, 'w')

	
	if not args.junction:
		while True:
			code=list(islice(inf,1))
			if code == []:
				print ('GTF file reading completed')
				ouf.close()
				inf.close()
				break
			code=code[0].split('\t')
			if geneID == []:
			###Initiate
				if code[2]=='gene':
					geneID=code[8].split()[1]
					gene=[int(x) for x in code[3:5]]
				if code[2]=='exon':
					geneID=code[8].split()[1]
					exons.append([int(x) for x in code[3:5]])
			else:
			###if there comes a new gene, then settle the previous gene & start a new
				if code[8].split()[1] != geneID:
					geneID=[]
					###if a gene has only 1 exon it won't have intron 
					if len(exons)>1:
						exons.sort()
						exons=merge_olp(exons)
					###Calculate introns
						introns=get_intron(gene, exons)
					###Write intron and everything into a new file
						for i in range(len(introns)):
							record=template
							record[2]='intron'
							record[3:5]=[str(i) for i in introns[i]]
							ouf.write('\t'.join(record))
		
					if code[2]=='gene':
						gene=[int(x) for x in code[3:5]]
						exons=[]
					elif code[2]=='exon':
						exons=[int(x) for x in code[3:5]]
						
		
		
				else:
					template=code[:]
					if code[2]=='gene':
						gene=[int(x) for x in code[3:5]]
					elif code[2]=='exon':
						exons.append([int(x) for x in code[3:5]])
		
			ouf.write('\t'.join(code))
	else:

		out_junctions=open(os.path.dirname(os.path.abspath(args.input))+'/junction_'+os.path.basename(args.input), 'w')
		while True:
			code=list(islice(inf,1))
			if code == []:
				print ('GTF file reading completed')
				ouf.close()
				inf.close()
				break
			code=code[0].split('\t')
			if geneID == []:
			###Initiate
				if code[2]=='gene':
					geneID=code[8].split()[1]
					gene=[int(x) for x in code[3:5]]
				if code[2]=='exon':
					geneID=code[8].split()[1]
					exons.append([int(x) for x in code[3:5]])
			else:
			###if there comes a new gene, then settle the previous gene & start a new
				if code[8].split()[1] != geneID:
					geneID=[]
					###if a gene has only 1 exon it won't have intron 
					if len(exons)>1:
						exons.sort()
						exons=merge_olp(exons)
					###Calculate introns
						introns=get_intron(gene, exons)
						junctions=get_junctions(gene, introns)
					###Write intron and everything into a new file
						for i in range(len(introns)):
							record=template
							record[2]='intron'
							record[3:5]=[str(i) for i in introns[i]]
							ouf.write('\t'.join(record))
					###Write junctions
						genename=template[8].split('"')[1]
						for i in junctions:
							out_junctions.write('%s\t%s\t%s\t%d\n' %(genename,template[0], template[6], i))
		
					if code[2]=='gene':
						gene=[int(x) for x in code[3:5]]
						exons=[]
					elif code[2]=='exon':
						exons=[int(x) for x in code[3:5]]
						
		
		
				else:
					template=code[:]
					if code[2]=='gene':
						gene=[int(x) for x in code[3:5]]
					elif code[2]=='exon':
						exons.append([int(x) for x in code[3:5]])
		
			ouf.write('\t'.join(code))
	
if __name__=='__main__':
	Main(args.input)