"""Literal Python conversion of fig8_17.m."""

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


td = np.array([])
zd = np.empty((0, 2))
tmax = 0.0


# right-hand side of DE
def f(t, y, p1, p2, p3):
    z = np.zeros(2)
    z[0] = p1 * (y[0] - y[0] ** 3 + y[1])
    z[1] = -(y[0] - p3 + p2 * y[1]) / p1
    return z


# RK4 method
def rk4(t, y0, h, n, p1, p2, p3):
    y = np.asarray(y0, dtype=float)
    ypoints = [y[0]]
    vpoints = [y[1]]
    for i in range(1, n):
        k1 = h * f(t[i - 1], y, p1, p2, p3)
        k2 = h * f(t[i - 1] + 0.5 * h, y + 0.5 * k1, p1, p2, p3)
        k3 = h * f(t[i - 1] + 0.5 * h, y + 0.5 * k2, p1, p2, p3)
        k4 = h * f(t[i], y + k3, p1, p2, p3)
        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)


def fspline(t, y, dyd, td_eval):
    """Direct replacement for the external MATLAB fspline routine."""
    dyd = np.asarray(dyd, dtype=float).ravel()
    return CubicSpline(t, y, bc_type=((1, dyd[0]), (1, dyd[1])))(td_eval)


def F(p):
    global td, zd, tmax
    p1 = p[0]; p2 = p[1]; p3 = p[2]
    m = 60
    y0 = -1
    v0 = 1
    z0 = np.array([y0, v0], dtype=float)
    t = np.linspace(0, tmax, m)
    k = t[1] - t[0]
    y, v = rk4(t, z0, k, m, p1, p2, p3)
    dy0 = f(0, z0, p1, p2, p3)
    dy1 = f(tmax, z0, p1, p2, p3)
    # These calls intentionally follow the original use of dy0 and dy1.
    yf = fspline(t, y, dy0, td)
    vf = fspline(t, v, dy1, td)
    s = np.linalg.norm(yf - zd[:, 0], 2) ** 2 + np.linalg.norm(vf - zd[:, 1], 2) ** 2
    return s


def fig8_17():
    #  fit solution of IVP to data to find parameters p1 and p2

    #  requires file splineA.m

    #  IVP:   y'=ff(y,v)  v'=gg(y,v)  for 0 < t < tmax
    #         y(0)=y0     v(0)=v0

    #  method: uses RK4 + cubic spline + fminsearch

    #  as written, the FitzHugh-Nagumo equations are used

    #  the functions ff(y,v) and gg(y,v) are entered in function f(t,y,p1,p2)

    global td, zd, tmax

    ####  generate data for regression
    ####  note that tmax, y0, v0 must also be entered in function F
    #  first find numerical solution for given values
    p1 = 3; p2 = 0.2; p3 = 0.2
    tmax = 20
    y0 = -1
    v0 = 1
    nt = 100
    t = np.linspace(0, tmax, nt)
    solution = solve_ivp(lambda tt, yy: f(tt, yy, p1, p2, p3), (0, tmax), [y0, v0], t_eval=t)
    t = solution.t
    z = solution.y.T
    dz0 = f(0, z[0, :], p1, p2, p3)
    dz1 = f(tmax, z[nt - 1, :], p1, p2, p3)

    # now randomize data:  nd = number of data points
    # nd=120
    nd = 60
    yd = np.zeros(nd); vd = np.zeros(nd); zd = np.zeros((nd, 2))
    # td=np.linspace(0.1,19.9,nd)
    tdd = tmax * np.random.rand(nd)
    td = np.sort(tdd)
    td[0] = 0
    qd = 0.4 * (np.random.rand(nd) - 0.5)
    qdd = 0.5 * (np.random.rand(nd) - 0.5)
    print(np.shape(t))
    print(np.shape(z[:, 0].T))
    print(np.shape(td))
    yd = CubicSpline(t, z[:, 0], bc_type=((1, dz0[0]), (1, dz1[0])))(td)
    vd = CubicSpline(t, z[:, 1], bc_type=((1, dz0[1]), (1, dz1[1])))(td)
    for in_ in range(1, nd):
        yd[in_] = (1 + qd[in_]) * yd[in_]
        vdd = vd[in_]
        if abs(qdd[in_]) > 0.2:
            vd[in_] = (1 + qdd[in_]) * vd[in_]
        else:
            vd[in_] = vd[in_] + 0.4 * (np.random.rand() - 0.5)
        zd[in_, 0] = yd[in_]
        zd[in_, 1] = vd[in_]

    result = minimize(
        F,
        [1, 1, 1],
        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)
    pd1 = sol[0]; pd2 = sol[1]; pd3 = sol[2]

    #  p1=c  p2=b  p3=a
    # z(1)=p1*(y(1)-y(1)^3+y(2));
    # z(2)=-(y(1)-p3+p2*y(2))/p1;
    print(f"\n  Computed Solution =  {pd1:e}  {pd2:e}   {pd3:e}\n")

    ### plot data and solution using calcuated parameter values
    plt.close("all")
    # get(gcf)
    # MATLAB Position [4 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)
    nt = 200
    t = np.linspace(0, tmax, nt)
    solution = solve_ivp(lambda tt, yy: f(tt, yy, pd1, pd2, pd3), (0, tmax), [y0, v0], t_eval=t)
    t = solution.t
    z = solution.y.T
    ax.plot(t, z[:, 0], "r", linewidth=1.6)
    ax.plot(t, z[:, 1], "b", linewidth=1.6)
    ax.axis([0, tmax, -1.5, 1.5])
    ax.legend([" v", " w"], loc="lower left", fontsize=14, prop={"weight": "bold", "size": 14})
    ax.grid(True)
    ax.set_yticks([-1.5, -0.5, 0.5, 1.5])
    ax.plot(td, yd, ".r", markersize=18)
    ax.plot(td, vd, "db", markersize=5, linewidth=1.5)
    ax.set_xlabel("t-axis", fontsize=14, fontweight="bold")
    # ax.set_ylabel('y-axis')
    ax.tick_params(labelsize=14)
    for label in ax.get_xticklabels() + ax.get_yticklabels():
        label.set_fontweight("bold")

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


if __name__ == "__main__":
    fig8_17()
