import scipy
import pylab

def RandomWalk(N=100, d=2):
    """
    Use scipy.cumsum and scipy.random.uniform to generate
    a d-dimensional random walk of length N, each of which has a
    random DeltaX and DeltaY between -1/2 and 1/2.
    You'll want to generate an array of shape (N,d), using (for example),
    scipy.random.uniform(min, max, shape).
    It is important to do the cumulative sum (cumsum) along the
    appropriate axis of the array.  The default behavior (axis=None)
    sums over all axes, reducing a multi-dimensional array of shape (N,d)
    to a one-dimensional array of length N*d.  To independently sum
    each of the d coordinates of each step in the walk, do the
    cumulative sum over axis=0.
    """
    return scipy.cumsum(scipy.random.uniform(-0.5,0.5,(N,d)), axis=0)

def PlotRandomWalkXT(N=100):
    """
    Plot X(t) for one-dimensional random walk 
    """
    X = RandomWalk(N,1)
    pylab.plot(X)
    pylab.show()

def PlotRandomWalkXY(N=100):
    """
    Plot X, Y coordinates of random walk where 
        X = walk[:,0]  # slicing to extract the 0 column on the first axis
        Y = walk[:,1]  # slicing to extract the 1 column on the first axis
    To make the X and Y axes the same length, 
    use pylab.figure(figsize=(8,8)) before pylab.plot(X,Y) and
    pylab.axis('equal') afterward.
    """
    walk = RandomWalk(N,d=2)
    X = walk[:,0]  # slicing to extract the 0 column on the first axis
    Y = walk[:,1]  # slicing to extract the 1 column on the first axis
    pylab.figure(figsize=(8,8))
    pylab.plot(X,Y)
    pylab.axis('equal')
    pylab.show()

def BoxCount(walk):
    w = walk
    h1 = scipy.histogram2d(w[:,0], w[:,1], [1,1])
    bbox = scipy.array([h1[1][1]-h1[1][0], h1[2][1]-h1[2][0]])
    histos = {}
    sizes = [1,2,4,8,16,32,64,128,256,512,1024]
    for n in sizes:
        if bbox[0] < bbox[1]:
            nx = n
            ny = int(nx*(bbox[1]/bbox[0]))
        else:
            ny = n
            nx = int(ny*(bbox[0]/bbox[1]))
        histos[n] = scipy.histogram2d(w[:,0], w[:,1], [nx,ny])
    boxes = dict([(min(bbox[0],bbox[1])/n, \
                   scipy.sum(histos[n][0]>0)) for n in histos])
    sk = sorted(boxes.keys())
    sv = [boxes[k] for k in sk]
    pylab.plot(sk, sv, 'bo-')
    pylab.xlabel('Box size (L)')
    pylab.ylabel('Number of boxes to cover (N)')
    pylab.loglog()
    return histos, boxes

def Endpoints(W=10000, N=10, d=2):
    """
    Returns a list of endpoints of W random walks of length N and
    dimension d. (In one dimension, this should return an array of
    one-element arrays, to be consistent with higher dimensions.)
    One can generate the random walks and then peel off the final positions,
    or one can generate the steps and sum them directly, since the
    intermediate positions are not needed here.
    scipy.random.uniform can be used to generate a 3-index array
    of the appropriate shape (containing W random walks of shape (N,d)).
    As described above in RandomWalk(), the sum (using scipy.sum)
    should be done along the appropriate axis of the randomly generated array.
    """
    return scipy.sum(scipy.random.uniform(-0.5,0.5,(W,N,d)), axis=1)

def PlotEndpoints(W=10000, N=10, d=2):
    """
    Plot endpoints of random walks.
    Use scipy.transpose to pull out X, Y. 
    To plot black points not joined by lines use pylab.plot(X, Y, 'k.')
    Again, use pylab.figure(figsize=(8,8)) before and
    pylab.axis('equal') afterward.
    """
    X, Y = scipy.transpose(Endpoints(W, N, d))
    pylab.figure(figsize=(8,8))
    pylab.plot(X,Y,'k.')
    pylab.axis('equal')
    pylab.show()

def HistogramRandomWalk(W=10000, N=10, d=1, bins=50):
    """
    Compares the histogram of random walks with the normal distribution
    predicted by the central limit theorem.
    #
    (1) Plots a histogram rho(x) of the probability that a random walk
    with N has endpoint X-coordinate at position x. 
    Uses pylab.hist(X, bins=bins, normed=1) to produce the histogram
    #
    (2) Calculates the RMS stepsize sigma for a random walk of length N
    (with each step uniform in [-1/2,1/2]
    Plots rho = (1/(sqrt(2 pi) sigma)) exp(-x**2/(2 sigma**2)) 
    for -3 sigma < x < 3 sigma on the same plot (i.e., before pylab.show).
    Hint: Create x using arange. Squaring, exponentials, and other operations
    can be performed on whole arrays, so typing in the formula for rho will
    work without looping over indices, except sqrt, pi, and exp need to be
    from the appropriate library (pylab, scipy, ...)
    """
    X = Endpoints(W, N, d)[:,0] 
    pylab.hist(X, bins=bins, normed=1)
    #
    sigma = scipy.sqrt(N/12.)
    x = scipy.arange(-3*sigma,3*sigma,sigma/bins)
    rho = (1/(scipy.sqrt(2*scipy.pi)*sigma))*scipy.exp(-x**2/(2*sigma**2))
    pylab.plot(x, rho, "k-")
    pylab.show()

def yesno():
    response = input('    Continue? (y/n) ')
    if len(response)==0:        # [CR] returns true
        return True
    elif response[0] == 'n' or response[0] == 'N':
        return False
    else:                       # Default
        return True
    
def demo():
    """Demonstrates solution for exercise: example of usage"""
    print("Random Walk Demo")
    print("Random Walk X vs. t")
    PlotRandomWalkXT(10000)
    if not yesno(): return
    print("Random Walk X vs. Y")
    PlotRandomWalkXY(10000)
    if not yesno(): return
    print("Endpoints of many random walks")
    print("N=1: square symmetry")
    PlotEndpoints(N=1)
    if not yesno(): return
    print("N=10: emergent circular symmetry")
    PlotEndpoints(N=10)
    if not yesno(): return
    print("Central Limit Theorem: Histogram N=10 steps")
    HistogramRandomWalk(N=10)
    if not yesno(): return
    print("1 step")
    HistogramRandomWalk(N=1)
    if not yesno(): return
    print("2 steps")
    HistogramRandomWalk(N=2)
    if not yesno(): return
    print("4 steps")
    HistogramRandomWalk(N=4)
    if not yesno(): return

if __name__=="__main__":
   demo()



# Copyright (C) Cornell University
# All rights reserved.
# Apache License, Version 2.0