import numpy as np
import matplotlib.pyplot as plt
from matplotlib.scale import FuncScale

# f: (t, y) -> R, dy/dt = f
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

f = lambda t, y: t**3 + y**3
threshold = 100
a = 0
b = 2

lims = [
    (0, 1.4, 1.8, 1, threshold),
    (1, 0.4, 0.6, 1, threshold),
    (-1, 0.4, 0.6, -threshold, -1),
]

y0, xm, xM, ym, yM = lims[0]

fig, ax= plt.subplots(1, 2)
for h in [0.1, 0.01, 0.001, 0.0001]:
    ax[0].plot(*euler(f, a, b, y0, h), label=f'$\\Delta t$={h}', linewidth=1)
    T, Y = improvedEuler(f, a, b, y0, h)
    ax[1].plot(T, Y, label=f'$\\Delta t$={h}', linewidth=1)
    print(f"Max x: {T[-1]}")

for axe in ax:
    axe.legend()
    axe.set_xlim(xm, xM)
    axe.set_ylim(ym, yM)
    #axe.set_yscale(FuncScale(axe, (lambda y: np.log(-y), lambda y: -np.exp(y))))

ax[0].set_title("Euler Method")
ax[1].set_title("Improved Euler Method")

#fig.suptitle(r"Numerically solving $y'=y^3 + t^3$ starting at (0, 0)")
plt.show()