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_15T():
    # this also has fig9_16T

    # surface/countour plots of surface
    # export: expand axis

    # Specify which example to run ifun=1,2
    ifun = 1

    fig = matlab_figure(None, 639, 368)
    ax = fig.add_subplot(111, projection="3d")
    set_color_order(ax)

    if ifun == 1:
        nn = 20
        xa = -1.4
        xb = 1.4
        ya = -1.1
        yb = 2
        xd = np.linspace(xa, xb, nn)
        yd = np.linspace(ya, yb, nn)
        X, Y = np.meshgrid(xd, yd)
        Z = 100 * (X**2 - Y) ** 2 + (X - 1) ** 2
        azim = 125
        elev = 25
    else:
        nn = 20
        xa = -1
        xb = 1
        ya = -1
        yb = 1
        xd = np.linspace(xa, xb, nn)
        yd = np.linspace(ya, yb, nn)
        X, Y = np.meshgrid(xd, yd)
        Z = (X - Y) ** 4 + 8 * X * Y - X + Y + 3
        azim = -25
        elev = 12

    ax.grid(True)
    ax.plot_surface(X, Y, Z, cmap="viridis", edgecolor="none")
    ax.contour(X, Y, Z, zdir="z", offset=float(np.min(Z)), cmap="viridis")
    ax.view_init(elev=elev, azim=azim)
    ax.set_xlabel("x-axis")
    ax.set_ylabel("y-axis")
    style_axes(ax, 16)
    ax.set_box_aspect((1.2, 1, 0.7))
    fig.tight_layout()
    finish_figure(fig, "fig9_15T.png")
    plt.show()


if __name__ == "__main__":
    fig9_15T()
