from commons.core.utils.FileUtils import FileUtils
import os
import unittest
from commons.launcher.LaunchTallymer import ConvertUtils

class Test_ConvertUtils(unittest.TestCase):

    def test_convertTallymerFormatIntoMapFormatAndGenerateData(self):
        inputFastaFileName = "input.fa"
        inputTallymerSearchFileName = "input.tallymer"
        expMapFileName = "exp.map"
        obsMapFileName = "obs.map"
        self._writeInputFasta(inputFastaFileName)
        self._writeInputTallymerSearch(inputTallymerSearchFileName)
        self._writeExpMap(expMapFileName)
        
        expOccNb = 4
        expDKmer2Occ = {"caagatgcgtaacggccata":1,
                        "aagatgcgtaacggccatac":1,
                        "agatgcgtaacggccataca":1,
                        "gatgcgtaacggccatacat":1}
        expPlotData = {4379: 17,
                       4380: 17,
                       4381: 16,
                       4382: 16}
        
        obsOccNb, obsDKmer2Occ, obsPlotData, obsPlotData2 = ConvertUtils.convertTallymerFormatIntoMapFormatAndGenerateData(inputFastaFileName, inputTallymerSearchFileName, obsMapFileName)
        
        self.assertEquals(expOccNb, obsOccNb)
        self.assertEquals(expDKmer2Occ, obsDKmer2Occ)
        self.assertEquals(expPlotData, obsPlotData)
        self.assertTrue(FileUtils.are2FilesIdentical(expMapFileName, obsMapFileName))
        
        os.remove(inputFastaFileName)
        os.remove(inputTallymerSearchFileName)
        os.remove(expMapFileName)
        os.remove(obsMapFileName)
        
    def test_convertTallymerFormatIntoMapFormatAndGenerateDataMultiFasta(self):
        inputFastaFileName = "input.fa"
        inputTallymerSearchFileName = "input.tallymer"
        expMapFileName = "exp.map"
        obsMapFileName = "obs.map"
        self._writeInputMultiFasta(inputFastaFileName)
        self._writeInputTallymerSearchMultiFasta(inputTallymerSearchFileName)
        self._writeExpMapMultiFasta(expMapFileName)
        
        expOccNb = 6
        expDKmer2Occ = {"caagatgcgtaacggccata":1,
                        "aagatgcgtaacggccatac":1,
                        "agatgcgtaacggccataca":1,
                        "gatgcgtaacggccatacat":1,
                        "agatgcgtaacggccataca":2,
                        "gatgcgtaacggccatacat":2}
        expPlotData = {4379: 17, 4380: 17, 4381: 16, 4382: 35}
        
        obsOccNb, obsDKmer2Occ, obsPlotData, obsPlotData2 = ConvertUtils.convertTallymerFormatIntoMapFormatAndGenerateData(inputFastaFileName, inputTallymerSearchFileName, obsMapFileName)
        
        self.assertEquals(expOccNb, obsOccNb)
        self.assertEquals(expDKmer2Occ, obsDKmer2Occ)
        self.assertEquals(expPlotData, obsPlotData)
        self.assertTrue(FileUtils.are2FilesIdentical(expMapFileName, obsMapFileName))
        
        os.remove(inputFastaFileName)
        os.remove(inputTallymerSearchFileName)
        os.remove(expMapFileName)
        os.remove(obsMapFileName)
    
    def _writeInputTallymerSearch(self, tallymerFormatFileName):
        with open(tallymerFormatFileName, 'w') as f:
            f.write("0\t+4378\t17\tcaagatgcgtaacggccata\n")
            f.write("0\t+4379\t17\taagatgcgtaacggccatac\n")
            f.write("0\t+4380\t16\tagatgcgtaacggccataca\n")
            f.write("0\t+4381\t16\tgatgcgtaacggccatacat\n")
             
    def _writeInputFasta(self, fastaFileName):
        with open(fastaFileName, 'w') as f:
            f.write(">Seq1\n")
            f.write("caagatgcgtaacggccata\n")
            f.write("aagatgcgtaacggccatac\n")
            f.write("agatgcgtaacggccataca\n")
            f.write("gatgcgtaacggccatacat\n")
            
    def _writeExpMap(self, mapFileName):
        with open(mapFileName, 'w') as f:
            f.write("caagatgcgtaacggccata_17\tSeq1\t4379\t4399\t17\n")
            f.write("aagatgcgtaacggccatac_17\tSeq1\t4380\t4400\t17\n")
            f.write("agatgcgtaacggccataca_16\tSeq1\t4381\t4401\t16\n")
            f.write("gatgcgtaacggccatacat_16\tSeq1\t4382\t4402\t16\n")
            
    def _writeInputTallymerSearchMultiFasta(self, tallymerFormatFileName):
        with open(tallymerFormatFileName, 'w') as f:
            f.write("0\t+4378\t17\tcaagatgcgtaacggccata\n")
            f.write("0\t+4379\t17\taagatgcgtaacggccatac\n")
            f.write("0\t+4380\t16\tagatgcgtaacggccataca\n")
            f.write("0\t+4381\t16\tgatgcgtaacggccatacat\n")
            f.write("1\t+4380\t16\tagatgcgtaacggccataca\n")
            f.write("1\t+4381\t35\tgatgcgtaacggccatacat\n")
            
    def _writeInputMultiFasta(self, fastaFileName):
        with open(fastaFileName, 'w') as f:
            f.write(">Seq1\n")
            f.write("caagatgcgtaacggccata\n")
            f.write("aagatgcgtaacggccatac\n")
            f.write("agatgcgtaacggccataca\n")
            f.write("gatgcgtaacggccatacat\n")
            f.write(">Seq2\n")
            f.write("agatgcgtaacggccataca\n")
            f.write("gatgcgtaacggccatacat\n")
            
    def _writeExpMapMultiFasta(self, mapFileName):
        with open(mapFileName, 'w') as f:
            f.write("caagatgcgtaacggccata_17\tSeq1\t4379\t4399\t17\n")
            f.write("aagatgcgtaacggccatac_17\tSeq1\t4380\t4400\t17\n")
            f.write("agatgcgtaacggccataca_16\tSeq1\t4381\t4401\t16\n")
            f.write("gatgcgtaacggccatacat_16\tSeq1\t4382\t4402\t16\n")
            f.write("agatgcgtaacggccataca_16\tSeq2\t4381\t4401\t16\n")
            f.write("gatgcgtaacggccatacat_35\tSeq2\t4382\t4402\t35\n")

if __name__ == "__main__":
    unittest.main()