import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d as a3

plt.close('all')
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.view_init(vertical_axis='y')
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_xlabel('Z')
ax.axes.set_xlim3d(left=-1, right=3.5) 
ax.axes.set_ylim3d(bottom=0, top=10) 
ax.axes.set_zlim3d(bottom=-0.5, top=3.5) 
ax.set_aspect('equal')
    
def a(elemList):
    return np.array(elemList)

def drawTriangle(centre, size=1, color='y'):
    # A neuron is a triangle
    left = centre + a([- size / 2, - size / 2, 0])
    right = centre + a([size / 2, - size / 2, 0])
    top = centre + a([0, size / 2, 0])
    coords = a([left, top, right])
    tri = a3.art3d.Poly3DCollection([coords])
    tri.set_color(color)
    tri.set_edgecolor('k')
    ax.add_collection3d(tri)

for col in range(3):
    drawTriangle(a([col, 10, 0]))