#Illustrating the plotting library with some simulations
#involving coin-tossing.

#This code uses the pylab plotting functions. You have a choice
#between using standard Python lists and using the enhanced
#array manipulation functions in numpy.  This code illustrates
#both approaches.



from pylab import *

#Caution: Some people object to importing the
#pylab package in this way.  The advantage is that
#you do not have to write the library name every time
#you invoke a function, but you can sometimes run into
#the problem of colliding namespaces. It is probably 
#more prudent to type 'import pylab' and invoke each
#function with the 'pylab.' prepended.


#Simulated coin tosses:  return a Python list of booleans, representing the
#result of n tosses of a coin with heads probability p.



def coins(p,n):
    """Returns a list of n tosses of a coin with Heads probability p.
    The return value is a list of booleans, with True denoting Heads"""
    tosses=[]
    for j in range(n):
        tosses.append(rand()<p)
    return tosses

#Here is a fancier (and more efficient) way do do this, using numpy arrays.

def coins2(p,n):
    """Returns an array of n tosses of a coin with Heads probability p.
    The return value is a numpy  sequence of booleans, with True denoting Heads"""
    return (rand(n)<p)



#Heads minus tails: A player tosses a  coin with heads probability p
#n times in succession. The function returns a list whose ith entry is the
#number of heads minus the number of tails in the first i tosses.  This is the
#same as the accumulated winnings of the player, if they win one penny
#for each toss of Heads and lose one for each toss of Tails.
#Does this difference tend to 0 as the number of tosses grows?



def hmt(p,n):
    """Returns a list of length n+1 giving at index i the accumulated number of
heads minus the number of tails after i tosses"""
    booltosses=coins(p,n)
    diffs=[0]
    hmt=0
    for j in range(n):
        hmt+=2*booltosses[j]-1 #note if b boolean, 2*b-1 is 1 or -1
        diffs.append(hmt)
    return diffs

#Here's the same thing done fancier, using numpy arrays (treating booleans
#as integers, multiplying  arrays by constants, adding constants to arrays,
#and using the cumulative sum function.

def hmt2(p,n):
    """Returns an array giving at each entry the accumulated number of
heads minus the number of tails"""
    v=cumsum(2*coins2(p,n)-1)
    #we prepend 0
    return concatenate(([0],v))
    
    
#This demonstrates how to use some of the plotting tools.  The result
#of the simulation in the preceding function is displayed graphically.
#At the heart is the function plot, which here takes two lists of
#equal length as successive x and y values, and draws the corresponding
#line graph.  We have supplemented this with calls to functions that add
#the x-axis to a plot, as well as a plot title and axis labels.  The
#triple quotes in the plot title enable you to include line breaks within
#a long title. You can use either lists, as is done here, or numpy arrays
#(using hmt2 in place of hmt) as the arguments to the plot function.

def plot_hmt(p,n):
    """Draws a plot of the accumulated winnings returned by hmt"""
    winnings=hmt(p,n)
    plot(range(n+1),winnings)
    axhline(y=0,xmin=0,xmax=n,color='k')
    title("""Winnings in """+ str(n)+""" tosses of
a coin with heads probability """+str(p))
 
    ylabel("winnings")
    show()

#Perform the simulation above k times, and superimpose the plots of
#all the simulations.  Note that successive calls to plot before the
#call to show add new lines to the plot,choosing a new color each time.
#This figure, with p=0.5, appears in the textbook.
    
def multiplot_hmt(p,n,k):
    """Draw k superimposed plots of the accumulated winnings returned
       by headminustails."""
    for j in range(k):
        winnings=hmt(p,n)
        plot(range(n+1),winnings,linewidth=0.5)
    axhline(y=0,xmin=0,xmax=n, color='k')
    title("""Winnings in """+ str(n)+""" tosses of
a coin with heads probability """+str(p))
    xlabel("tosses")
    ylabel("winnings")
    show()
    
        

#By now you should be convinced that for a fair coin, the absolute difference
#between the number of heads and tails does NOT approach 0.  The next simulation
#exhibits the real story,namely that the PROPORTION of heads to the total number
#of tosses approaches 1/2.

#Here we put aside Python lists and just use the array-manipulation capabilities
#of numpy.  The functions
#arange and cumsum return such arrays, and the division 
#statements perform componentwise division of arrays. Here we carry out the
#simulation with two different coins, having two different proabilities of heads.
#Observe that we can take the cumulative sum of an array of booleans!--True
#and False are treated as 1 and 0, respectively.
    
#Note that we have also added a legend to the plot, to distinguish which line
#represents which coin.
    
def heads_rel(p,q,n):
    tosses1=coins2(p,n)
    tosses2=coins2(q,n)
    one_to_n = arange(1.0,n+1,1.0)
    headsratio1=cumsum(tosses1)/one_to_n #see how easy this is!
    headsratio2=cumsum(tosses2)/one_to_n
    #Now draw the plot
    plot(one_to_n,headsratio1,label = "prob="+str(p))
    plot(one_to_n,headsratio2,label = "prob="+str(q))
    axhline(y=p,xmin=0,xmax=n,ls="--")
    axhline(y=q,xmin=0,xmax=n,ls="--")
    title("""Ratio of heads to tosses with """+ str(n)+""" tosses
of a two coins with different success probabilities""")
    legend(loc="lower right")
    show()

#Just to illustrate some other plotting options:  Here we
#toss a fair coin n times and draw a scatterplot of the result, with 0
#for tails and 1 for heads. A scatterplot may seem like an odd choice
#for data of this kind, but the resulting picture is informative:
#The number of large clusters that appear shows long streaks in the data,
#more than our intuition might suggest. Once again, we can treat the array
#of booleans as an array of integers without explicit conversion.
    
def coinscatter(n):
    tosses=coins2(0.5,n)
    
    #the parameter s in scatter is the marker size
    scatter(range(n),tosses,s=5)
    xlim((0,n))
    yticks(range(2),('Tails','Heads'))
    title("Distribution of heads and tails in "+ str(n)+ " tosses of a fair coin.")
    show()



#Toss m coins n times and draw a histogram of the number of heads.
#There is an option to make it cumulative.

#Here we make further use of our array-manipulation functions:  The rand
#function is used here to generate a 2-dimensional array of booleans.
#This is further converted into a 2-dimensional array of 0's and 1's, and
#then summed along each row to get a one-dimensional array of values between
#0 and m. This built-in hist function takes this array and a specification
#of bins, and draws a column chart of the number of samples in each bin.

def coinhist(p,m,n,cum=False):
    #n rows of m columns of coin-tosses
    tosses=(rand(n,m)<p)
    #sum along rows
    heads=sum(tosses,axis=1)
    #"bins" for the histogram.  We get a column chart of
    #the number of elements in heads between each successive pair
    #of bin markers
    bins=arange(-0.5,m+1.25,1.0)
    title(str(n) + " tosses of "+ str(m) +""" coins
with success probability """+str(p))
    xlim(-1,m+1)
    xlabel("Number of heads")
    ylabel("Number of tosses")
    hist(heads,bins,color="g",cumulative=cum,edgecolor='k')
    show()

#We'll repeat this last problem by displaying the histogram as a stem plot.
#The his function in matplotlib does not provide this as an option, so instead
#we use the histogram function in numpy, which produces the histogram data
#without plotting anything, and pass it to the stem function in matplotlib.
def coinhist_stem(p,m,n):
    #n rows of m columns of coin-tosses
    tosses=(rand(n,m)<p)
    #sum along rows
    heads=sum(tosses,axis=1)
    #"bins" for the histogram.  We get a column chart of
    #the number of elements in heads between each successive pair
    #of bin markers
    bins=arange(-0.5,m+1.5,1)
    h=histogram(heads,bins)
    print (h)
    stem(range(m+1),h[0],linefmt='k-',markerfmt='ko')
    title(str(n) + " tosses of "+ str(m) +""" coins
with success probability """+str(p))
    xlim(-1,m+1)
    xlabel("Number of heads")
    ylabel("Number of tosses")
    show()