import numpy as np
import matplotlib.pyplot as plt
from cycler import cycler

from plot_helpers import CO, bold_ticks
from splineA import splineA


def fig7_15():
    # requires splineA.m file
    # solve IVP using RK4 then use spline
    # y' = f(t,y) with y(0) = y0

    M = 8
    r = 10
    t0 = 0
    y0 = 0.1
    tmax = 1

    # exact solution
    nt = 100
    tt = np.linspace(t0, tmax, nt)
    a0 = (1 - y0) / y0
    exact = np.zeros(100)
    for it in range(100):
        exact[it] = 1 / (1 + a0 * np.exp(-r * tt[it]))

    tp = np.linspace(t0, tmax, M + 1)
    h = tp[1] - tp[0]

    y_rk4 = rk4(f, tp, y0, h, M + 1)
    dy0 = f(0, y0)
    dy1 = f(tmax, y_rk4[M])
    yf = splineA(tp, y_rk4, tt, [1, 1], [dy0, dy1])

    plt.close("all")
    # get(gcf)
    fig, ax = plt.subplots(figsize=(8.8, 3.3))
    fig.subplots_adjust(left=0.09, right=0.98, bottom=0.22, top=0.98)
    ax.set_prop_cycle(cycler(color=CO))

    ax.plot(tp, y_rk4, "or", markersize=8, linewidth=2)
    ax.plot(tt, yf, "b", linewidth=1.4)
    ax.plot(tt, exact, "--k", linewidth=1.4)
    ax.legend(
        [" RK4", " Spline", " Exact"],
        loc="lower right",
        prop={"size": 14, "weight": "bold"},
    )
    ax.set_xlabel("t-axis")
    ax.set_ylabel("Solution")
    ax.set_xlim(0, tmax)
    ax.set_ylim(0, 1.05)
    ax.grid(True)
    bold_ticks(ax, 14)
    for spine in ax.spines.values():
        spine.set_visible(True)
    plt.show()

    # plt.savefig('/Users/mark/Desktop/rksp.eps')


# right-hand side of DE
def f(t, y):
    r = 10
    z = r * y * (1 - y)
    return z


# RK4 method
def rk4(f, t, y0, h, n):
    y = y0
    ypoints = np.empty(n)
    ypoints[0] = y0
    for i in range(1, n):
        k1 = h * f(t[i - 1], y)
        k2 = h * f(t[i - 1] + 0.5 * h, y + 0.5 * k1)
        k3 = h * f(t[i - 1] + 0.5 * h, y + 0.5 * k2)
        k4 = h * f(t[i], y + k3)
        yy = y + (k1 + 2 * k2 + 2 * k3 + k4) / 6
        ypoints[i] = yy
        y = yy
    return ypoints


if __name__ == "__main__":
    fig7_15()
