# PastaParam.py
# This plots parametric plots in 3D of the pasta at http://nestedtori.com.
# This was tested and run in Python 3 using the Anaconda distribution,
# but should work in any installation that has matplotlib and numpy.

import matplotlib
import matplotlib.pyplot as plt
# get_ipython().magic('matplotlib qt')

import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

# a parametric plotting routine that handles all the mesh-gridding automatically
# allows specification of some colors: each element in the mesh can be assigned a
# different color using colorinfo. To do that, provide a function of two variables.

# parameter flist is a list of 3 functions of two variables (they can take lambdas)
def parm_plot_3d(flist, srange, trange,colorinfo=None,cmap=None,fig=plt.gcf(),ax=plt.gca()):
    s = np.linspace(*srange)
    t = np.linspace(*trange)
    s,t= np.meshgrid(s,t)
    X,Y,Z = [h(s,t) for h in flist]
    if colorinfo is None:
        return ax.plot_surface(X, Y, Z, rstride=1, cstride=1,cmap=cmap,
                                antialiased=True,linewidth=0,shade=True)

    colors = np.array([ [colorinfo(si,ti) for si in range(len(s)) ]
    for ti in range(len(t))])

    return ax.plot_surface(X, Y, Z, rstride=1, cstride=1, facecolors=colors,cmap=cmap,
                            antialiased=True,linewidth=0,shade=True)
# In[11]:

# striping, horizontally
def pasta_stripe_2(s,t):
    colortuple = ('#DDE267','#939130'))
    return colortuple[t%2]

# striping, vertically
def pasta_stripe_1(s,t):
    colortuple = ('#DDE267','#939130'))
    return colortuple[s%2]

# uniform pasta color
def pasta_smooth(s,t):
    return '#DDE267'
# striping with different colors
def wholewheat_stripe_2(s,t):
    colortuple3=('#966C00','#D49A0D')
    return colortuple3[t%2]

# uniform green
def spinach(s,t):
    return '#285D28'

# first plot: a shell
shell = parm_plot_3d([lambda s,t: (t/(2*np.pi)*(1-0.8*s**2))*np.cos(t),
lambda s,t: (t/(2*np.pi)*(1-0.8*s**2))*np.sin(t),
lambda s,t: s +.4], srange=(-1,1,70), trange=(np.pi/6,13*np.pi/6,100),
colorinfo=pasta_stripe_1)

# second plot: Cavatappi
beta = 0.1
alpha = 0.08
gamma = 1./16

cavatappi = parm_plot_3d([lambda s,t: (beta+alpha*np.cos(s))*np.cos(t)-1.3,
lambda s,t: (beta+alpha*np.cos(s))*np.sin(t),
lambda s,t: gamma*t + alpha*np.sin(s)],
srange=(0,2*np.pi,31),trange=(-4*np.pi,2*np.pi,100),
            ax=ax,colorinfo=pasta_stripe_1)

# third plot: Penne Rigate
penne_rigate = parm_plot_3d(srange=(-.5,.5,60),
trange=(0,2*np.pi,41),
flist=[lambda s,t: 0.2+0.1*np.cos(t),
lambda s,t: 1.7 - 0.1*np.sin(t),
lambda s,t: .6 +0.1*np.sin(t) + s],
ax=ax, colorinfo=wholewheat_stripe_1)

# Farfalle

farfalle = parm_plot_3d(srange=(-np.pi,np.pi,80),
trange=np.linspace(-3,3,80),
flist=[lambda s,t: .25*t + .025*np.sin(10*s),
lambda s,t: .25 * 2*s/3 *(1.2 - 1/(1+t**2)) + 1,
lambda s,t: .25 * np.sin(np.pi*s)/(2*np.pi*s) -.8],
ax=ax, colorinfo=pasta_smooth)


fusilli = [ parm_plot_3d(
[lambda s,t: s/3 * np.cos(t-n*2*np.pi/3)+1.5,
lambda s,t: s/3 * np.sin(t-n*2*np.pi/3)+1.5,
lambda s,t: t/10 + s**2/2],
srange=(0,0.5,40), trange=(-2*np.pi,2*np.pi,100),
ax=ax, colorinfo=spinach)
for n in range(3)]

# change viewing params
# ax2.set_zlim([-1,1])
ax.set_axis_off()
fig.savefig('pasta-experiment.pdf',dpi=300,transparent=True)