3D Plots [LAST FUNC LEFT]

from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
import numpy as np


def ellipsoid():
    # DONE!
    a, b, c = map(float, input("Введите коэффициенты a, b и c через пробел: ").split())
    coefs = (a, b, c)
    rx, ry, rz = (coefs)

    # Установка всех сферических углов:
    u = np.linspace(0, 2 * np.pi, 100)
    v = np.linspace(0, np.pi, 100)

    # Декартовы координаты, которые соответствуют сферическим углам
    # (уравнение эллипсоида):
    x = rx * np.outer(np.cos(u), np.sin(v))
    y = ry * np.outer(np.sin(u), np.sin(v))
    z = rz * np.outer(np.ones_like(u), np.cos(v))

    ax.plot_wireframe(x, y, z, rstride=6, cstride=6)

    # Регулировка осей, чтобы все они имели одинаковый промежуток
    max_radius = max(rx, ry, rz)
    for axis in 'xyz':
        getattr(ax, 'set_{}lim'.format(axis))((-max_radius, max_radius))


def oneline_hyperboloid():
    # DONE!
    a, b, c = map(float, input("Введите коэффициенты a, b и c через пробел: ").split())
    coefs = (a, b, c)
    rx, ry, rz = (coefs)  # 1 / np.sqrt(coefs)
    u = np.linspace(-2, 2, 100)
    v = np.linspace(0, 2 * np.pi, 60)
    [u, v] = np.meshgrid(u, v)
    axys = input("Введите направляющую ось x, y или z: ")
    axys.lower()
    while axys != "x" and axys != "y" and axys != "z":
        axys = input("Ошибка, введите x, y или z: ")
    if axys == "z":
        x = rx * np.cosh(u) * np.cos(v)
        y = ry * np.cosh(u) * np.sin(v)
        z = rz * np.sinh(u)
    elif axys == "y":
        z = rx * np.cosh(u) * np.cos(v)
        y = ry * np.cosh(u) * np.sin(v)
        x = rz * np.sinh(u)
    elif axys == "x":
        x = rx * np.cosh(u) * np.cos(v)
        z = ry * np.cosh(u) * np.sin(v)
        y = rz * np.sinh(u)

    ax.plot_wireframe(x, y, z, rstride=6, cstride=6)

    max_radius = max(rx * 4, ry * 4, rz * 4)
    for axis in 'xyz':
        getattr(ax, 'set_{}lim'.format(axis))((-max_radius, max_radius))


def twoline_hyperboloid():
    # DONE!
    a, b, c = map(float, input("Введите коэффициенты a, b и c через пробел: ").split())
    coefs = (a, b, c)
    rx, ry, rz = coefs
    u = np.linspace(-3, 3, 100)
    v = np.linspace(0, 2 * np.pi, 100)
    u, v = np.meshgrid(u, v)
    axys = input("Введите направляющую ось x, y или z: ")
    axys.lower()
    while axys != "x" and axys != "y" and axys != "z":
        axys = input("Ошибка, введите x, y или z: ")
    if axys == "z":
        x = a * np.sinh(u) * np.cos(v)
        y = b * np.sinh(u) * np.sin(v)
        z = c * np.cosh(u)
        ax.plot_wireframe(x / 10, y / 10, z / 10, rstride=10, cstride=10)
        ax.plot_wireframe(x / 10, y / 10, -z / 10, rstride=10, cstride=10)
    elif axys == "y":
        z = a * np.sinh(u) * np.cos(v)
        y = b * np.sinh(u) * np.sin(v)
        x = c * np.cosh(u)
        ax.plot_wireframe(x / 10, y / 10, z / 10, rstride=10, cstride=10)
        ax.plot_wireframe(-x / 10, y / 10, z / 10, rstride=10, cstride=10)
    elif axys == "x":
        x = a * np.sinh(u) * np.cos(v)
        z = b * np.sinh(u) * np.sin(v)
        y = c * np.cosh(u)
        ax.plot_wireframe(x / 10, y / 10, z / 10, rstride=10, cstride=10)
        ax.plot_wireframe(x / 10, -y / 10, -z / 10, rstride=10, cstride=10)

    max_radius = max(rx, ry, rz)
    for axis in 'xyz':
        getattr(ax, 'set_{}lim'.format(axis))((-max_radius, max_radius))


def konus():
    # DONE!
    a, b, c = map(float, input("Введите коэффициенты a, b и c через пробел: ").split())
    coefs = (a, b, c)
    rx, ry, rz = coefs
    u = np.linspace(0, 2 * np.pi, 100)
    v = np.linspace(0, 2, 100)
    t, r = np.meshgrid(u, v)
    axys = input("Введите направляющую ось x, y или z: ")
    axys.lower()
    while axys != "x" and axys != "y" and axys != "z":
        axys = input("Ошибка, введите x, y или z: ")
    if axys == "z":
        x = r * a * np.cos(t)
        y = r * b * np.sin(t)
        z = -r * c

        ax.plot_wireframe(x/2, y/2, z/2, rstride=6, cstride=6)
        ax.plot_wireframe(x/2, y/2, -z/2, rstride=6, cstride=6)

    elif axys == "y":
        z = r * a * np.cos(t)
        y = r * b * np.sin(t)
        x = -r * c

        ax.plot_wireframe(x / 2, y / 2, z / 2, rstride=6, cstride=6)
        ax.plot_wireframe(-x / 2, y / 2, z / 2, rstride=6, cstride=6)

    elif axys == "x":
        x = r * a * np.cos(t)
        z = r * b * np.sin(t)
        y = -r * c

        ax.plot_wireframe(x / 2, y / 2, z / 2, rstride=6, cstride=6)
        ax.plot_wireframe(x / 2, -y / 2, z / 2, rstride=6, cstride=6)

    max_radius = max(rx, ry, rz)
    for axis in 'xyz':
        getattr(ax, 'set_{}lim'.format(axis))((-max_radius, max_radius))


def elliptic_paraboloid():
    # problems with z axys camera
    a, b, c = map(float, input("Введите коэффициенты a, b и c через пробел: ").split())
    u = np.linspace(0, 2 * np.pi, 100)
    v = np.linspace(0, 2, 100)
    t, r = np.meshgrid(u, v)

    rx = r * a
    ry = r * b
    rz = r * c

    axys = input("Введите направляющую ось x, y или z: ")
    axys.lower()
    while axys != "x" and axys != "y" and axys != "z":
        axys = input("Ошибка, введите x, y или z: ")
    if axys == "z":
        x = rx * np.cos(t)
        y = ry * np.sin(t)
        z = rz ** 2

    elif axys == "y":
        z = rx * np.cos(t)
        y = ry * np.sin(t)
        x = rz ** 2

    elif axys == "x":
        x = rx * np.cos(t)
        z = ry * np.sin(t)
        y = rz ** 2

    ax.plot_wireframe(x, y, z, rstride=10, cstride=10)
    max_radius = (2 * max(a, b, c)) ** 2
    for axis in 'xyz':
        getattr(ax, 'set_{}lim'.format(axis))((-max_radius, max_radius))

def hyperbolic_paraboloid():
    # not sure
    a, b = map(float, input("Введите коэффициенты a и b через пробел: ").split())
    x = np.linspace(-2, 2, 100)
    y = np.linspace(-2, 2, 100)
    x, y = np.meshgrid(x, y)
    z = (x ** 2 / a ** 2) - (y ** 2 / b ** 2)

    ax.plot_wireframe(x, y, z, rstride=10, cstride=10)


def elliptic_cylinder():
    # DONE!
    a, b = map(float, input("Введите коэффициенты a и b через пробел: ").split())
    axys = input("Введите направляющую ось x, y или z: ")
    axys.lower()
    while axys != "x" and axys != "y" and axys != "z":
        axys = input("Ошибка, введите x, y или z: ")
    if axys == "z":
        x = np.linspace(-2, 2, 100)
        z = np.linspace(0, 10, 100)
        x, z = np.meshgrid(x, z)
        y = np.sqrt(2 ** 2 - x ** 2)

        x *= a / 2
        y *= b / 2
        z /= 10

        ax.plot_wireframe(x, y, z, rstride=10, cstride=10)
        ax.plot_wireframe(x, -y, z, rstride=10, cstride=10)

        max_radius = max(a, b, 1)
        for axis in 'xyz':
            getattr(ax, 'set_{}lim'.format(axis))((-max_radius, max_radius))
        ax.set_zlim(0, 1)

    elif axys == "x":
        x = np.linspace(-2, 2, 100)
        y = np.linspace(0, 10, 100)
        x, y = np.meshgrid(x, y)
        z = np.sqrt(2 ** 2 - x ** 2)

        x *= a / 2
        z *= b / 2
        y /= 10

        ax.plot_wireframe(x, y, z, rstride=10, cstride=10)
        ax.plot_wireframe(x, y, -z, rstride=10, cstride=10)

        max_radius = max(a, b, 1)
        for axis in 'xyz':
            getattr(ax, 'set_{}lim'.format(axis))((-max_radius, max_radius))
        ax.set_ylim(0, 1)
    elif axys == "y":
        y = np.linspace(-2, 2, 100)
        x = np.linspace(0, 10, 100)
        y, x = np.meshgrid(y, x)
        z = np.sqrt(2 ** 2 - y ** 2)

        y *= a / 2
        z *= b / 2
        x /= 1

        ax.plot_wireframe(x, y, z, rstride=10, cstride=10)
        ax.plot_wireframe(x, y, -z, rstride=10, cstride=10)

        max_radius = max(a, b, 1)
        for axis in 'xyz':
            getattr(ax, 'set_{}lim'.format(axis))((-max_radius, max_radius))
        ax.set_xlim(0, 1)


def hyperbolic_cylinder():
    # DONE!
    a, b = map(float, input("Введите коэффициенты a и b через пробел: ").split())
    coefs = (a, b)
    rx, ry = coefs
    u = np.linspace(-2, 2, 100)
    v = np.linspace(-2, 2, 100)
    [u, v] = np.meshgrid(u, v)
    axys = input("Введите направляющую ось x, y или z: ")
    axys.lower()
    while axys != "x" and axys != "y" and axys != "z":
        axys = input("Ошибка, введите x, y или z: ")
    if axys == "z":
        x = b * np.sinh(u)
        y = a * np.cosh(u)
        z = v

        ax.plot_wireframe(x, y, z, rstride=10, cstride=10)
        ax.plot_wireframe(x, -y, z, rstride=10, cstride=10)

        max_radius = max(rx * 4, ry * 4)
        for axis in 'xyz':
            getattr(ax, 'set_{}lim'.format(axis))((-max_radius, max_radius))
        ax.set_zlim(-2, 2)
    elif axys == "x":
        x = b * np.sinh(u)
        z = a * np.cosh(u)
        y = v

        ax.plot_wireframe(x, y, z, rstride=10, cstride=10)
        ax.plot_wireframe(x, y, -z, rstride=10, cstride=10)

        max_radius = max(rx * 4, ry * 4)
        for axis in 'xyz':
            getattr(ax, 'set_{}lim'.format(axis))((-max_radius, max_radius))
        ax.set_ylim(-2, 2)
    elif axys == "y":
        y = b * np.sinh(u)
        z = a * np.cosh(u)
        x = v

        ax.plot_wireframe(x, y, z, rstride=10, cstride=10)
        ax.plot_wireframe(x, y, -z, rstride=10, cstride=10)

        max_radius = max(rx * 4, ry * 4)
        for axis in 'xyz':
            getattr(ax, 'set_{}lim'.format(axis))((-max_radius, max_radius))
        ax.set_xlim(-2, 2)


def parabolic_cylinder():
    a = float(input("Введите коэффициент a: "))
    u = np.linspace(-2, 2, 100)
    v = np.linspace(-2, 2, 100)
    [u, v] = np.meshgrid(u, v)
    axys = input("Введите направляющую ось x, y или z: ")
    axys.lower()
    while axys != "x" and axys != "y" and axys != "z":
        axys = input("Ошибка, введите x, y или z: ")
    if axys == "z":
        x = u**2/np.pi
        y = u
        z = a * v

        ax.plot_wireframe(x, y, z, rstride=10, cstride=10)

        max_radius = a
        for axis in 'xyz':
            getattr(ax, 'set_{}lim'.format(axis))((-max_radius, max_radius))
        # ax.set_xlim(-2, 2)



def choose_type(type):
    if type == 1:
        ellipsoid()
    elif type == 2:
        oneline_hyperboloid()
    elif type == 3:
        twoline_hyperboloid()
    elif type == 4:
        konus()
    elif type == 5:
        elliptic_paraboloid()
    elif type == 6:
        hyperbolic_paraboloid()
    elif type == 7:
        elliptic_cylinder()
    elif type == 8:
        hyperbolic_cylinder()
    elif type == 9:
        parabolic_cylinder()

# Инициализируем фигуру
fig = plt.figure()

ax = fig.add_subplot(111, projection='3d')

# Выбираем тип поверхности, после чего считываем коэффициенты
print("Пожалуйста, выберите тип поверхности:\n1. Эллипсоид\n2. Однополсоный гиперболоид\n3. Двуполосный "
      "гиперболоид\n4. Коническая поверхность\n5. Эллиптический параболлоид\n6. Гиперболический параболлоид\n7. "
      "Эллиптический цилиндр\n8. Гиперболический цилиндр\n9. Параболлический цилиндр")
type = int(input())

choose_type(type)

# Настройка отрисовки
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")

plt.title("")

plt.show()