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

"""
Shows a swarm of particles with numpy and matplotlib.
"""

import numpy as np
import random
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.
    """
    x = np.zeros(npa)
    y = np.zeros(npa)
    xymax = int(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(x, y,'bo')
    strtitle = 'animating %d particles' % npa
    ax.set_title(strtitle)
    fig.canvas.draw()
    plt.pause(0.00001)
    for step in range(nst):
        for i in range(npa):
            d = random.randint(1,4)
            if d == 1: y[i] += 1   # north
            elif d == 2: y[i] -= 1 # south
            elif d == 3: x[i] += 1 # east
            elif d == 4: x[i] -= 1 # west
        if(step+1) % pls == 0:
            dots.set_xdata(x)
            dots.set_ydata(y)
            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.
    """
    random.seed(10)
    particles(3000, 4000, 20)

main()