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

from plot_helpers import CO, bold_ticks


icase = 1


def fig7_9(M=None):
    # comparison of the RK4 and other RK4s

    global icase

    if M is None:
        M = np.array([10, 100, 1000, 10000], dtype=int)
    else:
        M = np.asarray(M, dtype=int)

    plt.close("all")
    # get(gcf)
    fig, axes = plt.subplots(2, 1, figsize=(8.8, 5.6))
    fig.subplots_adjust(left=0.10, right=0.98, bottom=0.12, top=0.99, hspace=0.18)
    for ax in axes:
        ax.set_prop_cycle(cycler(color=CO))

    # logistic eq
    icase = 1
    tmax = 1
    y0 = 0.1
    a0 = (1 - y0) / y0
    r = 10
    # exact = 1/(1+a0*np.exp(-r*t))

    points = np.zeros(len(M), dtype=int)
    e_trap = np.zeros(len(M))
    e_rk4L = np.zeros(len(M))
    e_rk4 = np.zeros(len(M))
    e_rk4R = np.zeros(len(M))

    for im in range(len(M)):
        m = int(M[im])
        points[im] = m
        t = np.linspace(0, tmax, m + 1)
        k = t[1] - t[0]

        exact = np.zeros(m + 1)
        for i in range(m + 1):
            exact[i] = 1 / (1 + a0 * np.exp(-r * t[i]))

        # trap method
        y_trap = traps(t, y0, k, m + 1)
        e_trap[im] = np.linalg.norm(exact - y_trap, np.inf)

        # Labatto RK4
        y_rk4L = rk4L(t, y0, k, m + 1)
        e_rk4L[im] = np.linalg.norm(exact - y_rk4L, np.inf)

        # RK4
        y_rk4 = rk4(t, y0, k, m + 1)
        e_rk4[im] = np.linalg.norm(exact - y_rk4, np.inf)

        # RK4 Ralston min error
        y_rk4R = rk4R(t, y0, k, m + 1)
        e_rk4R[im] = np.linalg.norm(exact - y_rk4R, np.inf)

    ax = axes[0]
    ax.loglog(points, e_rk4, "--bo", markersize=8, linewidth=1.4)
    ax.loglog(points, e_rk4L, "--r*", markersize=8, linewidth=1.4)
    ax.loglog(points, e_trap, "--kd", markersize=8, linewidth=1.4)
    ax.grid(True)
    # ax.set_xlim(10, 1e3); ax.set_ylim(1e-8, 1)
    ax.set_yticks([1e-15, 1e-12, 1e-9, 1e-6, 1e-3, 1])
    ax.legend(
        [" RK4-C", " RK4-L", " Trap"],
        loc="lower left",
        prop={"size": 14, "weight": "bold"},
    )

    # ax.loglog(points, e_rk4R, "--gs", markersize=8, linewidth=1.4)

    # ax.set_xlabel("Number of Time Points")
    ax.set_ylabel("Error")
    bold_ticks(ax, 14)

    # forced problem
    icase = 2
    y0 = 1
    tmax = 3
    z = 5
    # exact = np.exp(t)*np.cos(z*t**2)

    e_trap2 = np.zeros(len(M))
    e_rk4L2 = np.zeros(len(M))
    e_rk42 = np.zeros(len(M))
    e_rk4R2 = np.zeros(len(M))

    for im in range(len(M)):
        m = int(M[im])
        points[im] = m
        t = np.linspace(0, tmax, m + 1)
        k = t[1] - t[0]

        exact2 = np.zeros(m + 1)
        for i in range(m + 1):
            exact2[i] = np.exp(t[i]) * np.cos(z * t[i] ** 2)

        # trap method
        y_trap = traps(t, y0, k, m + 1)
        e_trap2[im] = np.linalg.norm(exact2 - y_trap, np.inf)

        # Labatto RK4
        y_rk4L2 = rk4L(t, y0, k, m + 1)
        e_rk4L2[im] = np.linalg.norm(exact2 - y_rk4L2, np.inf)

        # RK4
        y_rk42 = rk4(t, y0, k, m + 1)
        e_rk42[im] = np.linalg.norm(exact2 - y_rk42, np.inf)

        # RK4 Ralston min error
        y_rk4R2 = rk4R(t, y0, k, m + 1)
        e_rk4R2[im] = np.linalg.norm(exact2 - y_rk4R2, np.inf)

    ax = axes[1]
    ax.loglog(points, e_rk42, "--bo", markersize=8, linewidth=1.4)
    ax.loglog(points, e_rk4L2, "--r*", markersize=8, linewidth=1.4)
    ax.loglog(points, e_trap2, "--kd", markersize=8, linewidth=1.4)
    ax.grid(True)
    # ax.set_xlim(10, 1e3); ax.set_ylim(1e-8, 1)
    ax.set_yticks([1e-15, 1e-12, 1e-9, 1e-6, 1e-3, 1])
    ax.legend(
        [" RK4-C", " RK4-L", " Trap"],
        loc="lower left",
        prop={"size": 14, "weight": "bold"},
    )

    # ax.loglog(points, e_rk4R2, "--gs", markersize=8, linewidth=1.4)

    ax.set_xlabel("Number of Time Points")
    ax.set_ylabel("Error")
    bold_ticks(ax, 14)

    plt.show()

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


def f(t, y):
    global icase
    if icase == 1:
        g = 10 * y * (1 - y)
    else:
        z = 5
        g = y - 2 * z * t * np.exp(t) * np.sin(z * t ** 2)
    return g


# Labatto RK4
def rk4L(t, y0, k, n):
    a2 = (1 - 1 / np.sqrt(5)) / 2
    a3 = (1 + 1 / np.sqrt(5)) / 2
    b21 = a2
    c1 = 1 / 12
    c2 = 5 / 12
    c3 = 5 / 12
    c4 = 1 / 12
    b31 = -1 / 4 - (3 / 20) * np.sqrt(5)
    b32 = (1 / 4) * np.sqrt(5) + 3 / 4
    b41 = -1 / 4 + (5 / 4) * np.sqrt(5)
    b42 = -5 / 4 - (3 / 4) * np.sqrt(5)
    b43 = 5 / 2 - (1 / 2) * np.sqrt(5)
    ypoints = np.empty(n)
    ypoints[0] = y0
    y = y0
    for i in range(1, n):
        k1 = k * f(t[i - 1], y)
        k2 = k * f(t[i - 1] + a2 * k, y + b21 * k1)
        k3 = k * f(t[i - 1] + a3 * k, y + b31 * k1 + b32 * k2)
        k4 = k * f(t[i], y + b41 * k1 + b42 * k2 + b43 * k3)
        yy = y + c1 * k1 + c2 * k2 + c3 * k3 + c4 * k4
        ypoints[i] = yy
        y = yy
    return ypoints


# RK4 method
def rk4(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


# trap method
def traps(t, y0, k, n):
    tol = 5 * np.finfo(float).eps
    y = y0
    ypoints = np.empty(n)
    ypoints[0] = y0
    fold = y + 0.5 * k * f(t[0], y)
    yc = y0
    yb = 1.0001 * y0
    for i in range(1, n):
        # secant method
        ya = 1.01 * yc
        fa = ya - 0.5 * k * f(t[i], ya) - fold
        yb = yc
        fb = yb - 0.5 * k * f(t[i], yb) - fold
        err = 1.0
        iterations = 0
        while err > tol:
            denominator = fb - fa
            if denominator == 0 or iterations >= 100:
                break
            yc = yb - fb * (yb - ya) / denominator
            fc = yc - 0.5 * k * f(t[i], yc) - fold
            err = abs(yc - yb)
            ya = yb
            fa = fb
            yb = yc
            fb = fc
            iterations += 1
        y = yc
        fold = y + 0.5 * k * f(t[i], y)
        ypoints[i] = y
    return ypoints


# RK4 min error (ralston)
def rk4R(t, y0, h, n):
    y = y0
    ypoints = np.empty(n)
    ypoints[0] = y0
    a2 = 2 / 5
    a3 = 7 / 8 - 3 * np.sqrt(5) / 16
    b21 = a2
    c1 = 1 / 2 + (1 - 2 * (a2 + a3)) / (12 * a2 * a3)
    c2 = (2 * a3 - 1) / (12 * a2 * (a3 - a2) * (1 - a2))
    c3 = (1 - 2 * a2) / (12 * a3 * (a3 - a2) * (1 - a3))
    c4 = 1 / 2 + (2 * (a2 + a3) - 3) / (12 * (1 - a2) * (1 - a3))
    b32 = a3 * (a3 - a2) / (2 * a2 * (1 - 2 * a2))
    b42 = (
        (1 - a2)
        * (a2 + a3 - 1 - (2 * a3 - 1) ** 2)
        / (2 * a2 * (a3 - a2) * (6 * a2 * a3 - 4 * (a2 + a3) + 3))
    )
    b43 = (
        (1 - 2 * a2)
        * (1 - a2)
        * (1 - a3)
        / (a3 * (a3 - a2) * (6 * a2 * a3 - 4 * (a2 + a3) + 3))
    )
    b31 = a3 - b32
    b41 = 1 - b42 - b43
    for i in range(1, n):
        k1 = h * f(t[i - 1], y)
        k2 = h * f(t[i - 1] + a2 * h, y + b21 * k1)
        k3 = h * f(t[i - 1] + a3 * h, y + b31 * k1 + b32 * k2)
        k4 = h * f(t[i], y + b41 * k1 + b42 * k2 + b43 * k3)
        yy = y + c1 * k1 + c2 * k2 + c3 * k3 + c4 * k4
        ypoints[i] = y  # literal conversion of the original MATLAB line
        y = yy
    return ypoints


if __name__ == "__main__":
    fig7_9()
