import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp

def complex_ode_system(t, y, c):
    """dz/dt = z^2 + c, split into real and imaginary parts."""
    x, y_ = y
    z = x + 1j * y_
    dz = z**2 + c
    return [dz.real, dz.imag]

def escape_time(c, R_max=10.0, t_max=10.0):
    """
    Return escape time for given complex c.
    If no escape within t_max, return t_max (or inf).
    """
    # Initial condition: z(0)=0
    y0 = [0.0, 0.0]
    # Event: |z| > R_max
    def event(t, y):
        return np.hypot(y[0], y[1]) - R_max
    event.terminal = True
    event.direction = 1

    sol = solve_ivp(complex_ode_system, [0, t_max], y0, args=(c,),
                    events=event, rtol=1e-6, atol=1e-9)
    if sol.t_events[0].size > 0:
        return sol.t_events[0][0]   # escape time
    else:
        return t_max                 # no escape (or bounded)

# Grid of c values
real_vals = np.linspace(-2.0, 1.0, 400)
imag_vals = np.linspace(-1.5, 1.5, 400)
T = np.zeros((len(real_vals), len(imag_vals)))

for i, re in enumerate(real_vals):
    for j, im in enumerate(imag_vals):
        c = re + 1j * im
        T[i, j] = escape_time(c)

# Plot escape time as a 2D map (fractal)
plt.figure(figsize=(10, 8))
plt.imshow(T.T, extent=[real_vals[0], real_vals[-1], imag_vals[0], imag_vals[-1]],
           origin='lower', cmap='hot', aspect='auto')
plt.colorbar(label='Escape time')
plt.xlabel('Re(c)')
plt.ylabel('Im(c)')
plt.title('Fractal Escape Map for dz/dt = z² + c')
plt.show()