"""Literal Python conversion of fig8_3.m."""

from pathlib import Path
from time import perf_counter

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


xd = np.array([])
yd = np.array([])


def Einf(v1, v2):
    global xd, yd
    g = np.linalg.norm(v1 + v2 * xd - yd, np.inf)
    return g


def EEinf(v):
    global xd, yd
    g = np.linalg.norm(v[0] + v[1] * xd - yd, np.inf)
    return g


def E2(v1, v2):
    global xd, yd
    g = np.linalg.norm(v1 + v2 * xd - yd, 2) ** 2
    return g


def E22(v):
    global xd, yd
    g = np.linalg.norm(v[0] + v[1] * xd - yd, 2) ** 2
    return g


def einf_surface(A1, A2):
    """Vectorized equivalent of the original nested Einf loops."""
    global xd, yd
    E = np.zeros_like(A1, dtype=float)
    for k in range(len(xd)):
        E = np.maximum(E, np.abs(A1 + A2 * xd[k] - yd[k]))
    return E


def e2_surface(A1, A2):
    """Vectorized equivalent of the original nested E2 loops."""
    global xd, yd
    E = np.zeros_like(A1, dtype=float)
    for k in range(len(xd)):
        E = E + (A1 + A2 * xd[k] - yd[k]) ** 2
    return E


def E2_surf(rM):
    global xd, yd
    a = len(xd); b = 2 * np.sum(xd); c = np.dot(xd, xd)
    d = -2 * np.sum(yd); e = -2 * np.dot(xd, yd); f_ = np.dot(yd, yd)
    theta = 0.5 * abs(np.arctan(b / (a - c)))
    print(theta / np.pi)
    if b >= 0 and a >= c:
        ct = np.cos(theta)
        st = np.sin(theta)
    elif b >= 0 and a < c:
        ct = -np.cos(theta)
        st = np.sin(theta)
    elif b < 0 and a >= c:
        ct = np.cos(theta)
        st = -np.sin(theta)
    else:
        ct = -np.cos(theta)
        st = -np.sin(theta)

    ### transform to E=A*xp^2+B*yp^2+fb
    ### where x=x0+xp*ct-yp*ct, y=y0+xp*ct+yp*st
    x0 = (2 * c * d - b * e) / (b ** 2 - 4 * a * c)
    y0 = (2 * a * e - b * d) / (b ** 2 - 4 * a * c)
    A = a * ct ** 2 + b * ct * st + c * st ** 2
    B = a * st ** 2 - b * ct * st + c * ct ** 2
    fb = f_ + a * x0 ** 2 + b * x0 * y0 + c * y0 ** 2 + d * x0 + e * y0
    n0 = 100
    # rM=1/3
    r = np.linspace(0.0001, rM, n0)
    ang = np.linspace(0, 2 * np.pi, n0)
    R, ANG = np.meshgrid(r, ang, indexing="ij")
    xp = np.sqrt(B) * R * np.cos(ANG)
    yp = np.sqrt(A) * R * np.sin(ANG)
    A1 = x0 + xp * ct - yp * ct
    A2 = y0 + xp * ct + yp * st
    # E=a*A1**2+b*A1*A2+c*A2**2+d*A1+e*A2+f
    E = A * xp ** 2 + B * yp ** 2 + fb
    return A1, A2, E


def set_bold_ticks(ax, fontsize):
    ax.tick_params(labelsize=fontsize)
    for label in ax.get_xticklabels() + ax.get_yticklabels():
        label.set_fontweight("bold")
    if hasattr(ax, "get_zticklabels"):
        for label in ax.get_zticklabels():
            label.set_fontweight("bold")
            label.set_fontsize(fontsize)


def fig8_3(surface_n=1100, contour_n=2000, e2_contour_n=100):
    global xd, yd

    # plot error functionS for least square error

    # this takes maybe 15 sec to run
    start = perf_counter()

    #  generate data
    # n=50
    # xd=np.random.rand(n)
    # yd=np.zeros(n)
    # a1=3; a2=1
    # for i in range(n):
    #     yd[i]=a1*xd[i]+a2+(np.random.rand()-0.5)*(a1*xd[i]+a2)
    # surf_data=np.vstack((xd,yd))
    # np.savetxt('/Users/mark/Desktop/surf_data.txt',surf_data)

    surf_data = np.loadtxt(Path(__file__).with_name("surf_data.txt"))
    xd = surf_data[0, :]
    yd = surf_data[1, :]

    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],
    ])

    # E_infty surface
    v2a = -2; v2b = 8; v1a = -4; v1b = 6
    # n0=4000
    n0 = surface_n
    v2 = np.linspace(v2a, v2b, n0)
    v1 = np.linspace(v1a, v1b, n0)
    A1, A2 = np.meshgrid(v1, v2)
    E = einf_surface(A1, A2)

    fig1 = plt.figure(figsize=(5.6, 4.6), layout="constrained")
    ax1 = fig1.add_subplot(111, projection="3d")
    surf = ax1.plot_surface(A1, A2, E, cmap="copper", rcount=min(200, n0), ccount=min(200, n0))
    ax1.contour(A1, A2, E, zdir="z", offset=np.min(E), cmap="copper")
    ax1.view_init(elev=16, azim=-42)
    ax1.set_xlabel("v_1", fontsize=18, fontweight="bold")
    ax1.set_ylabel("v_2", fontsize=18, fontweight="bold")
    ax1.set_zlabel(r"$E_\infty$", fontsize=18, fontweight="bold")
    set_bold_ticks(ax1, 18)
    # fig1.savefig('/Users/mark/Desktop/surf1.eps')

    # fminsearch(@EM, a, optimset('TolFun',eps,'TolX',eps))
    vinfmim = minimize(EEinf, [2, 2], method="Nelder-Mead").x
    print(f"vinfmim = {vinfmim}")

    # E_infty contours
    n0 = contour_n
    v2 = np.linspace(v2a, v2b, n0)
    v1 = np.linspace(v1a, v1b, n0)
    A1, A2 = np.meshgrid(v1, v2)
    E = einf_surface(A1, A2)

    fig2, ax2 = plt.subplots(figsize=(5.6, 4.6), layout="constrained")
    ax2.set_prop_cycle(color=co)
    ax2.contour(A1, A2, E, 12, colors="b", linewidths=1.6)
    ax2.plot(vinfmim[0], vinfmim[1], "r*", markersize=14, linewidth=2)
    ax2.set_xlabel("v_1-axis", fontsize=18, fontweight="bold")
    ax2.set_ylabel("v_2-axis", fontsize=18, fontweight="bold")
    ax2.set_yticks([-4, -2, 0, 2, 4, 6, 8])
    set_bold_ticks(ax2, 18)
    # fig2.savefig('/Users/mark/Desktop/cont1.eps')

    # E_2 surface
    n0 = surface_n
    v2 = np.linspace(v2a, v2b, n0)
    v1 = np.linspace(v1a, v1b, n0)
    A1, A2 = np.meshgrid(v1, v2)
    E = einf_surface(A1, A2)
    rM = 0.7
    A1, A2, EE2 = E2_surf(rM)

    fig3 = plt.figure(figsize=(5.6, 4.6), layout="constrained")
    ax3 = fig3.add_subplot(111, projection="3d")
    ax3.plot_surface(A1, A2, EE2, cmap="copper", rcount=100, ccount=100)
    ax3.contour(A1, A2, EE2, zdir="z", offset=np.min(EE2), cmap="copper")
    ax3.view_init(elev=16, azim=-42)
    ax3.set_xlabel("v_1", fontsize=18, fontweight="bold")
    ax3.set_ylabel("v_2", fontsize=18, fontweight="bold")
    ax3.set_zlabel(r"$E_2$", fontsize=18, fontweight="bold")
    set_bold_ticks(ax3, 18)
    # fig3.savefig('/Users/mark/Desktop/surf2.eps')

    # E_2 contours
    v2a = -2; v2b = 8; v1a = -4; v1b = 6
    n0 = e2_contour_n
    v2 = np.linspace(v2a, v2b, n0)
    v1 = np.linspace(v1a, v1b, n0)
    A1, A2 = np.meshgrid(v1, v2)
    EE2 = e2_surface(A1, A2)

    # fminsearch(@EM, a, optimset('TolFun',eps,'TolX',eps))
    v2mim = minimize(E22, [2, 2], method="Nelder-Mead").x
    print(f"v2mim = {v2mim}")

    fig4, ax4 = plt.subplots(figsize=(5.6, 4.6), layout="constrained")
    ax4.set_prop_cycle(color=co)
    levels = [30, 60, 115, 190, 320, 480, 700, 900, 1200, 1550, 1900, 2400]
    ax4.contour(A1, A2, EE2, levels, colors="b", linewidths=1.6)
    ax4.plot(v2mim[0], v2mim[1], "r*", markersize=14, linewidth=2)
    ax4.set_xlabel("v_1-axis", fontsize=18, fontweight="bold")
    ax4.set_ylabel("v_2-axis", fontsize=18, fontweight="bold")
    ax4.set_yticks([-4, -2, 0, 2, 4, 6, 8])
    set_bold_ticks(ax4, 18)
    # fig4.savefig('/Users/mark/Desktop/cont2.eps')

    print(f"Elapsed time = {perf_counter() - start:.6f} seconds")
    plt.show()


if __name__ == "__main__":
    fig8_3()
