import HTSeq
import math
import pickle


#Get transcript lengths
transcripts=HTSeq.FastaReader('Transcripts.fasta')
lengths={}
for t in transcripts:
	lengths[t.name.split('|')[1]]=len(t.seq)


#Make dictionary of splicing variants

vardic={}

for t in transcripts:
	sp=t.name.split('|')
	if sp[0] not in vardic:
		vardic[sp[0]]=[sp[1]]
	else:
		vardic[sp[0]].append(sp[1])

import os
for subdir, dirs, files in os.walk('.'):
	for f in files:
		s=f.find('1.sam')
		if s!=-1:

			genexp={}
			gencount={}
			transabs={}
			fragments=0

			alms=HTSeq.SAM_Reader(f)
			for bundle in HTSeq.pair_SAM_alignments(alms,bundle=True):
				fragments+=1
				genes=[]
				transcripts=[]
	
				for m in bundle:
					first_al, second_al=m
					if first_al.aligned and second_al.aligned:
						if first_al.iv.chrom==second_al.iv.chrom:
							if first_al.inferred_insert_size<600:
								genes.append(first_al.iv.chrom.split('|')[0])
								transcripts.append(first_al.iv.chrom.split('|')[1])
	
				if len(genes)>0:
					ambiguous=0
					for g in genes:
						if g!=genes[0]:
							ambiguous=1
					if ambiguous==0:
						l=0
						n=len(transcripts)
						for t in transcripts:
							l+=lengths[t]
						avlen=float(l)/float(n)
						fpk=1000/float(avlen)
			
						if genes[0] not in genexp:
							genexp[genes[0]]=fpk
							gencount[genes[0]]=1
						else:
							genexp[genes[0]]+=fpk
							gencount[genes[0]]+=1
							
						for tr in vardic[genes[0]]:
							if tr not in transcripts:
								if tr not in transabs:
									transabs[tr]=1
								else:
									transabs[tr]+=1

			gen_fpkm={}

			for g in genexp:
				gen_fpkm[g]=genexp[g]*float(1000000)/float(fragments)
			name=f[:s]+'.pkl'


			with open(name,'wb') as output:
				pickle.dump(gen_fpkm, output)


	
	
					
					
				

