# load arrays with coordinates of points on a helix, then plot them
# this file is unit02_draw_helix.py 

# import libraries. Python may complain about the first one, but you really do 
# need it.

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

# total number of turns the helix will make

NumberTurns = 6

# pitch (meters per turn)

pitch = 0.3

# radius

Radius = 1.0

# points to plot per turn

PointsPerTurn = 60

# create the array of angles for successive points. the arguments to linspace
# are (1) first value; (2) last value; (3) number of equally-spaced values
# to put into the array.

ThetaArray = np.linspace(0, NumberTurns*2*np.pi, NumberTurns*PointsPerTurn)

# Now get x,y,z for each point to plot.

x = np.cos(ThetaArray)
y = np.sin(ThetaArray)
z = np.linspace(0, NumberTurns*pitch, NumberTurns*PointsPerTurn)

# now create a (blank) figure so we can set some of its attributes.

fig = plt.figure()

# "gca" is "get current axes." set the projection attribute to 3D. 

ax = fig.gca(projection='3d')

# set the x, y, and z axis limits of the plot axes

ax.set_xlim(-1, 1)
ax.set_ylim(-1, 1)
ax.set_zlim(0, 2) 
 
# label the axes and give the plot a title
 
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
ax.set_title("Helix-- George Gollin") 

# now plot the helix.
 
ax.plot(x, y, z)

# now save the plot to a png (portable network graphics) file
plt.savefig("HelixPlot.png")