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

# Define the Painlevé I equation as a system of first-order ODEs
def painleve_I(t, y):
    # y[0] is y, y[1] is y'
    return [y[1], 6 * y[0]**2 + t]

# Set up the plotting range
x_min, x_max = -12, 3

# Values of k to plot
k_values = [-0.45143, -0.45142, 0, 1, 1.8518, 1.8519]

# Create the figure
plt.figure(figsize=(10, 6))

for i, k in enumerate(k_values):
    # Initial conditions at x=0
    y0 = [0, k]
    
    t_span_forward = [0, x_max]
    sol_forward = solve_ivp(painleve_I, t_span_forward, y0, method='DOP853', 
                          dense_output=True, rtol=1e-10, atol=1e-10)
    
    t_forward = np.linspace(0, x_max, 1000)
    y_forward = sol_forward.sol(t_forward)[0]
    
    t_span_backward = [0, x_min]
    sol_backward = solve_ivp(painleve_I, t_span_backward, y0, method='DOP853', 
                           dense_output=True, rtol=1e-10, atol=1e-10)
    
    t_backward = np.linspace(0, x_min, 1000)
    y_backward = sol_backward.sol(t_backward)[0]
    
    t = np.concatenate([t_backward[::-1], t_forward[1:]])
    y = np.concatenate([y_backward[::-1], y_forward[1:]])
    plt.plot(t, y, label=f"$w'(0) = {k}$")

y_parabola = np.linspace(-1.5, 1.5, 1000)
x_parabola = -6 * y_parabola**2

plt.plot(x_parabola, y_parabola, 'k-', linewidth=2, label='$6y^2 + x = 0$')

plt.axhline(y=0, color='gray', linestyle='-', alpha=0.3)
plt.axvline(x=0, color='gray', linestyle='-', alpha=0.3)
plt.grid(True, alpha=0.3)
plt.title("Painlevé I: $y'' = 6y^2 + x$")
plt.xlabel('$x$')
plt.ylabel('$y$')
plt.legend()
plt.xlim(x_min, x_max)
plt.ylim(-2, 2)

# Show the plot
plt.tight_layout()
plt.savefig('painleve_I.svg')
plt.show()