import numpy as np
import matplotlib.pyplot as plt


def RK_n(y0, t, n=4):
    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]
    
    #forward RK
    for i in range(t0_idx, steps - 1):
        #n-th order Runge-Kutta method
        k = np.zeros(n)
        k[0] = h * ODE(t[i], y[i])
        for j in range(1, n):
            t_j = t[i] + h * (j / n)
            y_j = y[i] + k[j - 1] / n
            k[j] = h * ODE(t_j, y_j)
        y[i + 1] = y[i] + np.sum(k) / n
        
    #backward RK
    for i in range(t0_idx - 1, -1, -1):
        #n-th order Runge-Kutta method
        k = np.zeros(n)
        k[0] = h * ODE(t[i], y[i])
        for j in range(1, n):
            t_j = t[i] + h * (j / n)
            y_j = y[i] + k[j - 1] / n
            k[j] = h * ODE(t_j, y_j)
        y[i] = y[i + 1] - np.sum(k) / n
    return y
    
case_num = 3
if case_num == 1:
    plt.title("RK4 result with IVPs for large t")
    initial_values = [-1, 0, 1]
    t = np.linspace(0, 2, 1000)
    for i, y0 in enumerate(initial_values):
        res = RK_n(y0, t, 4)
        plt.plot(t, res, label=f'y(0)={y0}, n=4')
    plt.ylim(-5,5)
    #observation: the solutions diverge as t increases and seem to have an asymptote which depends on y(0)
    # and converge to a same solution as t decreases.
    
elif case_num == 2:
    plt.title("RK4 result with IVPs for small t")
    initial_values = [-1, 0, 1]
    t = np.linspace(-0.5, 0.5, 1000)
    for i, y0 in enumerate(initial_values):
        res = RK_n(y0, t, 4)
        plt.plot(t, res, label=f'y(0)={y0}, n=4')
    plt.ylim(-5,5)

elif case_num == 3:
    plt.title("Comparing different order RK methods")
    order = [1, 2, 4, 5]
    t = np.linspace(-1, 1.4, 1000)
    res_5 = RK_n(0, t, 5)
    for n in order:
        res = RK_n(0, t, n)
        if n == 1:
            label = 'Euler - RK5'
        elif n == 2:
            label = 'Improved Euler - RK5'
        else:
            label = f'RK{n} - RK5'
        plt.plot(t, res - res_5, label=label)
    plt.legend()
    
elif case_num == 4:
    plt.title("RK4 with IVP from 0 to 1")
    initial_values = np.linspace(0, 1, 5)
    t = np.linspace(-3, 10, 1000)
    for i, y0 in enumerate(initial_values):
        res = RK_n(y0, t, 4)
        plt.plot(t, res, label=f'y(0)={y0}')
    plt.ylim(-5,10)
    plt.legend()
    
elif case_num == 5:
    fig, axs = plt.subplots(2, 2, figsize=(10, 8))
    plt.title("RK4 with different y-scale in the negatives")
    scale = [5, 10, 20, 50]
    t = np.linspace(-10, 1, 1000)
    res = RK_n(0, t, 4)
    for i, s in enumerate(scale):
        ax = axs[i // 2, i % 2]
        ax.plot(t, res / s, label=f'scale={s}')
        ax.set_ylim(-1, scale[i])
        ax.axhline(0, color='black', lw=0.5)
        ax.axvline(0, color='black', lw=0.5)
        ax.set_title(f'RK4 for t < {s}')
    # observation: There is no asymptote in the negative y-scale
    
elif case_num == 6:
    plt.title("RK4 with different step sizes")
    steps = [5, 10, 100, 1000]
    for n in steps:
        t = np.linspace(0, 1, n)
        res = RK_n(0, t, 4)
        plt.plot(t, res, label=f'{n} steps')
    plt.legend()
    

plt.axhline(0, color='black', lw=0.5)
plt.axvline(0, color='black', lw=0.5)
plt.show()
# plt.legend()