##################################
#                                #
# Last modified 09/11/2011       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
import math
from sets import Set

def run():

    if len(sys.argv) < 2:
        print 'usage: python %s gtf outfilename' % sys.argv[0]
        print '       this script only works on GTF files with transcript coordinates explicitly defined'
        sys.exit(1)

    gtf = sys.argv[1]
    outfilename = sys.argv[2]

    TranscriptDict={}
    TranscriptPrefixDict={}

    lineslist  = open(gtf)
    for line in lineslist:
        if line[0]=='#':
            continue
        fields=line.strip().split('\t')
        if fields[2]!='transcript':
            continue
        chr=fields[0]
        left=fields[3]
        right=fields[4]
        strand=fields[6]
        transcriptID=fields[8].split('transcript_id "')[1].split('";')[0]
        if TranscriptDict.has_key(transcriptID):
            pass
        else:
            TranscriptDict[transcriptID]=0
        TranscriptDict[transcriptID]+=1
        IDfields=transcriptID.split('.')
        prefix=''
        for i in range(len(IDfields)-1):
            prefix = prefix + IDfields[i] + '.'
        prefix=prefix[0:-1]
        if TranscriptPrefixDict.has_key(prefix):
            pass
        else:
            TranscriptPrefixDict[prefix]=[]
        try:
            TranscriptPrefixDict[prefix].append(int(IDfields[1]))
        except:
            TranscriptPrefixDict[prefix].append(IDfields[1])

    DuplicatesToNewNames={}

    for transcriptID in TranscriptDict.keys():
        if TranscriptDict[transcriptID] > 1:
            prefix=''
            IDfields=transcriptID.split('.')
            for i in range(len(IDfields)-1):
                prefix = prefix + IDfields[i] + '.'
            prefix=prefix[0:-1]
            DuplicatesToNewNames[transcriptID]={}
            for i in range(1,TranscriptDict[transcriptID]):
                maxending = max(TranscriptPrefixDict[prefix]) + 1
                TranscriptPrefixDict[prefix].append(maxending)
                DuplicatesToNewNames[transcriptID][maxending]=''

#        Duplicates=False
#        for transcriptID in GeneDict[geneID].keys():
#            if len(GeneDict[geneID][transcriptID]) > 1:
#                DuplicatesToNewNames[(geneID,transcriptID)]=''
#                Duplicates=True
#        if Duplicates:
#            endings=[]
#            for transcriptID in GeneDict[geneID].keys():
#                try:
#                    endings.append(int(transcriptID.split('.')[-1]))
#                except:
#                    continue
#            print geneID, transcriptID, len(GeneDict[geneID][transcriptID]), endings
#            for transcriptID in GeneDict[geneID].keys():
#                if DuplicatesToNewNames.has_key((geneID,transcriptID)):
#                     end=max(endings)+1
#                     endings.append(end)
#                     fields=transcriptID.split('.')
#                     newtranscriptID=''
#                     for i in range(len(fields)-1): 
#                         newtranscriptID=newtranscriptID+fields[i]+'.'
#                     newtranscriptID=newtranscriptID+str(end)
#                     DuplicatesToNewNames[(geneID,transcriptID)]=newtranscriptID

    print DuplicatesToNewNames

    outfile = open(outfilename, 'w')

    lineslist  = open(gtf)
    current=''
    substitute=''
    seenDict={}
    doSub=False
    i=0
    for line in lineslist:
        if line[0]=='#':
            outfile.write(line)
            continue
        i+=1
        fields=line.strip().split('\t')
        transcriptID=fields[8].split('transcript_id "')[1].split('";')[0]
        if fields[2] == 'transcript':
            if seenDict.has_key(transcriptID):
                doSub=True
                prefix=''
                IDfields=transcriptID.split('.')
                for i in range(len(IDfields)-1):
                    prefix = prefix + IDfields[i] + '.'
                prefix=prefix[0:-1]
                suffix = DuplicatesToNewNames[transcriptID].keys()[0]
                substitute = prefix + '.' + str(suffix)
                print transcriptID, substitute, DuplicatesToNewNames[transcriptID].keys()
                del DuplicatesToNewNames[transcriptID][suffix]
                print DuplicatesToNewNames[transcriptID].keys()
            else:
                doSub=False
        seenDict[transcriptID]=''
        if doSub:
            line=line.replace(transcriptID+'";',substitute+'";')
            outfile.write(line)
        else:
            outfile.write(line)

    outfile.close()
        
run()