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

from plot_helpers import CO, bold_ticks


def fig7_11():
    # using RK4 on the system y' = f(t,y) with y(0) = y0
    # this also contains Figure 7.12

    t0 = 0
    y0 = np.array([np.pi / 4, 0], dtype=float)
    tmax = 4000
    # k = 1/3
    k = 1 / 2
    # k = 1/200
    n = round(tmax / k)
    print(n)
    t = np.linspace(t0, tmax, n)

    # solve using rk4
    y, v = rk4(t, y0, k, n)

    # determine the number of zero crossings
    # k=1/200 => 1000 crossing for 0<t<3266
    # k=1/2 => 1016 crossing for 0<t<3266
    # nzeros=0
    # for it in range(1,n):
    #     if y[it]*y[it-1] < 0:
    #         nzeros=nzeros+1
    # nzeros_RK4=nzeros

    fig1, ax1 = plt.subplots(figsize=(8.8, 3.3))
    fig1.subplots_adjust(left=0.09, right=0.98, bottom=0.22, top=0.98)
    ax1.set_prop_cycle(cycler(color=CO))

    ax1.plot(t, y, "-b", linewidth=2)
    T1 = 40
    ax1.set_xlim(0, T1)
    ax1.set_ylim(-0.85, 0.85)
    ax1.plot([0, T1], [np.pi / 4, np.pi / 4], "--r", linewidth=2)
    ax1.set_yticks([-np.pi / 4, 0, np.pi / 4])
    ax1.set_yticklabels([" ", "0", " "])
    say = r"$-\pi/4$"
    ax1.text(-3, -np.pi / 4, say, fontsize=18)
    say = r"$\pi/4$"
    ax1.text(-2.6, np.pi / 4, say, fontsize=18)
    ax1.set_xticks([0, 10, 20, 30, 40])
    ax1.grid(True)
    ax1.set_xlabel("t-axis")
    # ax1.set_ylabel(r"$\theta$-axis")
    ax1.set_ylabel("Theta")
    bold_ticks(ax1, 14)
    for spine in ax1.spines.values():
        spine.set_visible(True)

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

    fig2, ax2 = plt.subplots(figsize=(8.8, 3.3))
    fig2.subplots_adjust(left=0.08, right=0.995, bottom=0.22, top=0.98)
    ax2.set_prop_cycle(cycler(color=CO))

    tr = np.linspace(0, tmax, 100)
    aa = 1e-4 * np.log(7 * np.pi / 3)
    R = np.zeros(100)
    for ir in range(100):
        R[ir] = 0.25 * np.pi * np.exp(-aa * tr[ir])

    ax2.plot(t, y, "-b", linewidth=0.01, color="#00BFFF")
    # ax2.plot(tr, R, "r", linewidth=2.5)
    ax2.plot([0, tmax], [np.pi / 4, np.pi / 4], "--r", linewidth=2)
    ax2.set_yticks([-np.pi / 4, 0, np.pi / 4])
    # ax2.set_yticklabels(["-pi/4", "0", "pi/4"])
    ax2.set_yticklabels([" ", "0", " "])
    say = r"$-\pi/4$"
    ax2.text(-300, -np.pi / 4, say, fontsize=18)
    say = r"$\pi/4$"
    ax2.text(-250, np.pi / 4, say, fontsize=18)
    ax2.grid(True)
    ax2.set_xlabel("t-axis")
    ax2.set_ylabel("Theta")
    bold_ticks(ax2, 14)
    for spine in ax2.spines.values():
        spine.set_visible(True)

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


# right-hand side of DE
def f(t, y):
    z = np.array([y[1], -np.sin(y[0])])
    return z


# RK4 method
def rk4(t, y0, h, n):
    y = np.asarray(y0, dtype=float)
    ypoints = np.empty(n)
    vpoints = np.empty(n)
    ypoints[0] = y0[0]
    vpoints[0] = y0[1]
    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[0]
        vpoints[i] = yy[1]
        y = yy
    return ypoints, vpoints


if __name__ == "__main__":
    fig7_11()
