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

from plot_helpers import CO, bold_ticks


def fig7_6():
    # comparison of trap, euler and b euler
    # y' = ry(1-y) with y(0) = y0

    M = 30
    r = 10
    t0 = 0
    y0 = 0.01
    tmax = 3
    t = np.linspace(t0, tmax, M + 1)

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

    # euler
    h = t[1] - t[0]
    hr = h * r
    y = y0
    euler = [y0]
    for i in range(1, M + 1):
        yy = y + hr * y * (1 - y)
        euler.append(yy)
        y = yy
    euler = np.asarray(euler)

    # backward euler
    y = y0
    beuler = [y0]
    for i in range(1, M + 1):
        yy = -0.5 * (1 - hr - np.sqrt((1 - hr) ** 2 + 4 * hr * y)) / hr
        beuler.append(yy)
        y = yy
    beuler = np.asarray(beuler)

    # trapezoidal
    y = y0
    trap = [y0]
    for i in range(1, M + 1):
        b = 2 * hr * y * (1 + 0.5 * hr * (1 - y))
        yy = (-1 + 0.5 * hr + np.sqrt((1 - 0.5 * hr) ** 2 + b)) / hr
        trap.append(yy)
        y = yy
    trap = np.asarray(trap)

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

    ax = axes[0]
    ax.grid(True)
    ax.set_xlim(0, tmax)
    ax.set_ylim(0, 1.1)
    ax.plot(tt, exact, "k", linewidth=1.2)
    ax.plot(t, euler, "--bo", markersize=8, linewidth=1.6)
    ax.plot(t, beuler, "--d", color=[0, 0.5, 0], markersize=8, linewidth=1.6)
    ax.plot(t, trap, "--r*", markersize=8, linewidth=1.6)
    ax.set_yticks([0, 0.5, 1])
    ax.set_xticks([0, 1, 2, 3])
    ax.legend(
        [" Exact", " Euler", " B Euler", " Trap"],
        loc="lower right",
        prop={"size": 14, "weight": "bold"},
    )
    # ax.set_xlabel("t-axis")
    ax.set_ylabel("Solution")
    bold_ticks(ax, 14)
    for spine in ax.spines.values():
        spine.set_visible(True)

    ax = axes[1]
    M = 30
    r = 10
    t = np.linspace(0, tmax, M + 1)

    # leapfrog solution
    h = t[1] - t[0]
    hr = h * r
    y = y0
    yy = 1 / (1 + a0 * np.exp(-hr))
    leap = [y, yy]
    for i in range(2, M + 1):
        yyy = y + 2 * hr * yy * (1 - yy)
        leap.append(yyy)
        y = yy
        yy = yyy
    leap = np.asarray(leap)

    ax.plot(tt, exact, "k", linewidth=1.2)
    ax.grid(True)
    ax.plot(t, leap, "--*b", linewidth=1.6)

    M = 150
    r = 10
    t = np.linspace(0, tmax, M + 1)

    # leapfrog solution
    h = t[1] - t[0]
    hr = h * r
    y = y0
    yy = 1 / (1 + a0 * np.exp(-hr))
    leap = [y, yy]
    for i in range(2, M + 1):
        yyy = y + 2 * hr * yy * (1 - yy)
        leap.append(yyy)
        y = yy
        yy = yyy
    leap = np.asarray(leap)
    ax.plot(t, leap, "-r", linewidth=1.6)

    ax.legend(
        [" Exact", " Leap (M=30)", " Leap (M=150)"],
        loc="upper left",
        prop={"size": 14, "weight": "bold"},
    )
    ax.set_yticks([-0.5, 0, 0.5, 1, 1.5, 2])
    ax.set_xlim(0, tmax)
    ax.set_ylim(-0.5, 2)
    ax.set_xticks([0, 1, 2, 3])
    ax.set_xlabel("t-axis")
    ax.set_ylabel("Solution")
    bold_ticks(ax, 14)

    plt.show()

    # plt.savefig('/Users/holmes/Desktop/etal.eps')


if __name__ == "__main__":
    fig7_6()
