import numpy as np
import matplotlib.pyplot as plt

# import random_walk_2d as rw

""" Settings: """
# Mesh sizes:
dx = 0.4
dt = 0.04
# Population size:
N = 1000
# Initial conditions:
x0 = np.zeros((2, N))
t0 = 0.0
# Number of discrete time steps:
K = 20
# Axes limits for figure:
xmin = -8.0
xmax = 8.0

""" Function to update the figure: """
def _update_figure(fig, X, t, xmin, xmax):
    fig.clear()
    plt.scatter(X[:, 0], X[:, 1])
    plt.title("Population at time t = %.3f"%t)
    plt.xlim([xmin, xmax])
    plt.ylim([xmin, xmax])

    return fig

# """ Instantiate the class: """
# RW = rw.RandomWalk2D(dx, dt, N, x0)

# """ Simulation: """
# X = RW.positions()
# fig = plt.figure(figsize=(6, 4))
# fig = _update_figure(fig, X, 0.0, xmin, xmax)

# for ii in range(1, K+1):
#     # Update positions:
#     RW.simulate()
#     X = RW.positions()
#     t = RW.time()
#     # Update figure:
#     fig = _update_figure(fig, X, t, xmin, xmax)
#     plt.pause(0.25)

# plt.show()