import math
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import LinearLocator


def f1(x):
    return 2 * math.sin(x) - x ** 2 / 10


def showf1():
    t = np.linspace(0, 4, 1000)
    plt.plot(t, [f1(x) for x in t])
    plt.show()


def f2(x, y):
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2


def f3(x, y):
    return (1.5 - x + x * y) ** 2 + (2.25 - x + x * y ** 2) ** 2 + (2.625 - x + x * y ** 3) ** 2


def showf2():
    fig = plt.figure(figsize=plt.figaspect(1.))
    ax = fig.add_subplot(1, 1, 1, projection='3d')
    X = np.arange(-1, 1, 0.05)
    Y = np.arange(-1, 1, 0.05)
    X, Y = np.meshgrid(X, Y)
    # R = np.sqrt(X ** 2 + Y ** 2)
    Z = f2(X, Y)
    # Z = np.sin(R)

    surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1,
                           cmap=plt.cm.YlGnBu_r,
                           linewidth=0, antialiased=False)
    fig.colorbar(surf, shrink=0.5, aspect=5)
    plt.show()


def showf3():
    fig = plt.figure(figsize=plt.figaspect(1.))
    ax = fig.add_subplot(1, 1, 1, projection='3d')
    X = np.arange(-5, 5, 0.25)
    Y = np.arange(-5, 5, 0.25)
    X, Y = np.meshgrid(X, Y)
    # R = np.sqrt(X ** 2 + Y ** 2)
    Z = f3(X, Y)
    # Z = np.sin(R)

    surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1,
                           cmap=plt.cm.YlGnBu_r,
                           linewidth=0, antialiased=False)
    fig.colorbar(surf, shrink=0.5, aspect=5)
    plt.show()


if __name__ == '__main__':
    showf3()