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

from plot_helpers import CO, bold_ticks


def fig7_7(M_values=None):
    # comparison of various DE solvers
    # y' = f(t,y) with y(0) = y0

    # logistic equation
    t0 = 0
    tmax = 3
    y0 = 0.01
    r = 10
    a0 = (1 - y0) / y0
    exact = 1 / (1 + a0 * np.exp(-r * tmax))

    # example 2
    # t0 = 0
    # tmax = 3
    # y0 = 1
    # exact = np.exp(tmax)*np.cos(tmax**2)

    if M_values is None:
        M_values = [20, 100, 1000, 10000, 100000]

    points = []
    e_euler = []
    e_beuler = []
    e_trap = []
    e_rk2H = []
    e_rk2E = []

    for m in M_values:
        points.append(m)
        t = np.linspace(t0, 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]))

        # euler
        y_euler = eulers(t, y0, k, m + 1)
        e_euler.append(np.linalg.norm(exact - y_euler, np.inf))

        # backward euler
        y_beuler = beulers(t, y0, k, m + 1)
        e_beuler.append(np.linalg.norm(exact - y_beuler, np.inf))

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

        # RK2: Heun
        y_rk2H = rk2H(t, y0, k, m + 1)
        e_rk2H.append(np.linalg.norm(exact - y_rk2H, np.inf))

        # RK2: min error
        y_rk2E = rk2E(t, y0, k, m + 1)
        e_rk2E.append(np.linalg.norm(exact - y_rk2E, np.inf))

    points = np.asarray(points)
    e_euler = np.asarray(e_euler)
    e_beuler = np.asarray(e_beuler)
    e_trap = np.asarray(e_trap)

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

    ax.loglog(points, e_euler, "--bo", markersize=8, linewidth=1.6)
    ax.set_yticks([1e-10, 1e-5, 1])
    ax.set_xlim(10, 1e5)
    ax.set_ylim(7e-11, 1)
    ax.loglog(
        points,
        e_beuler,
        "--d",
        color=[0, 0.5, 0],
        markersize=8,
        linewidth=1.6,
    )
    ax.grid(True)
    ax.grid(False, which="minor")
    ax.loglog(points, e_trap, "--r*", markersize=8, linewidth=1.6)
    ax.legend(
        [" Euler", " B Euler", " Trap"],
        loc="lower left",
        prop={"size": 14, "weight": "bold"},
    )

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

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


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


# euler method
def eulers(t, y0, k, n):
    y = y0
    ypoints = np.empty(n)
    ypoints[0] = y0
    for i in range(1, n):
        yy = y + k * f(t[i - 1], y)
        ypoints[i] = yy
        y = yy
    return ypoints


# backward euler method
def beulers(t, y0, h, n):
    tol = 5 * np.finfo(float).eps
    y = y0
    fold = f(t[0], y)
    ypoints = np.empty(n)
    ypoints[0] = y0
    for i in range(1, n):
        # secant method
        yb = y
        fb = yb - h * f(t[i], yb) - y
        yc = y + 2 * h * (fold + f(t[i], y + h * fold))
        fc = yc - h * f(t[i], yc) - y
        iterations = 0
        while abs(yc - yb) > tol:
            ya = yb
            fa = fb
            yb = yc
            fb = fc
            denominator = fb - fa
            if denominator == 0 or iterations >= 100:
                break
            yc = yb - fb * (yb - ya) / denominator
            fc = yc - h * f(t[i], yc) - y
            iterations += 1
        y = yc
        fold = f(t[i], y)
        ypoints[i] = y
    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


# RK2: Heun
def rk2H(t, y0, h, n):
    b = 0.5
    a = 1 - b
    alpha = 0.5 / b
    beta = alpha
    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] + alpha * h, y + beta * k1)
        yy = y + a * k1 + b * k2
        ypoints[i] = yy
        y = yy
    return ypoints


# RK2: min error
def rk2E(t, y0, h, n):
    b = 3 / 4
    a = 1 - b
    alpha = 2 / 3
    beta = alpha
    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] + alpha * h, y + beta * k1)
        yy = y + a * k1 + b * k2
        ypoints[i] = yy
        y = yy
    return ypoints


if __name__ == "__main__":
    fig7_7()
