from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np

from plot_utils import finish_figure, matlab_figure, set_color_order, style_axes


def fig9_4():
    fig2 = matlab_figure(2, 560, 420)
    ax2 = fig2.add_subplot(111, projection="3d")
    set_color_order(ax2)

    # %% the following code changes coordinates to align with elliptic contours;
    # %% it is longer than it should be because it is written for the general
    # %% quadratic form E=ax^2+bxy+cy^2+dx+ey+f
    a11 = 3
    a12 = -2
    a22 = 4
    b1 = 1
    b2 = -1
    a = a11
    b = a12 / 2
    c = a22
    d = -b1
    e = -b2
    f = 0

    theta = 0.5 * abs(np.arctan(b / (a - c)))
    print("theta/pi =", 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 = 30
    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

    Emax = fb + A * B * rM**2
    Emin = fb
    print("Emax =", Emax)
    print("Emin =", Emin)

    ax2.plot_surface(A1, A2, E, cmap="viridis", edgecolor="none")
    ax2.contour(A1, A2, E, zdir="z", offset=float(np.min(E)), cmap="viridis")
    ax2.view_init(elev=20, azim=60)
    ax2.set_xlabel("x")
    ax2.set_ylabel("y")
    ax2.set_zlabel("F(x,y)")
    ax2.set_yticks([-0.5, 0, 0.5])
    ax2.set_xticks([-0.4, 0, 0.4, 0.8])
    style_axes(ax2, 14)
    ax2.set_box_aspect((1, 1, 0.7))
    fig2.tight_layout()

    # exportgraphics(gcf,'/Users/mark/Desktop/surf.eps')
    finish_figure(fig2, "surf.png")

    fig1 = matlab_figure(1, 560, 420)
    ax1 = fig1.add_subplot(111)
    set_color_order(ax1)

    # contour(A1,A2,E,[Emin+1 50 60 80 100 140 180 220 300 380 500 600],'b','LineWidth',1.3)
    ax1.contour(A1, A2, E, levels=[-0.11, 0.035, 0.3, 0.55, 0.8, 1.1], colors="b", linewidths=1.3)
    # contour(A1,A2,E,'LineWidth',1.3)
    ax1.set_xlabel("x-axis", fontsize=14, fontweight="bold")
    ax1.set_ylabel("y-axis", fontsize=14, fontweight="bold")
    ax1.set_yticks([-0.5, 0, 0.5])
    ax1.set_xticks([-0.4, 0, 0.4, 0.8])
    style_axes(ax1, 14)
    fig1.tight_layout()

    # exportgraphics(gcf,'/Users/mark/Desktop/cont.eps')
    finish_figure(fig1, "cont.png")
    plt.show()


if __name__ == "__main__":
    fig9_4()
