"""Literal Python conversion of fig3_6.m."""

import numpy as np
import matplotlib.pyplot as plt
from cycler import cycler


def fig3_6():
    plt.close("all")
    # get(gcf)
    # MATLAB Position [4 431 560 420], scaled to fit the plot window.
    fig1, ax1 = plt.subplots(num=1, figsize=(7.2, 5.4), constrained_layout=True)
    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],
    ])
    ax1.set_prop_cycle(cycler(color=co.tolist()))

    # compute and then plot contours
    xa = -3
    xb = 3
    ya = -3
    yb = 3
    n0 = 200
    x = np.linspace(xa, xb, n0)
    y = np.linspace(ya, yb, n0)
    A1, A2 = np.meshgrid(x, y)
    F = np.zeros((n0, n0))
    for i in range(n0):
        for j in range(n0):
            F[i, j] = f(x[i], y[j])
    # ax1.contour(A1,A2,F,[0.1,0.3,0.5,0.8,1.4,1.9,2.5,6,20],colors='k',linewidths=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.contour(A1,A2,F,colors='k')
    ax1.set_xlabel("x-axis", fontsize=14, fontweight="bold")
    ax1.set_ylabel("y-axis", fontsize=14, fontweight="bold")
    ax1.tick_params(labelsize=14)
    for label in ax1.get_xticklabels() + ax1.get_yticklabels():
        label.set_fontweight("bold")
    ax1.set_yticks([-1, -0.5, 0, 0.5, 1])

    # MATLAB clear * starts the second calculation with fresh variables.
    # get(gcf)
    # MATLAB Position [2 925 560 420], scaled to fit the plot window.
    fig2 = plt.figure(num=2, figsize=(7.2, 5.4), constrained_layout=True)
    ax2 = fig2.add_subplot(111, projection="3d")
    ax2.set_prop_cycle(cycler(color=co.tolist()))

    xa = -3
    xb = 3
    ya = -3
    yb = 3
    n0 = 60
    x = np.linspace(xa, xb, n0)
    y = np.linspace(ya, yb, n0)
    A1, A2 = np.meshgrid(x, y)
    F = np.zeros((n0, n0))
    for i in range(n0):
        for j in range(n0):
            F[i, j] = f(x[i], y[j])
    # surfc(A1,A2,F)
    ax2.plot_surface(A1, A2, F, cmap="viridis", linewidth=0, antialiased=True)
    zmin = float(np.min(F))
    ax2.contour(A1, A2, F, zdir="z", offset=zmin, linewidths=1.0)
    ax2.set_zlim(bottom=zmin)
    ax2.view_init(elev=16, azim=-120)
    ax2.set_xlabel("x", fontsize=14, fontweight="bold")
    ax2.set_ylabel("y", fontsize=14, fontweight="bold")
    ax2.set_zlabel("F(x,y)", fontsize=14, fontweight="bold")
    ax2.tick_params(labelsize=14, width=1)
    for label in ax2.get_xticklabels() + ax2.get_yticklabels() + ax2.get_zticklabels():
        label.set_fontweight("bold")
    ax2.grid(True)
    # H = gca; H.LineWidth = 1
    for axis in (ax2.xaxis, ax2.yaxis, ax2.zaxis):
        axis.line.set_linewidth(1)

    plt.show()
    return (fig1, ax1), (fig2, ax2)


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 - 5 * x + 3 * y ** 2) * np.exp(-x ** 2 - y ** 2)
    return z


if __name__ == "__main__":
    fig3_6()
