#!/net/python/bin/python
from numpy import *
import pylab as p

def fac(x):
    if x < 0:
        return 0.0
    if x == 0:
        return 1.0
    else:
        return x*fac(x-1)

def Binomial(n,p,v):
    return fac(n)/( fac(v)*fac(n-v) ) * p**v * (1-p)**(n-v)

def Poisson(mu,v):
    return exp(-mu) * (mu**v)/(fac(v))

def Gaussian(x, mu, sig2):
    return exp( -(x-mu)**2/(2*sig2) ) / sqrt(2*pi*sig2)

def Double(seq):
    newseq = [0]
    for e in seq:
        newseq.append(e)
        newseq.append(e)
    newseq.append(0)
    return newseq

def DoubleR(seq):
    newseq = [seq[0]-0.5]
    for e in seq:
        newseq.append(e-0.5)
        newseq.append(e+0.5)
    newseq.append(seq[-1]+0.5)
    return newseq

def PlotDistribution(x,y):

    #x_step = (x[-1]-x[0])/( len(x)-1 )
    x_step = 1.0
    
    new_x = [x[0]-x_step/2]
    for e in x:
        new_x.append(e-x_step/2)
        new_x.append(e+x_step/2)
    new_x.append(x[-1]+x_step/2)
    
    new_y = [0]
    for e in y:
        new_y.append(e)
        new_y.append(e)
    new_y.append(0)

    p.plot( new_x,new_y )

def PlotGaussian(x,mu,sig2):
    y = [ Gaussian(e,mu,sig2) for e in x ]
    p.plot(x,y)

if __name__ == "__main__":
    
    for (P,N) in [ (0.4,5) , (0.4,50) ]:

        MU = P*N

        Range = arange(N+1)
        RangeG = arange(-1,(N+1),0.01)
        
        Bin = [ Binomial(N,P,V) for V in Range ]
        Pos = [ Poisson(MU,V) for V in Range ]
        
        sig2Bin = 1/(2*pi* (max(Bin))**2)
        sig2Pos = 1/(2*pi* (max(Pos))**2)
        
        PlotDistribution(Range,Bin)
        PlotGaussian(RangeG,MU,sig2Bin)
        p.title('Binomial, p=%s, N=%i'% (str(P),N) )
        p.savefig('Binomial%s_%i.ps'% (str(P),N) )
        p.close()
        
        PlotDistribution(Range,Pos)
        PlotGaussian(RangeG,MU,sig2Pos)
        p.title('Poisson, p=%s, N=%i'% (str(P),N) )
        p.savefig('Poisson%s_%i.ps'% (str(P),N) )
        p.close()