from __future__ import annotations

import os

import matplotlib.pyplot as plt
import numpy as np

from plot_utils import finish_figure, matlab_figure, set_color_order, style_axes


def f(x, y):
    # z=-x^4+x^6/6-10*x*y+y+y^4;
    # z=100*(x^2-y)^2+(x-1)^2;
    z = (x**2 + 4 * y**2 - 1) ** 2 + (4 * x**2 + y**2 - 1) ** 2
    return z


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

    # compute and then plot contours
    xa = -1
    xb = 1
    ya = -1
    yb = 1
    n0 = int(os.getenv("FIG9_1_N0", "2000"))
    x = np.linspace(xa, xb, n0)
    y = np.linspace(ya, yb, n0)
    A1, A2 = np.meshgrid(x, y)
    F = f(A1, A2)
    # contour(A1,A2,F,'k')
    # contour(A1,A2,F,[ 0.1 0.3 0.5 0.8 1.4 1.9 2.5 6 20],'k','LineWidth',1.2)
    ax1.contour(A1, A2, F, levels=[0.04, 0.3, 0.5, 0.8, 1.5, 1.95, 4.3, 8, 12, 16, 22], linewidths=1.2)
    ax1.set_xlabel("x-axis", fontsize=14, fontweight="bold")
    ax1.set_ylabel("y-axis", fontsize=14, fontweight="bold")
    ax1.set_yticks([-1, -0.5, 0, 0.5, 1])
    style_axes(ax1, 14)
    fig1.tight_layout()

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

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

    xa = -0.7
    xb = 0.7
    ya = -0.7
    yb = 0.7
    n0 = int(os.getenv("FIG9_1_SURFACE_N0", "60"))
    x = np.linspace(xa, xb, n0)
    y = np.linspace(ya, yb, n0)
    A1, A2 = np.meshgrid(x, y)
    F = f(A1, A2)
    ax2.plot_surface(A1, A2, F, cmap="viridis", edgecolor="none")
    ax2.contour(A1, A2, F, zdir="z", offset=float(np.min(F)), cmap="viridis")
    ax2.view_init(elev=16, azim=-32)
    ax2.set_xlabel("x")
    ax2.set_ylabel("y")
    ax2.set_zlabel("F(x,y)")
    style_axes(ax2, 14)
    ax2.set_box_aspect((1, 1, 0.7))
    fig2.tight_layout()

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


if __name__ == "__main__":
    fig9_1()
