#format python

"""PmfSimulation - peptide length distribution using GnuPlotPy
"""

import unittest
import Gnuplot
"""This module require ImitProFound, enzyme module made by yong27's company
"""
import ImitProFound, enzyme

class PeptideStat:
    def __init__(self,aTax='',aStep=20):
        self.tax=aTax
        self.seqIter=self.getFastaSeq()
        self.step=20
        self.peptides=[]
        self.histo=[0]*50000
    def getFastaSeq(self):
        return ImitProFound.getFastaSeqIter('fasta_'+self.tax)
    def digestEnzymeBy(self,anEnzyme):
        while 1:
            fastaSeq = self.seqIter.next()
            if fastaSeq is None: break
            self.peptides.append(ImitProFound.PeptideMap(fastaSeq,anEnzyme))

    def getHistoData(self):
        if not self.peptides:
            print "you must execute diegstEnzymeBy(anEnzyme)"
            sys.exit()
        for protein in self.peptides:
            for peptide in protein.peptideMasses:
                self.histo[int(peptide/self.step)]+=1
        return self.histo
    def getHistoDataForGnuplot(self,aMethod):
        gpinput=[];i=0
        for each in aMethod():
            gpinput.append([i,each])
            i+=1
        return gpinput
    def getHistoDataByAaLen(self):
        if not self.peptides:
            print "you must execute diegstEnzymeBy(anEnzyme)"
            sys.exit()
        aaLen=[0]*5000
        for protein in self.peptides:
            for peptide in protein.peptSeq:
                aaLen[len(peptide)]+=1
        return aaLen


def main():
    pt=PeptideStat('Escherichia_coli')
    #pt=PeptideStat('All')
    pt.digestEnzymeBy(enzyme.Trypsin())

    g=Gnuplot.Gnuplot()
    g("set term png")
    g("set output 'plot.png'")
    g("set data style l")
    g("set xrange[0:50]")
    #g.plot(pt.getHistoDataForGnuplot(pt.getHistoData))
    g.plot(pt.getHistoDataForGnuplot(pt.getHistoDataByAaLen))

class TestPeptideStat(unittest.TestCase):
    def test1(self):
        pt=PeptideStat('Escherichia_coli')
        pt.digestEnzymeBy(enzyme.Trypsin())
        print pt.getHistoData()

if __name__=='__main__':
    #unittest.main(argv=('','-v'))
    main()