import sys
import optparse
from cistematic.genomes import Genome
from commoncode import getGeneInfoDict, getConfigParser, getConfigOption, getConfigBoolOption

print "getGOgenes: version 3.2"

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

    usage = "usage: python %s genome GOID1 [GOID2 ....] [--outfile outfilename] [--append] [--restrict genefile]"

    parser = makeParser(usage)
    (options, args) = parser.parse_args(argv[1:])

    if len(args) < 2:
        print usage
        sys.exit(1)

    genome = args[0]

    GOIDlist = []
    for arg in args:
        if "GO:" in arg:
            GOIDlist.append(arg)

    getGOgenes(genome, GOIDlist, options.outfilename, options.restrictfilename, options.append)


def makeParser(usage=""):
    parser = optparse.OptionParser(usage=usage)
    parser.add_option("--outfile", dest="outfilename")
    parser.add_option("--append", action="store_true", dest="append")
    parser.add_option("--restrict", dest="restrictfilename")

    configParser = getConfigParser()
    section = "getGOgenes"
    outfilename = getConfigOption(configParser, section, "outfilename", None)
    restrictfilename = getConfigOption(configParser, section, "restrictfilename", None)
    append = getConfigBoolOption(configParser, section, "append", False)

    parser.set_defaults(outfilename=outfilename, restrictfilename=restrictfilename, append=append)

    return parser


def getGOgenes(genome, GOIDlist, outfilename=None, restrictfilename=None, append=False):
    writeOut = False
    if outfilename is not None:
        writeOut = True

    restrict = False
    if restrictfilename is not None:
        restrict = True
    
    hg = Genome(genome)

    print sys.argv
    print GOIDlist

    firstGeneList = []
    for GOID in GOIDlist:
        testList = hg.allGIDsbyGOID(GOID)
        print "GOID: %s (%d)" % (GOID, len(testList))
        firstGeneList += testList

    geneDict = {}
    for gid in firstGeneList:
        geneDict[gid] = 1

    geneList = geneDict.keys()
    print len(geneList)
    geneInfoList = getGeneInfoDict(genome)

    if writeOut:
        if append:
            outfile = open(outfilename, "a")
        else:
            outfile = open(outfilename, "w")

        for GOID in GOIDlist:
            outfile.write("#%s\n" % GOID) 

    restrictList = []
    restrictDict = {}
    if restrict:
        restrictFile = open(restrictfilename)
        for line in restrictFile:
            fields = line.strip().split()
            restrictList.append(fields[0])
            restrictDict[fields[0]] = line

    outList = []
    symbolDict = {}
    for gid in geneList:
        symbol = "LOC" + gid
        if restrict and gid not in restrictList:
            continue

        try:
            symbol = geneInfoList[gid][0][0]
        except:
            pass

        if restrict:
            symbolDict[symbol] = restrictDict[gid]

        outList.append(symbol)

    outList.sort()
    for symbol in outList:
        if writeOut:
            if restrict:
                outfile.write(symbolDict[symbol])
            else:
                outfile.write(symbol + "\n")
        else:
            print symbol

    if writeOut:
        outfile.close()


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