##################################
#                                #
# Last modified 2023/07/18       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
from sets import Set
import math

def run():

    if len(sys.argv) < 2:
        print 'usage: python %s gtf outfile' % sys.argv[0]
        print 'NOTE: only use this script for non-spliced entries!!!'
        sys.exit(1)

    GTF = sys.argv[1]
    outfile = sys.argv[2]

    SeenDict = {}

    outfile = open(outfile, 'w')

    GeneDict={}
    linelist = open(GTF)
    for line in linelist:
        if line.startswith('#'):
            outfile.write(line)
            continue
        fields=line.strip().split('\t')
        CDSEXON = fields[2]
        if SeenDict.has_key(CDSEXON):
            pass
        else:
            SeenDict[CDSEXON] = {}
            SeenDict[CDSEXON]['geneID'] = {}
            SeenDict[CDSEXON]['geneName'] = {}
            SeenDict[CDSEXON]['transcriptID'] = {}
            SeenDict[CDSEXON]['transcriptName'] = {}
        geneID = fields[8].split('gene_id "')[1].split('";')[0].replace('/','').replace('(','').replace(')','').replace(' ','_')
        transcriptID = fields[8].split('transcript_id "')[1].split('";')[0].replace('/','').replace('(','').replace(')','').replace(' ','_')
        geneName = fields[8].split('gene_name "')[1].split('";')[0].replace('/','').replace('(','').replace(')','').replace(' ','_')
        transcriptName = fields[8].split('transcript_name "')[1].split('";')[0].replace('/','').replace('(','').replace(')','').replace(' ','_')
        if SeenDict[CDSEXON]['geneID'].has_key(geneID):
            SeenDict[CDSEXON]['geneID'][geneID] += 1
        else:
            SeenDict[CDSEXON]['geneID'][geneID] = 1
        if SeenDict[CDSEXON]['geneName'].has_key(geneName):
            SeenDict[CDSEXON]['geneName'][geneName] += 1
        else:
            SeenDict[CDSEXON]['geneName'][geneName] = 1
        if SeenDict[CDSEXON]['transcriptID'].has_key(transcriptID):
            SeenDict[CDSEXON]['transcriptID'][transcriptID] += 1
        else:
            SeenDict[CDSEXON]['transcriptID'][transcriptID] = 1
        if SeenDict[CDSEXON]['transcriptName'].has_key(transcriptName):
            SeenDict[CDSEXON]['transcriptName'][transcriptName] += 1
        else:
            SeenDict[CDSEXON]['transcriptName'][transcriptName] = 1
        outline = fields[0] + '\t' + fields[1] + '\t' + fields[2] + '\t' + fields[3] + '\t' + fields[4] + '\t' + fields[5] + '\t' + fields[6] + '\t' + fields[7] + '\t'
        outline = outline + 'gene_id "' + geneID + str(SeenDict[CDSEXON]['geneID'][geneID])
        outline = outline + '"; transcript_id "' + transcriptID + str(SeenDict[CDSEXON]['transcriptID'][transcriptID])
        outline = outline + '"; gene_name "' + geneName + str(SeenDict[CDSEXON]['geneName'][geneName])
        outline = outline + '"; transcript_name "' + transcriptName + str(SeenDict[CDSEXON]['transcriptName'][transcriptName])
        outline = outline + '";'
        outfile.write(outline + '\n')

    outfile.close()
   
run()