import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# Parameters
steps = 500
step_size = 0.1
#np.random.seed(42)

# Generate random walks for real and imaginary parts
real_walk = np.cumsum(np.random.normal(0, step_size, steps))
imag_walk = np.cumsum(np.random.normal(0, step_size, steps))

# Complex path Z(t)
Z = real_walk + 1j * imag_walk

# Reciprocal (singularity when Z -> 0)
reciprocal = 1 / (Z + 1e-12)  # avoid division by zero

# Animation setup
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

# Left: complex plane walk
ax1.plot(real_walk, imag_walk, 'b-', alpha=0.5, linewidth=1)
point1, = ax1.plot([], [], 'ro', markersize=6)
ax1.set_xlabel('Real')
ax1.set_ylabel('Imag')
ax1.set_title('Random Walk Z(t) in Complex Plane')
ax1.grid(True)

# Right: magnitude of reciprocal over time
ax2.set_xlim(0, steps)
ax2.set_ylim(0, 10)
line2, = ax2.plot([], [], 'r-', linewidth=1)
scat2, = ax2.plot([], [], 'bo', markersize=4)
ax2.set_xlabel('Step t')
ax2.set_ylabel('|1/Z(t)|')
ax2.set_title('Singularity (blow-up) when Z(t) near 0')
ax2.grid(True)

# Precompute the magnitude of reciprocal for all steps
mag_recip = np.abs(reciprocal)

def update(frame):
    # Update left plot: current position
    point1.set_data([real_walk[frame]], [imag_walk[frame]])
    # Update right plot: trace of |1/Z| up to current frame
    line2.set_data(np.arange(frame+1), mag_recip[:frame+1])
    scat2.set_data([frame], [mag_recip[frame]])
    return point1, line2, scat2

ani = FuncAnimation(fig, update, frames=steps, interval=20, blit=True)
plt.tight_layout()
plt.show()

# Optionally save the animation:
# ani.save('singularity_random_walk.gif', writer='pillow', fps=50)

# Print where the singularity is most severe
min_idx = np.argmin(np.abs(Z))
print(f"Closest approach to origin at step {min_idx}: |Z| = {np.abs(Z[min_idx]):.4f}")
print(f"Maximum |1/Z| = {np.max(mag_recip):.2f}")