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

def random_walk(n, rng):
    p = 0.25
    steps = rng.rand(n)
    xsteps = np.zeros(len(steps))
    ysteps = np.zeros(len(steps))
    xsteps[(steps>=0*p)&(steps<1*p)] = +1
    xsteps[(steps>=1*p)&(steps<2*p)] = -1
    ysteps[(steps>=2*p)&(steps<3*p)] = +1
    ysteps[(steps>=3*p)&(steps<4*p)] = -1
    x = np.hstack(((0.0,), np.cumsum(xsteps)))
    y = np.hstack(((0.0,), np.cumsum(ysteps)))
    return x, y

if __name__ == "__main__":
    fig, ax = plt.subplots(figsize=cm_size(figure_width, 0.65*figure_width))
    fig.subplots_adjust(**adjust_fs(fig, right=1.0))
    rng = np.random.RandomState(42281)
    nmax = 500
    lcs = [colors['blue'], colors['red'], colors['orange'],
           colors['yellow'], colors['green']]
    for k in range(len(lcs)):
        x, y = random_walk(nmax, rng)
        ls = dict(**lsAm)
        ls['color'] = lcs[k%len(lcs)]
        ax.plot(x, y, **ls)
    ax.plot([0], [0], **psMarker)
    ax.set_xlabel('Position $x_n$')
    ax.set_ylabel('Position $y_n$')
    ax.set_xlim(-30, 30)
    ax.set_ylim(-20, 20)
    ax.set_xticks(np.arange(-30, 31, 10))
    ax.set_yticks(np.arange(-20, 21, 10))
    fig.savefig("randomwalktwo.pdf")
    plt.close()