##################################
#                                #
# Last modified 10/15/2012       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import numpy
from sets import Set

try:
	import psyco
	psyco.full()
except:
	pass

def run():

    if len(sys.argv) < 5:
        print 'usage: python %s candidate_lincRNA.gtf annotation_gtf ORF PhyloCSF PFAM_search_output outfilename' % sys.argv[0]
        sys.exit(1)

    gtf = sys.argv[1]
    annotation = sys.argv[2]
    ORF = sys.argv[3]
    PhyloCSF = sys.argv[4]
    PFAM = sys.argv[5]
    outputfileprefix = sys.argv[6]

    TranscriptDict={}

    print gtf
        
    candidatelncRNACoverageDict={}

    listoflines = open(gtf)
    for line in listoflines:
        if line.startswith('#'):
            continue
        fields=line.strip().split('\t')
        transcriptID=fields[8].split('transcript_id "')[1].split('";')[0]
        if TranscriptDict.has_key(transcriptID):
            pass
        else:
            TranscriptDict[transcriptID]={}
            TranscriptDict[transcriptID]['AnnotationOverlap'] = False
        if fields[2] == 'exon':
            chr = fields[0]
            if candidatelncRNACoverageDict.has_key(chr):
                pass
            else:
                candidatelncRNACoverageDict[chr]={}
            left = int(fields[3])
            right = int(fields[4])
            for i in range(left,right):
                candidatelncRNACoverageDict[chr][i]=0

    print annotation

    listoflines = open(annotation)
    for line in listoflines:
        if line.startswith('#'):
            continue
        fields=line.strip().split('\t')
        if fields[2] == 'exon':
            chr = fields[0]
            if candidatelncRNACoverageDict.has_key(chr):
                left = int(fields[3])
                right = int(fields[4])
                for i in range(left,right):
                    if candidatelncRNACoverageDict[chr].has_key(i):
                        candidatelncRNACoverageDict[chr][i]=1

    print gtf

    listoflines = open(gtf)
    for line in listoflines:
        if line.startswith('#'):
            continue
        fields=line.strip().split('\t')
        transcriptID=fields[8].split('transcript_id "')[1].split('";')[0]
        if fields[2] == 'exon':
            chr = fields[0]
            left = int(fields[3])
            right = int(fields[4])
            for i in range(left,right):
                if candidatelncRNACoverageDict[chr][i]==1:
                    TranscriptDict[transcriptID]['AnnotationOverlap']=True
                    break

    print PhyloCSF
    
    listoflines = open(PhyloCSF)
    for line in listoflines:
        if line.startswith('#'):
            continue
        fields=line.strip().split('\t')
        transcriptID = fields[0]
        if TranscriptDict.has_key(transcriptID):
            pass
        else:
            print transcriptID, 'not found in GTF file'
            continue
        if fields[1] == 'exception':
            continue
        if fields[1] == 'abort':
            TranscriptDict[transcriptID]['PhyloCSF'] = 'abort'
        elif fields[1] == 'failure':
            TranscriptDict[transcriptID]['PhyloCSF'] = 'Failure'
        else:
            TranscriptDict[transcriptID]['PhyloCSF'] = float(fields[2])
        print TranscriptDict[transcriptID]['PhyloCSF']

    print PFAM

    listoflines = open(PFAM)
    for line in listoflines:
        if line.startswith('#'):
            continue
        fields=line.strip().split('\t')
        transcriptID = fields[0].split(':')[0]
        if TranscriptDict.has_key(transcriptID):
            pass
        else:
            continue
        if TranscriptDict.has_key(transcriptID):
            TranscriptDict[transcriptID]['PFAM']=True
    
    print ORF

    listoflines = open(ORF)
    TranscriptIDFieldid = 2
    ORFLengthFieldID = 10
    for line in listoflines:
        fields=line.strip().split('\t')
        if line.startswith('#'):
            TranscriptIDFieldid = fields.index('TranscriptID')
            ORFLengthFieldID = fields.index('protein_length')
            continue
        transcriptID = fields[TranscriptIDFieldid]
        if TranscriptDict.has_key(transcriptID):
            pass
        else:
            continue
        ORFlength = fields[ORFLengthFieldID]
        TranscriptDict[transcriptID]['ORFLength'] = ORFlength

    outfile = open(outputfileprefix + '.list', 'w')
    outfile.write('#TranscriptID\tPhyloCSF\tPFAM\tAnnotationOverlap\tORFLength\n')
    PhyloCSFRejected=0
    PFAMRejected=0
    AnnotationRejected=0
    lincRNADict={}
    for transcriptID in TranscriptDict.keys():
        if 'PhyloCSF' in TranscriptDict[transcriptID].keys():
            PhyloCSF = TranscriptDict[transcriptID]['PhyloCSF']
        else:
            PhyloCSF = 'N/A'
        if 'PFAM' in TranscriptDict[transcriptID].keys():
            PFAM = 'Yes'
        else:
            PFAM = 'No'
        if TranscriptDict[transcriptID]['AnnotationOverlap']:
            AnnotationOverlap = 'Yes'
        else:
            AnnotationOverlap = 'No'
        if TranscriptDict[transcriptID].has_key('ORFLength'):
            ORFLength = TranscriptDict[transcriptID]['ORFLength']
        else:
            ORFLength = 'N/A'
        outline = transcriptID + '\t' + str(PhyloCSF) + '\t' + PFAM + '\t' + AnnotationOverlap + '\t' + ORFLength
        outfile.write(outline + '\n')

    outfile.close()

run()