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

from plot_helpers import CO, arrowhead, bold_ticks


def fig7_10():
    # RK4 for the system y' = f(t,y) with y(0) = y0
    # y=[y(1) y(2)]   f=[f1 f2]

    M = 400

    T = 10
    t = np.linspace(0, T, M + 1)
    k = t[1] - t[0]

    # solve using rk4
    y0 = np.array([-1, -6], dtype=float)
    u, v = rk4(t, y0, k, M + 1)

    # solve using rk4
    y0 = np.array([3, 2], dtype=float)
    u2, v2 = rk4(t, y0, k, M + 1)

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

    ax.plot(u, v, "-b", linewidth=1.5)
    ax.plot(u[0], v[0], "b.", markersize=40)
    ax.plot(u2, v2, "-r", linewidth=1.5)
    ax.plot(u2[0], v2[0], ".r", markersize=40)

    ar = 1
    ii = 39
    i = ii - 1
    arrowhead(ax, [u[i], u[ii]], [v[i], v[ii]], "b", scale=14 * ar, linewidth=3)

    ar = 1
    ii = 29
    i = ii - 1
    arrowhead(ax, [u2[i], u2[ii]], [v2[i], v2[ii]], "r", scale=14 * ar, linewidth=3)

    # ax.set_xlim(-3, 3); ax.set_ylim(-6, 3)
    ax.set_yticks([-6, -3, 0, 3])
    # ax.set_yticklabels([" ", "0", " "])
    ax.grid(True)
    ax.set_xlabel("u-axis")
    ax.set_ylabel("v-axis")
    bold_ticks(ax, 14)
    for spine in ax.spines.values():
        spine.set_visible(True)
    plt.show()

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


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


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


if __name__ == "__main__":
    fig7_10()
