#!/usr/bin/python
# Central limit theorem demonstrations
def convolve3( T , p , \
               title='Central limit demonstration' , \
               when_to_stop_pausing=20 ):
    """
    Show the probability distribution
    of a sum of t random variables
    (t=1....T)
    each of which is drawn from distribution p. [p is a list of floats]
    p is an arbitrary distribution over integers (0,1,2,3,....I_max)

    It is recommended, but not required, that p should be
    positive or zero and should sum to 1.0.

    example:
    >>> from convolve1 import *
    >>> convolve1( 30, [0.1, 0, 0, 0.4, 0.2, 0.3] )

    convolve3 is just like convolve1 except it does a log-scale plot too
    """

    print "Convolve version 1.2"
##    from Numeric import *  ## supplies convolve
    import scipy as S
    import Gnuplot, Gnuplot.funcutils
    g = Gnuplot.Gnuplot(persist=1) #(debug=1, persist=1)
    g.title(title) 
    g('set data style impulse') 
    gl = Gnuplot.Gnuplot(persist=1)
    gl.title(title) 
    gl('set data style impulse') 
    gl('set logscale y') 

    y  = S.array(p) # we will store the current distribution in y
    p0 = S.array(p) # we will convolve y with p0
    Number = 1
    keep_pausing = 1 
    for t in range(T):
        # Do the plot, ensuring the xrange is a bit bigger
        # than the width of the plots, and that the yrange is
        # exactly snug.
        g( "set xrange [-0.5:%g]" % (len(y)-0.5 ) )
        g( "set yrange [0.0:%g]" % (max(y)*1.01) )
        g.title( title+" - t = %d" % (t+1) ) 
        g.plot(y)
        gl( "set xrange [-0.5:%g]" % (len(y)-0.5 ) )
        gl( "set yrange [%g:%g]" % (min([y[0],y[-1]])*0.5, max(y)*1.5) )
        gl.title( title+" - t = %d" % (t+1) ) 
        gl.plot(y)
        # end of plot
        if(keep_pausing and (t<when_to_stop_pausing)):
            if ( raw_input('Press return to continue (or 0 to stop pausing)...\n') == '0' ):
                keep_pausing = 0 
            pass
        # This is the key step: 
        y = S.convolve( y , p0 )
        # count the number of variables we have summed 
        Number += 1
        pass
    pass

def test():    
    convolve3( 3, [0.166, 0.166, 0.166, 0.166, 0.166, 0.166] , 'Dice' )
    raw_input('Press return to continue ...\n')
    convolve3( 1000, [0.9, 0.1] , 'Bent coin (0.9,0.1)' , 50 )
    raw_input('Press return to continue ...\n') 
    convolve3( 300, [0.1, 0.0, 0.0, 0.0, 0.2, 0.0, 0.4, 0.3] , '[1,0,0,0,2,0,4,3]' )
    raw_input('Press return to continue ...\n')
    convolve3( 500, [0.5,0.5] , 'Fair coin (0.5,0.5)' )
    pass

if __name__ == '__main__': test()