#!/usr/bin/python

import sys
import gzip

if len(sys.argv) < 2:
	print "FPKM calculator.\nUsage:\n"
	print "fpkm.py annotation.gtf countmatrix.tsv\n"
	print "Results goes to standard output\n"
	print "annotation.gtf:      A GTF file containing the gene names and positions"
	print "countmatrix.tsv:     A count matrix created from multiple HTSeq outputs"
	exit(0)

genes = dict()

if sys.argv[1].endswith("gz"):
	gtffile = gzip.open(sys.argv[1])
else:
	gtffile = open(sys.argv[1])

for i in gtffile:
	if i.startswith("#"):
		continue
	fields = i.rstrip().split()
	if fields[2] != "gene":
		continue
	gene_name = fields[13].replace('"', "")[:-1]
	length    = int(fields[4]) - int(fields[3])
	genes[gene_name] = length


table = open(sys.argv[2])
colnum = len(table.readline().rstrip().split())
colsum = [0] * colnum

for i in table:
	fields = i.rstrip().split()
	# skipping rows with unknown gene name
	if fields[0] not in genes:
		sys.stderr.write(fields[0] + " not in annotation\n")
		continue
	for index in range(colnum):
		colsum[index] += int(fields[index + 1])
table.close()

table = open(sys.argv[2])
print table.readline().rstrip()

for i in table:
	fields = i.rstrip().split()
	if fields[0] not in genes:
		continue
	gene_length = genes[fields[0]]
	for index in range(colnum):
		fields[index + 1] = str(1e9 * float(fields[index + 1]) / ( colsum[index] * gene_length))
	print "\t".join(fields)
