# L-4 MCS 507 Mon 28 Aug 2023 : walkvector.py

"""
This is the vectorized version of walk.py.
"""

import numpy as np
import matplotlib
import matplotlib.pyplot as plt

def particles(npa, nst, pls):
    """
    Shows random particle movement with
    npa : number of particles,
    nst : number of time steps,
    pls : how many steps for next plot.
    """
    xpa = np.zeros(npa)
    ypa = np.zeros(npa)
    xymax = 3*np.sqrt(nst)
    xymin = -xymax
    plt.ion()
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.set_xlim(xymin, xymax)
    ax.set_ylim(xymin, xymax)
    dots, = ax.plot(xpa, ypa, 'bo')
    strtitle = 'animating %d particles' % npa
    ax.set_title(strtitle)
    fig.canvas.draw()
    plt.pause(0.00001)
    moves = np.random.randint(1, 5, nst*npa)
    moves.shape = (nst, npa)
    for step in range(nst):
        this_move = moves[step, :]
        ypa += np.where(this_move == 1, 1, 0)
        ypa -= np.where(this_move == 2, 1, 0)
        xpa += np.where(this_move == 3, 1, 0)
        xpa -= np.where(this_move == 4, 1, 0)
        if(step+1) % pls == 0:
            dots.set_xdata(xpa)
            dots.set_ydata(ypa)
            strtitle = '%d particles after %d steps' % (npa, step+1)
            ax.set_title(strtitle)
            fig.canvas.draw()
            plt.pause(0.01)

def main():
    """
    Fixes the seed for the random numbers
    and starts the particle simulation.
    """
    np.random.seed(10)
    particles(3000, 4000, 20)

main()