import numpy as np
import matplotlib.pyplot as plt
from plotstyle import *

def random_walk(n, p, rng):
    steps = rng.rand(n)
    steps[steps>=p] = 1
    steps[steps<p] = -1
    x = np.hstack(((0.0,), np.cumsum(steps)))
    return x

if __name__ == "__main__":
    fig, ax = plt.subplots()
    fig.subplots_adjust(**adjust_fs(fig, right=1.0))
    rng = np.random.RandomState(52281)
    nmax = 80
    nn = np.linspace(0.0, nmax, 200)
    ax.fill_between(nn, np.sqrt(nn), -np.sqrt(nn), **fsAa)
    ax.axhline(0.0, **lsAm)
    lcs = [colors['red'], colors['orange'],
           colors['yellow'], colors['green']]
    n = np.arange(0.0, nmax+1, 1.0)
    for k in range(12):
        x = random_walk(nmax, 0.5, rng)
        ls = dict(**lsAm)
        ls['color'] = lcs[k%len(lcs)]
        ax.plot(n, x, **ls)
    ax.set_xlabel('Iteration $n$')
    ax.set_ylabel('Position $x_n$')
    ax.set_xlim(0, nmax)
    ax.set_ylim(-20, 20)
    ax.set_yticks(np.arange(-20, 21, 10))
    fig.savefig("randomwalkone.pdf")
    plt.close()