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

def GaussianFunc(x,mu=0.0,sig=1.0):
    return exp( -((x-mu)**2)/(2*sig**2) ) / ( sig * sqrt(2.0*pi) )

def Gaussian(N=10000,mu=0,sig=1,method='Box-Muller'):
    if method not in ['Box-Muller','vonNeumann', 'Metropolis', 'HWcode']:
        print 'Bad Method'
        return

    if method == 'Box-Muller':
        rand = random.random(N+1) #add one in case N is odd

        for i in range((N+1)/2):
            r1 = rand[2*i]
            r2 = rand[2*i+1]
            rand[2*i] = mu + sig * sqrt(-2*log(r1))*sin(2*pi*r2)
            rand[2*i+1] = mu + sig * sqrt(-2*log(r1))*cos(2*pi*r2)

        return rand[:-1] #truncate back to N values
    
    if method == 'vonNeumann':
        xmin = -5*sig
        xmax = 5*sig
        ymax = 1/sqrt(2*pi*sig**2)

        randlist = []
        count = 0
        while True:
            r = random.random(2)
            
            r[1] = r[1] * ymax
            #---------------------------------------------
            #exponential distribution, reflected about origin
            #  is a good weight function for (mu,sig)=(0,1)
            if count <= N/2:
                r[0] = -2 * log(r[0])
            else:
                r[0] = 2 * log(r[0])
            #--------------------------------------------
            #Or implement a flat weight function
            #
            #r[0] = xmin + (xmax-xmin) * r[0]
            #--------------------------------------------
                
            if r[1] < GaussianFunc(r[0],mu,sig):
                randlist.append(r[0])
                count += 1
            if count >= N:
                return randlist

    if method == 'Metropolis':
        initial = 0
        step_size = 1

        randlist = [initial]
        count = 1
        while True:
            rand = random.random()
            x_i = randlist[-1]
            x_t = x_i + step_size * (2*rand-1)
            r = GaussianFunc(x_t,mu,sig)/GaussianFunc(x_i,mu,sig)

            if r >= random.random():
                randlist.append(x_t)
                count += 1

            if count >= N:
                return randlist

    if method == 'HWcode':
        n=12.0
        return [sum(random.random(n))-n/2 for i in range(N)]


def plothist(seq, avg_per_bin = 200):
    
    MAX = max(seq)
    MIN = min(seq)
    N = len(seq)/avg_per_bin
    
    hist = zeros(N)

    for val in seq:
        index = N * (val-MIN)/(MAX-MIN)
        if index == N: index = N-1
        hist[index] += 1

    p.plot(MIN+(0.5+arange(N))*(MAX-MIN)*1.0/N,hist/( avg_per_bin*(MAX-MIN) ))


if __name__ == '__main__':

    def PlotWeight():
        x = arange(100.0) * 0.2 - 10.0
        p.plot( x,GaussianFunc(x) )
        #p.plot(x, 0.5*e**(-abs(0.5*x)) )
        p.plot(x,0.4-0.02*x**2)
        p.axis([-6,6,0,0.55])

    PlotWeight()
    p.show()
    
    method = ['Box-Muller','vonNeumann', 'Metropolis', 'HWcode']
    subplots = [221,222,223,224]

    x = arange(100.0) * 0.2 - 10.0
    
    for i in range(4):
        p.subplot(subplots[i])
        c=clock()
        plothist( Gaussian(10000,0,1,method[i]) )
        c=clock()-c
        p.plot( x,GaussianFunc(x,0,1) )
        p.axis([-10,10,0,0.45])
        p.title(method[i])
        p.text(5,0.35,'%.2f s' % c)
    p.savefig('4plots.ps')