'''
Created on Aug 25, 2010

@author: sau
'''
import unittest
import string
import os
from erange import chksnp

dbPath = "/Users/sau/work/snpdb/hg18"

class TestChksnp(unittest.TestCase):
    """ First entries from snpDB using select func, name, start, stop from snp where chrom="1" limit 4;
        unknown|rs10218492|690|691
        unknown|rs10218493|766|767
        unknown|rs10218527|789|790
        unknown|rs28853987|800|801

        Entry from altSnpDB not in sndDB
        unknown|rs17160650|81751|81752
    """

    snpDB = "%s/dbSNP128.db" % dbPath
    altSnpDB = "%s/snp129cDNA.db" % dbPath
    inputFileName = "testChkSNP_input.txt"
    outputFileName = "testChkSNP_output.txt"

    def setUp(self):
        pass


    def tearDown(self):
        pass


    def testChkSNPFile(self):
        infile = open(self.inputFileName, "w")
        infile.write("# header line\n")
        snpEntry = string.join(["foo", "foo", "chr1", "691"], "\t")
        infile.write("%s\n" % snpEntry)
        snpEntry = string.join(["foo2", "foo2", "chr1", "81752"], "\t")
        infile.write("%s\n" % snpEntry)
        infile.close()

        chksnp.chkSNPFile(self.snpDB, self.inputFileName, self.outputFileName)
        outfile = open(self.outputFileName, "r")
        line = outfile.readline()
        result = "foo\tfoo\tchr1\t691\trs10218492\tunknown\n"
        self.assertEquals(result, line)
        result = "foo2\tfoo2\tchr1\t81752\tN\\A\tN\\A\n"
        line = outfile.readline()
        self.assertEquals(result, line)
        outfile.close()
        os.remove(self.outputFileName)

        chksnp.chkSNPFile(self.snpDB, self.inputFileName, self.outputFileName, snpDBList=[self.altSnpDB])
        outfile = open(self.outputFileName, "r")
        line = outfile.readline()
        result = "foo\tfoo\tchr1\t691\trs10218492\tunknown\n"
        self.assertEquals(result, line)
        result = "foo2\tfoo2\tchr1\t81752\trs17160650\tunknown\n"
        line = outfile.readline()
        self.assertEquals(result, line)
        outfile.close()

        os.remove(self.inputFileName)
        os.remove(self.outputFileName)


    def testMain(self):
        infile = open(self.inputFileName, "w")
        infile.write("# header line\n")
        snpEntry = string.join(["foo", "foo", "chr1", "691"], "\t")
        infile.write("%s\n" % snpEntry)
        snpEntry = string.join(["foo2", "foo2", "chr1", "81752"], "\t")
        infile.write("%s\n" % snpEntry)
        infile.close()

        argv = ["chksnp", self.snpDB, self.inputFileName, self.outputFileName]
        chksnp.main(argv)
        outfile = open(self.outputFileName, "r")
        line = outfile.readline()
        result = "foo\tfoo\tchr1\t691\trs10218492\tunknown\n"
        self.assertEquals(result, line)
        result = "foo2\tfoo2\tchr1\t81752\tN\\A\tN\\A\n"
        line = outfile.readline()
        self.assertEquals(result, line)
        outfile.close()

        os.remove(self.inputFileName)
        os.remove(self.outputFileName)


    def testChkSNP(self):
        snpPropertiesList = []
        dbList = [self.snpDB]
        self.assertEquals({}, chksnp.chkSNP(dbList, snpPropertiesList))

        snpPropertiesList = ["# header line"]
        snpEntry = string.join(["foo", "foo", "chr1", "691"], "\t")
        snpPropertiesList.append(snpEntry)
        snpEntry = string.join(["foo2", "foo2", "chr1", "81752"], "\t")
        snpPropertiesList.append(snpEntry)
        dbList = [self.snpDB, self.altSnpDB]
        result = {("1", 691): "foo\tfoo\tchr1\t691\trs10218492\tunknown",
                  ("1", 81752): "foo2\tfoo2\tchr1\t81752\trs17160650\tunknown"}
        self.assertEquals(result, chksnp.chkSNP(dbList, snpPropertiesList))


    def testGetSNPLocationInfo(self):
        snpPropertiesList = []
        snpEntry = string.join(["foo", "foo", "chr1", "20"], "\t")
        snpPropertiesList.append(snpEntry)
        snpLocationList, snpDict = chksnp.getSNPLocationInfo(snpPropertiesList)
        self.assertEquals([("1", 20)], snpLocationList)
        self.assertEquals({("1", 20): "foo\tfoo\tchr1\t20"}, snpDict)

        snpPropertiesList = ["# header line"]
        snpEntry = string.join(["foo", "foo", "chr1", "20"], "\t")
        snpPropertiesList.append(snpEntry)
        snpLocationList, snpDict = chksnp.getSNPLocationInfo(snpPropertiesList)
        self.assertEquals([("1", 20)], snpLocationList)
        self.assertEquals({("1", 20): "foo\tfoo\tchr1\t20"}, snpDict)


    def testDoNotProcessLine(self):
        self.assertTrue(chksnp.doNotProcessLine("#anything"))
        self.assertFalse(chksnp.doNotProcessLine("line to process"))


    def testAnnotateSNPFromDB(self):
        snpLocationList = [("1", 691), ("1", 81752)]
        snpDict = {("1", 691): "foo\tfoo\tchr1\t691",
                   ("1", 81752): "foo2\tfoo2\tchr1\t81752"}
        result = {("1", 691): "foo\tfoo\tchr1\t691\trs10218492\tunknown",
                  ("1", 81752): "foo2\tfoo2\tchr1\t81752\tN\\A\tN\\A"}
        self.assertEquals(result, chksnp.annotateSNPFromDB(snpLocationList, snpDict, self.snpDB))

        snpLocationList = [("1", 691), ("1", 81752)]
        snpDict = {("1", 691): "foo\tfoo\tchr1\t691",
                   ("1", 81752): "foo2\tfoo2\tchr1\t81752"}
        result = {("1", 691): "foo\tfoo\tchr1\t691\tN\\A\tN\\A",
                  ("1", 81752): "foo2\tfoo2\tchr1\t81752\trs17160650\tunknown"}
        self.assertEquals(result, chksnp.annotateSNPFromDB(snpLocationList, snpDict, self.altSnpDB))


    def testAnnotateSNPFromDBList(self):
        snpLocationList = []
        snpDict = {}
        dbList = [self.snpDB]
        self.assertEquals({}, chksnp.annotateSNPFromDBList(snpLocationList, snpDict, dbList))

        snpLocationList = [("1", 21)]
        snpDict = {("1", 21): "foo\tfoo\tchr1\t21"}
        dbList = [self.snpDB]
        result = {("1", 21): "foo\tfoo\tchr1\t21\tN\\A\tN\\A"}
        self.assertEquals(result, chksnp.annotateSNPFromDBList(snpLocationList, snpDict, dbList))

        snpLocationList = [("1", 21)]
        snpDict = {("1", 21): "foo\tfoo\tchr1\t21"}
        dbList = [self.snpDB]
        result = {("1", 21): "foo\tfoo\tchr1\t21\tN\\A\tN\\A"}
        self.assertEquals(result, chksnp.annotateSNPFromDBList(snpLocationList, snpDict, dbList, cachePages=10000))

        snpLocationList = [("1", 691)]
        snpDict = {("1", 691): "foo\tfoo\tchr1\t691"}
        dbList = [self.snpDB]
        result = {("1", 691): "foo\tfoo\tchr1\t691\trs10218492\tunknown"}
        self.assertEquals(result, chksnp.annotateSNPFromDBList(snpLocationList, snpDict, dbList))

        snpLocationList = [("1", 691), ("1", 81752)]
        snpDict = {("1", 691): "foo\tfoo\tchr1\t691",
                   ("1", 81752): "foo2\tfoo2\tchr1\t81752"}
        dbList = [self.snpDB]
        result = {("1", 691): "foo\tfoo\tchr1\t691\trs10218492\tunknown",
                  ("1", 81752): "foo2\tfoo2\tchr1\t81752\tN\\A\tN\\A"}
        self.assertEquals(result, chksnp.annotateSNPFromDBList(snpLocationList, snpDict, dbList))

        snpLocationList = [("1", 691), ("1", 81752)]
        snpDict = {("1", 691): "foo\tfoo\tchr1\t691",
                   ("1", 81752): "foo2\tfoo2\tchr1\t81752"}
        dbList = [self.snpDB, self.altSnpDB]
        result = {("1", 691): "foo\tfoo\tchr1\t691\trs10218492\tunknown",
                  ("1", 81752): "foo2\tfoo2\tchr1\t81752\trs17160650\tunknown"}
        self.assertEquals(result, chksnp.annotateSNPFromDBList(snpLocationList, snpDict, dbList))


def suite():
    suite = unittest.TestSuite()
    suite.addTest(unittest.makeSuite(TestChksnp))

    return suite


if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testName']
    unittest.main()