from __future__ import division
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
from numpy import pi, linspace, cos, sin


def x_(u, v):
    x = -2 + 2*cos(v) - cos(u)
    x[v<3*pi] = -2 + (2 + cos(u[v<3*pi]))*cos(v[v<3*pi])
    x[v<2*pi] = cos(u[v<2*pi])*(2.5 - 1.5*cos(v[v<2*pi]))
    return x

def y_(u, v):
    y = sin(u)
    y[v<2*pi] = sin(u[v<2*pi])*(2.5 - 1.5*cos(v[v<2*pi]))
    return y

def z_(u, v):
    z = -3*v + 12*pi
    z[v<3*pi] = (2 + cos(u[v<3*pi]))*sin(v[v<3*pi]) + 3*pi
    z[v<2*pi] = 3*v[v<2*pi] - 3*pi
    z[v<pi] = -2.5*sin(v[v<pi])
    return z

 
u = linspace(pi, 2*pi, 10)
v = linspace(0, 4*pi, 38)
u, v = np.meshgrid(u, v)

x = x_(u, v)
y = y_(u, v)
z = -z_(u, v)


fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap="YlGnBu_r", alpha=0.8,
                lw=1)
ax.set_xlim(-5, 5)
ax.set_ylim(-5, 5)
ax.set_zlim(-9, 1)
ax.view_init(azim=130, elev=30)
plt.axis("off")
plt.savefig("KleinBottle-cut.svg", bbox_inches="tight", transparent=True)
plt.show()