import numpy as np
import matplotlib.pyplot as plt

threshold = 100

def RK4(y0, t):
    ODE = lambda t, y: t**3 + y**3
    #note: t is a vector of time points, y0 is the initial value
    steps = len(t)
    y = np.zeros(steps)
    h = (t[-1] - t[0]) / (steps - 1)
    if t[0] > 0:
        raise ValueError("t[0] must be less than or equal to 0")
    
    #find the index of the first non-negative time point
    t0_idx = np.where(t >= 0)[0][0]
    y[t0_idx] = y0  # set the initial condition at t[0]
    
    for i in range(t0_idx, steps - 1):
        k1 = h * ODE(t[i], y[i])
        k2 = h * ODE(t[i] + h / 2, y[i] + k1 / 2)
        k3 = h * ODE(t[i] + h / 2, y[i] + k2 / 2)
        k4 = h * ODE(t[i] + h, y[i] + k3)
        y[i + 1] = y[i] + (k1 + 2 * k2 + 2 * k3 + k4) / 6
    return y    

def euler(f, a, b, y0, h):
    xs = [a]
    ys = [y0]
    x = a
    y = y0
    while x < b:
        y += h * f(x, y)
        x += h
        if abs(y) > threshold:
            break
        xs.append(x)
        ys.append(y)
    return xs, ys

def improvedEuler(f, a, b, y0, h):
    xs = [a]
    ys = [y0]
    
    x = a
    y = y0
    while x <= b:
        slope1 = f(x, y)
        y_predictor = y + h * slope1
        slope2 = f(x + h, y_predictor)
        y = y + (h/2) * (slope1 + slope2)
        x = x + h
        if abs(y) > threshold:
            break
        xs.append(x)
        ys.append(y)
    return xs, ys

lims = [
    (0, 1.55, 1.75, 1, threshold),
    (1, 0.4, 0.6, 1, threshold),
    (-1, 0.4, 0.6, -threshold, -1),
]

ODE = lambda t, y: t**3 + y**3
fig, axes = plt.subplots(3, 3)

col_titles = ["Euler's Method", "Improved Euler's", "Runge-Kutta"]
row_titles = ["y(0) = 0", "y(0) = 1", "y(0) = -1"]

# Add column titles
for col_idx, col_title in enumerate(col_titles):
    axes[0][col_idx].set_title(col_title, fontsize=12)

for row_idx, row_title in enumerate(row_titles):
    axes[row_idx][0].set_ylabel(row_title, fontsize=12, rotation=90, labelpad=20)

for idx, (y0, xm, xM, ym, yM) in enumerate(lims):
    for h in [0.1, 0.01, 0.001, 0.0001]:
        Te, Ye = euler(ODE, 0, xM, y0, h)
        Ti, Yi = improvedEuler(ODE, 0, xM, y0, h)
        Tr = np.linspace(0, xM, int((xM - 0)/h))
        Yr = RK4(y0, Tr)
        
        axes[idx][0].plot(Te, Ye, label=f'$\Delta t$={h}', linewidth=1)
        axes[idx][1].plot(Ti, Yi, label=f'$\Delta t$={h}', linewidth=1)
        axes[idx][2].plot(Tr, Yr, label=f'$\Delta t$={h}', linewidth=1)
    for idy in range(3):
        axes[idx][idy].set_xlim(xm, xM)
        axes[idx][idy].set_ylim(ym, yM)
        axes[idx][idy].legend()

fig.suptitle("Numeric Solutions for Each ODE Near the Asymptote")
plt.show()