#
#  utrChanges.py
#  ENRAGE
#

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

import sys
from commoncode import getMergedRegions, getLocusByChromDict
from cistematic.genomes import Genome

print "utrChanges: version 1.4"


def main(argv=None):
    if not argv:
        argv = sys.argv

    if len(argv) < 4:
        print "usage: python %s genome acceptedfile outfile" % argv[0]
        sys.exit(1)

    genome = argv[1]
    acceptfile =  argv[2]
    outfile = argv[3]

    utrChanges(genome, acceptfile, outfile)


def utrChanges(genome, acceptfile, outFileName):
    acceptDict = getMergedRegions(acceptfile, maxDist=0, keepLabel=True, verbose=True)
    outfile = open(outFileName, "w")

    hg = Genome(genome)

    origLocusByChromDict = getLocusByChromDict(hg, keepSense=True)
    newLocusByChromDict = getLocusByChromDict(hg, additionalRegionsDict=acceptDict, keepSense=True)

    new3utr = 0
    new5utr = 0
    changedGene = 0

    for chrom in origLocusByChromDict:
        for (gstart, gstop, gid, glen, sense) in origLocusByChromDict[chrom]:   
            for (newstart, newstop, newgid, newlen, newsense) in newLocusByChromDict[chrom]:
                if gid == newgid:
                    changedBoundary = False
                    new3p = "F"
                    new5p = "F"
                    if newstart < gstart:
                        if sense == "R":
                            new3utr += 1
                            new3p = "T"
                            changedBoundary = True
                        elif sense == "F":
                            new5utr += 1
                            new5p = "T"
                            changedBoundary = True
                        else:
                            print sense

                    if newstop > gstop:
                        if sense == "R":
                            new5utr += 1
                            new5p = "T"
                            changedBoundary = True
                        elif sense == "F":
                            new3utr += 1
                            new3p = "T"
                            changedBoundary = True
                        else:
                            print sense

                    if changedBoundary:
                        changedGene += 1
                        outfile.write("%s\tchr%s\t%d\t%d\t%s\tchr%s\t%d\t%d\t%s\t%s\n" % (gid, chrom, gstart, gstop, sense, chrom, newstart, newstop, new5p, new3p))

                    continue

    outfile.close()
    print "%d new 5'utr" % new5utr
    print "%d new 3'utr" % new3utr
    print "%s affected genes" % changedGene


if __name__ == "__main__":
    main(sys.argv)