from __future__ import division
from scipy.interpolate import splev
from numpy import (sqrt, linspace, sin, cos, array, hstack, transpose, pi,
                   ones_like, zeros_like, where, r_)
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.tri import Triangulation

def splevn(t, tks):
    """Like splev but control points are vectors."""
    return zip(*[splev(t, [tks[0], p, tks[-1]]) for p in zip(*tks[1])])

# Knots, weighted control points, degree:
arcNURBS = [[0,0,0,0.25,0.25,0.5,0.5,0.75,0.75,1,1,1],
            [[1,0,1],
             [sqrt(0.5), sqrt(0.5), sqrt(0.5)],
             [0,1,1],
             [-sqrt(0.5), sqrt(0.5), sqrt(0.5)],
             [-1,0,1],
             [-sqrt(0.5), -sqrt(0.5), sqrt(0.5)],
             [0,-1,1],
             [sqrt(0.5), -sqrt(0.5), sqrt(0.5)],
             [1,0,1]], 2]
theta = linspace(0, 2*pi, 7)
pts = transpose([cos(theta), sin(theta)])
w = 1.0 / array([1.,2,1,2,1,2,1])
arcNURBS = [array([0.,0,0,1,1,2,2,3,3,3]) / 3.0,
            hstack([pts, w[:,None]]),
            2]
Pw = array(arcNURBS[1])
t = linspace(arcNURBS[0][0], arcNURBS[0][-1], 6*15+1)

x, y, w = transpose(splevn(t, arcNURBS))

fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(111, projection='3d')
ax.view_init(azim=-50, elev=20)
for setRange in (ax.set_xlim, ax.set_ylim):
    setRange(-1.5, 1.5)
ax.set_zlim(-0.5, 1.5)

# Shade in the cone:
xs = hstack([[0], x/w])
ys = hstack([[0], y/w])
ws = hstack([[0], ones_like(w)])
tris = []
for i in range(len(x)):
    tris.append([0, (i + 1) % len(xs), (i + 2) % len(xs)])


color = (0.9,)*3
cone = ax.plot_trisurf(xs, ys, tris, ws, color=color, edgecolor=color,
                       alpha=1.0, lw=0)

Pw = transpose(arcNURBS[1])
ax.plot(Pw[0]/Pw[2], Pw[1]/Pw[2], Pw[2]/Pw[2], 'k.-',
        label='NURBS control w/o weights')
p = ax.plot(Pw[0], Pw[1], Pw[2], '.-', color='b',
            label='3D control polygon')[0]
p.set_zorder(-1000)
ax.plot(x, y, w, 'b', label='3D parabolas')
ax.plot(x/w, y/w, ones_like(w), color=(0.85,0.1,0.1), linewidth=2,
        label='NURBS circle at w = 1')
if True:
    ax.plot(Pw[0], Pw[1], ones_like(Pw[2]), '.-', color='b', alpha=0.2)
    ax.plot(x, y, ones_like(w), 'b', alpha=0.2)

if not True:
    plots = []
    plots += ax.plot(Pw[0]/Pw[2], Pw[1]/Pw[2], 0*Pw[2]/Pw[2], 'k.-', alpha=0.2)
    plots += ax.plot(Pw[0], Pw[1], 0*Pw[2], '.-', color='b', alpha=0.2)
    plots += ax.plot(x, y, 0*w, 'b', alpha=0.2)
    plots += ax.plot(x/w, y/w, zeros_like(w), color=(0.85,0.1,0.1),
                     linewidth=2, alpha=0.2)
    for i, p in enumerate(plots):
        p.set_zorder(-1000-i)

DRAW_RULES_ON_CONE = False
if DRAW_RULES_ON_CONE:
    for i in r_[:len(t):10]:
        ax.plot([0, x[i], x[i]/w[i]], [0, y[i], y[i]/w[i]], [0, w[i], 1],
                'k.--')


for i in where(Pw[-1] != 1)[0]:
    ax.plot([0, Pw[0,i]/Pw[2,i]], [0, Pw[1,i]/Pw[2,i]], [0, 1], 'k.--',
            alpha=0.5)


plt.xlabel(r"$x$", fontsize=18)
plt.ylabel(r"$y$", fontsize=18)
ax.set_zlabel(r"$z$", fontsize=18)
plt.legend()
plt.savefig("NURBS-circle-3D.svg", transparent=True, bbox_inches="tight")
plt.show()