#
#  makeSNPtrack.py
#  ENRAGE
#
# This script maps all the qualified SNC sites on to the genome browser 
# Output format: bed
# Written by: Wendy Lee
# Written on: August 18th, 2008
# Last Modified: December 14th, 2008 by Ali Mortazavi

import sys

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

    print "makeSNPtrack: version 1.3"

    if len(argv) < 4:
        print "usage: python %s snpfile trackname trackoutfile" % argv[0]
        sys.exit(1)

    snpfile = argv[1]
    track = argv[2]
    outfile = argv[3]

    makeSNPtrack(snpfile, track, outfile)


def makeSNPtrack(snpfilename, track, outfilename):

    snpfile = open(snpfilename, "r")
    writeSNPsBedfile(snpfile, track, outfilename)
    snpfile.close()


def writeSNPsBedfile(snpPropertiesList, track, outfilename):

    outfile = open(outfilename, "w")
    header = getHeader(track)
    outfile.write(header)

    for line in snpPropertiesList:
        if doNotProcessLine(line):
            continue

        fields = line.strip().split()
        outline = getBedOutputLine(fields)
        outfile.write(outline)

    outfile.close()


def getHeader(track):
    header = "track name=%s description=%s visibility=2 itemRgb=\"On\"\n" % (track, track)
    return header


def doNotProcessLine(line):
    return line[0] == "#"


def getBedOutputLine(snpPropertiesList):
    chromosome = snpPropertiesList[2]
    readStart = int(snpPropertiesList[3]) - 1
    readStop = readStart + 1
    readName = snpPropertiesList[7]
    color = getSNPColor(readName)
    score = "0"
    sense = "+"
    outline = "%s\t%d\t%d\t%s\t%s\t%s\t-\t-\t\t%s\n" % (chromosome, readStart, readStop, readName, score, sense, color)

    return outline


def getSNPColor(readName):
    baseColor = {"A": "200, 0, 255",
                 "T": "200, 0, 255",
                 "C": "200, 0, 255",
                 "G": "200, 0, 255"
    }

    specialColors = {"A-G": "255, 0, 0",
                     "T-C": "0, 0, 255"
    }

    if readName in specialColors.keys():
        color = specialColors[readName]
    else:
        try:
            color = baseColor[readName[-1]]
        except (IndexError, KeyError):
            color = "200, 0, 255"

    return color


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