#format python
"""PrimerDesign script using Primer3 by yong27, 2005-05-17

priming conditions refer to http://www.rfcgr.mrc.ac.uk/Software/EMBOSS/Apps/eprimer3.html
"""
import unittest, os, glob
from cStringIO import StringIO
from Bio import Fasta
from Bio.Emboss import Primer

class PrimerDesigner:
    temp_dir = 'temp'
    def __init__(self, ifile):
        self.fi = Fasta.Iterator(ifile, Fasta.RecordParser())
        self.pg = None

    def getNextFasta(self):
        return self.fi.next()

    def saveAllFastas(self):
        self.cleanUp()
        os.mkdir(self.temp_dir)
        for fasta in self.fi:
            fasta.title = fasta.title.replace(':','-')
            besid = '__'.join(fasta.title.split())
            file(os.path.join(self.temp_dir, besid), 'w').write(str(fasta))

    def cleanUp(self):
        if os.path.exists(self.temp_dir):
            for fname in os.listdir(self.temp_dir):
                os.remove(os.path.join(self.temp_dir, fname))
            os.rmdir(self.temp_dir)

    def doPrimer3(self, **keywds):
        self.saveAllFastas()
        for fname in os.listdir(self.temp_dir):
            fname = os.path.join(self.temp_dir, fname)
            cmds = ['eprimer3']
            cmds.extend(['-sequence', fname])
            cmds.extend(['-outfile', fname+'.eprimer3'])
            for k,v in keywds.iteritems():
                cmds.extend(['-'+k, v])
            os.system(' '.join(cmds))
        self.pg = self.primerGenerator()

    def getNextPrimer(self):
        return self.pg.next()

    def primerGenerator(self):
        ppser = Primer.Primer3Parser()
        for fname in glob.glob(os.path.join(self.temp_dir, '*.eprimer3')):
            try:
                primer = ppser.parse(file(fname)).primers[0]
            except IndexError:
                continue
            yield (os.path.splitext(os.path.split(fname)[1])[0], primer)

    def getPrimerInfos(self, primer, direction):
        seq = getattr(primer, direction+'_seq')
        tm = getattr(primer, direction+'_tm')
        return [direction, seq, str(tm), str(primer.size)]

    def writeXls(self, xfile):
        for fname, primer in self.pg:
            words = fname.split('__')
            xfile.write('\t'.join(words+self.getPrimerInfos(primer, 'forward'))+'\n')
            xfile.write('\t'.join(words+self.getPrimerInfos(primer, 'reverse'))+'\n')


class BreakPrimerDesigner(PrimerDesigner):
    def __init__(self, ifile):
        PrimerDesigner.__init__(self, ifile)
    def writeXls(self, xfile):
        xfile.write('\t'.join([
            'BAC', 'BES', 'FPC start', 'FPC end', 'Human match',
            'Direction', 'Primer sequence', 'Primer TM', 'Product size'
            ])+'\n')
        for fname, primer in self.pg:
            besid, bacid, start, end, humanMatch = fname.split('__')
            xfile.write('\t'.join([
                bacid, besid, start, end, humanMatch,
                ]+self.getPrimerInfos(primer, 'forward'))+'\n')
            xfile.write('\t'.join([
                bacid, besid, start, end, humanMatch,
                ]+self.getPrimerInfos(primer, 'reverse'))+'\n')


class ByunsBreakPrimerDesigner:
    def __init__(self, ifile):
        self.pd = BreakPrimerDesigner(ifile)
        self.pd.doPrimer3(otm='55', mintm='53')
    def write(self, wfile):
        self.pd.writeXls(wfile)


class TestPrimerDesigner(unittest.TestCase):
    def setUp(self):
        inputFile = StringIO("""\
>bE78F16SP6 CH242-78F16 46 72 18:11487196-11487708
GGTTGATGGTTCCCTGGTTTCATAAAACCTCTCAGGAACATTAGACTATGTCCGTGGAGA
GAAGTAAAACCCAATTATCCTCAGTGCCAGTAGTTTTTATGTCTCAGAAGAGATTCTGCT
GGAAGGATCACTGGGCATTTCAAGATGTTTGGAAACACTGCACTGATGGAGTTATCCTCA
GAGGGGTGTCATTAAATGTTTGTTTACGATTGTTGGTGGGTCTATCAAACCCACTTATTT
CTGCCAGCGTTAGATTTTTGACTCAATTAGGAAAAAAAAAGTATTAGAGGGAACTAACTT
TGCAGCAAAGGGCAGATACTCTTAGGTACTGAAAACTTGCTGAATTCTGTAGGAATAGGA
ACCTTGGGAATTATGCTTCCGTAGGAGAAATAATCAGAAAGGTCATTTTAGTCACTTAGG
AGTCTGTGGTGCCACTGGAATAATTAAAAACATGATAATTTAAACTGACTGTAGAGGTAA
TGGAAGAAACATGAAATTTAGTTAGACTAGGGTAGTTGAGTTTGTTCAAAAATCAAATGG
AAATCTTAGGGCACTACGGGAGCAGTTGAGAGTTGATCACAGCTGAGCAAGGCTGGTCTT
CTTCCTGCAGGTCCCGGCACCTTTGAGGATCTGACAAAGGCTGGCTGTCTTCTCTCAAGA
CAAGGGGCACGGGACCCCCAGTG
>bE78F16T7 CH242-78F16 46 72 18:11085776-11086061
TTGCTCTAACAGCCCTGAAATCCCCTGCAGCAGATTACTCTGATATAGGANGAANNGGCC
CCCNTTNCCNTTCACTGCTCTTTTTGACTTCCAAAAATAATATTCTGGCCCAGCTGTTCG
GTCAGGGCCTCGGTTTATTTGTAGTTAGCCAGAGAGGGGACCATAAATGAAATTCAGATG
CCGGGAGGGAGGAATGTCAGTCCTTCCAAATCCAACACTGAAAAGCCTTTCTTGTTTTTG
AGCTGAATAAAGAACGTCGGGCAATGGAACTTTGACCCAACATGCCATCCAGTCCAAAGA
GGAACGGCCAGAAAAAAGTGGACTTGTCTGGAGATGGTTCTCAGCCTCTAAGCCAGCCAG
CTCAGTCCCTTGTTTTAGCCCCCAGGGTATCTGGGCAGCCACAGGTCTGAGACACCCTTG
GTAGCTGAGGGCTCCCCTGGGCGCACAGGGCATGAGGGTGCCACCACGAAGAGGGGCCAC
AGGGCCCTTTTGTCAGCAAACAGAGCCTGGGCTCCAGGCTCCAGGTCTTGCAACAAATCC
CACCTCACC
""")
        self.pd = PrimerDesigner(inputFile)

    def testGetNextFasta(self):
        expected = "bE78F16SP6 CH242-78F16 46 72 18:11487196-11487708"
        self.assertEquals(expected, self.pd.getNextFasta().title)
        expected = "bE78F16T7 CH242-78F16 46 72 18:11085776-11086061"
        self.assertEquals(expected, self.pd.getNextFasta().title)

    def testSaveAllFastas(self):
        self.pd.saveAllFastas()
        expected = ['bE78F16SP6__CH242-78F16__46__72__18-11487196-11487708',
                    'bE78F16T7__CH242-78F16__46__72__18-11085776-11086061',
                   ]
        self.assertEquals(expected, os.listdir('temp'))

    def testDoPrimer3(self):
        self.pd.doPrimer3()
        expected = ['bE78F16SP6__CH242-78F16__46__72__18-11487196-11487708',
                    'bE78F16T7__CH242-78F16__46__72__18-11085776-11086061',
                    'bE78F16SP6__CH242-78F16__46__72__18-11487196-11487708.eprimer3',
                    'bE78F16T7__CH242-78F16__46__72__18-11085776-11086061.eprimer3',
                   ]
        expected.sort()
        self.assertEquals(expected, sorted(os.listdir('temp')))

    def testGetNextPrimer(self):
        self.pd.doPrimer3()
        fname, primer = self.pd.getNextPrimer()
        self.assertEquals('TCTGTGGTGCCACTGGAATA', primer.forward_seq)
        self.assertEquals('bE78F16SP6__CH242-78F16__46__72__18-11487196-11487708',fname)

    def testWriteXls(self):
        self.pd.doPrimer3()
        xfile = StringIO()
        self.pd.writeXls(xfile)
        self.assertEquals('bE78F16SP6', xfile.getvalue().splitlines()[1].split('\t')[0])

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