"""Literal Python conversion of fig8_16.m."""

import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import CubicSpline
from scipy.optimize import minimize


td = np.array([])
yd = np.array([])
tmax = 0.0


def f(t, y, alpha, beta):
    dy = alpha * y * (beta - y)
    return dy


# RK4 method
def rk4(t, y0, h, n, alpha, beta):
    y = y0
    ypoints = [y0]
    for i in range(1, n):
        k1 = h * f(t[i - 1], y, alpha, beta)
        k2 = h * f(t[i - 1] + 0.5 * h, y + 0.5 * k1, alpha, beta)
        k3 = h * f(t[i - 1] + 0.5 * h, y + 0.5 * k2, alpha, beta)
        k4 = h * f(t[i], y + k3, alpha, beta)
        yy = y + (k1 + 2 * k2 + 2 * k3 + k4) / 6
        ypoints.append(yy)
        y = yy
    return np.asarray(ypoints)


def F(p):
    global td, yd, tmax
    alpha = p[0]; beta = p[1]
    y0 = 1
    m = 6
    t = np.linspace(0, tmax, m + 1)
    k = t[1] - t[0]
    y_rk4 = rk4(t, y0, k, m + 1, alpha, beta)
    dy0 = f(0, y0, alpha, beta)
    dy1 = f(tmax, y_rk4[m], alpha, beta)
    # MATLAB splineA(...,[1 1],[dy0 dy1]) is a clamped cubic spline.
    yf = CubicSpline(t, y_rk4, bc_type=((1, dy0), (1, dy1)))(td)
    s = np.linalg.norm(yf - yd, 2) ** 2
    return s


def fig8_16():
    # fit logistic equation using RK4 + cubic spline + nead-medler

    global td, yd, tmax

    #  generate data for regression
    a = 1; alpha0 = 2; beta0 = 3
    n = 80
    tmax = 2
    # td=np.linspace(0,tmax,n)
    tdd = tmax * np.random.rand(n)
    td = np.sort(tdd)
    qd = 0.2 * (np.random.rand(n) - 0.5)
    yd = np.zeros(n)
    td[0] = 0; yd[0] = a
    for in_ in range(1, n):
        yd[in_] = (1 + qd[in_]) * a * beta0 / (
            a + (beta0 - a) * np.exp(-alpha0 * beta0 * td[in_])
        )

    result = minimize(
        F,
        [1, 2],
        method="Nelder-Mead",
        options={"maxfev": 1_000_000, "maxiter": 1_000_000, "xatol": 1e-6},
    )
    sol = result.x
    fv = result.fun
    ee = result.status
    o = result
    print(sol, fv, ee, o)
    alpha = sol[0]; beta = sol[1]
    print(f"\n  Computed Solution =  {alpha:e}  {beta:e}\n")

    nt = 100
    y0 = a
    tt = np.linspace(0, tmax, nt)
    k = tt[1] - tt[0]
    yy = rk4(tt, y0, k, nt, alpha, beta)

    plt.close("all")
    # get(gcf)
    # MATLAB Position [1 1115 658 230], scaled to fit the plot window.
    fig, ax = plt.subplots(figsize=(6.58, 3.0), layout="constrained")
    co = np.array([
        [0, 0, 1],
        [0, 0.5, 0],
        [1, 0, 0],
        [0, 0.75, 0.75],
        [0.75, 0, 0.75],
        [0.75, 0.75, 0],
        [0.25, 0.25, 0.25],
    ])
    ax.set_prop_cycle(color=co)
    ax.plot(tt, yy, "r", linewidth=1.6)
    ax.grid(True)
    ax.axis([0, tmax, 1, 3.5])
    ax.plot(td, yd, "b.", markersize=20)
    ax.set_xlabel("t-axis", fontsize=14, fontweight="bold")
    ax.set_ylabel("y-axis", fontsize=14, fontweight="bold")
    ax.tick_params(labelsize=14)
    for label in ax.get_xticklabels() + ax.get_yticklabels():
        label.set_fontweight("bold")

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


if __name__ == "__main__":
    fig8_16()
