import numpy as np
import matplotlib.pyplot as plt

o1_parent = np.concatenate([[0], np.arange(0,4),
    [0], np.arange(5,8),
    [0], np.arange(9,12),
    [0], np.arange(13,16),
    [0], np.arange(17,20)])
w_pos = [1,1,1]
t_len = 1 # per segment
i_len = 1.5
m_len = 1.6
r_len = 1.4
l_len = 0.8
t_theta = 3*np.pi/4.
i_theta = 5*np.pi/8.
m_theta = np.pi/2.
r_theta = 3*np.pi/8.
l_theta = np.pi/4.
jitter_theta = np.pi/12.
t_pos_0 = [w_pos[0]+t_len*np.sin(t_theta),w_pos[1]+t_len*np.cos(t_theta),w_pos[2]]
t_pos_1 = [t_pos_0[0]+t_len*np.sin(t_theta+jitter_theta),t_pos_0[1]+t_len*np.cos(t_theta+jitter_theta),w_pos[2]]
t_pos_2 = [t_pos_1[0]+t_len*np.sin(t_theta),t_pos_1[1]+t_len*np.cos(t_theta),w_pos[2]]
t_pos_3 = [t_pos_2[0]+t_len*np.sin(t_theta-jitter_theta),t_pos_2[1]+t_len*np.cos(t_theta-jitter_theta),w_pos[2]]
i_pos_0 = [w_pos[0]+i_len*np.sin(i_theta),w_pos[1]+i_len*np.cos(i_theta),w_pos[2]]
i_pos_1 = [i_pos_0[0]+i_len*np.sin(i_theta+jitter_theta),i_pos_0[1]+i_len*np.cos(i_theta+jitter_theta),w_pos[2]]
i_pos_2 = [i_pos_1[0]+i_len*np.sin(i_theta),i_pos_1[1]+i_len*np.cos(i_theta),w_pos[2]]
i_pos_3 = [i_pos_2[0]+i_len*np.sin(i_theta-jitter_theta),i_pos_2[1]+i_len*np.cos(i_theta-jitter_theta),w_pos[2]]
m_pos_0 = [w_pos[0]+m_len*np.sin(m_theta),w_pos[1]+m_len*np.cos(m_theta),w_pos[2]]
m_pos_1 = [m_pos_0[0]+m_len*np.sin(m_theta+jitter_theta),m_pos_0[1]+m_len*np.cos(m_theta+jitter_theta),w_pos[2]]
m_pos_2 = [m_pos_1[0]+m_len*np.sin(m_theta),m_pos_1[1]+m_len*np.cos(m_theta),w_pos[2]]
m_pos_3 = [m_pos_2[0]+m_len*np.sin(m_theta-jitter_theta),m_pos_2[1]+m_len*np.cos(m_theta-jitter_theta),w_pos[2]]
r_pos_0 = [w_pos[0]+r_len*np.sin(r_theta),w_pos[1]+r_len*np.cos(r_theta),w_pos[2]]
r_pos_1 = [r_pos_0[0]+r_len*np.sin(r_theta+jitter_theta),r_pos_0[1]+r_len*np.cos(r_theta+jitter_theta),w_pos[2]]
r_pos_2 = [r_pos_1[0]+r_len*np.sin(r_theta),r_pos_1[1]+r_len*np.cos(r_theta),w_pos[2]]
r_pos_3 = [r_pos_2[0]+r_len*np.sin(r_theta-jitter_theta),r_pos_2[1]+r_len*np.cos(r_theta-jitter_theta),w_pos[2]]
l_pos_0 = [w_pos[0]+l_len*np.sin(l_theta),w_pos[1]+l_len*np.cos(l_theta),w_pos[2]]
l_pos_1 = [l_pos_0[0]+l_len*np.sin(l_theta+jitter_theta),l_pos_0[1]+l_len*np.cos(l_theta+jitter_theta),w_pos[2]]
l_pos_2 = [l_pos_1[0]+l_len*np.sin(l_theta),l_pos_1[1]+l_len*np.cos(l_theta),w_pos[2]]
l_pos_3 = [l_pos_2[0]+l_len*np.sin(l_theta-jitter_theta),l_pos_2[1]+l_len*np.cos(l_theta-jitter_theta),w_pos[2]]
t_pos_array = np.vstack([t_pos_0, t_pos_1, t_pos_2, t_pos_3])
i_pos_array = np.vstack([i_pos_0, i_pos_1, i_pos_2, i_pos_3])
m_pos_array = np.vstack([m_pos_0, m_pos_1, m_pos_2, m_pos_3])
r_pos_array = np.vstack([r_pos_0, r_pos_1, r_pos_2, r_pos_3])
l_pos_array = np.vstack([l_pos_0, l_pos_1, l_pos_2, l_pos_3])
pred_3D = np.vstack([w_pos, t_pos_array, i_pos_array, m_pos_array, r_pos_array, l_pos_array]).T

plt.ion()
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.gca(projection='3d')
for segment in range(pred_3D.shape[1]):
    ax.plot([pred_3D[0,segment], pred_3D[0,o1_parent[segment]]],
        [pred_3D[1,segment], pred_3D[1,o1_parent[segment]]],
        [pred_3D[2,segment], pred_3D[2,o1_parent[segment]]])