# Re-import libraries and reinitialize code due to reset
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

np.random.seed(0)

# Parameters
T = 1.0  # Total time
N = 2500*4  # Number of steps
dt = T / N  # Time step size
n_particles = 1  # Number of particles for simplicity

# Generate 2D Brownian motion
increments_x = np.random.normal(0, np.sqrt(dt), (N, n_particles))
increments_y = np.random.normal(0, np.sqrt(dt), (N, n_particles))
x = np.cumsum(increments_x, axis=0)
y = np.cumsum(increments_y, axis=0)

# Time points
t = np.linspace(0, T, N)

# Prepare the figure
fig, ax = plt.subplots(figsize=(6, 6))
ax.set_facecolor('black')
ax.set_xlim(x.min(), x.max())
ax.set_ylim(y.min(), y.max())
#ax.set_title("2D Brownian Motion")
# ax.set_xlabel("X")
# ax.set_ylabel("Y")

# Initialize the plot elements
line, = ax.plot([], [], 'w', lw=1.0, label="Path")
point, = ax.plot([], [], 'ro', label="Current Position")
ax.legend()

# Update function for animation
def update(frame):
    line.set_data(x[:4*frame, 0], y[:4*frame, 0])  # Update path
    point.set_data(x[4*frame-1, 0], y[4*frame-1, 0])  # Update current position
    return line, point

# Create the animation
ani = FuncAnimation(fig, update, frames=N//4, interval=1, blit=True)

if True:
    # save as mp4
    ani.save("brownian-motion-2d-time.mp4", writer="ffmpeg", fps=60)

# Display the animation
plt.show()