import plotly.offline as py
import plotly.tools as tls
from plotly.graph_objs import *
import numpy as np
from numpy import pi, cos, sin, exp, log, sqrt
from random import random
py.init_notebook_mode()

zero=np.array([0,0,0])

def rotation_matrix(axis, theta):
    #based on http://stackoverflow.com/a/6802723
    """
    Return the rotation matrix associated with counterclockwise rotation about
    the given axis by theta radians.
    """
    axis = np.asarray(axis)
    theta = np.asarray(theta)
    axis = axis/sqrt(np.dot(axis, axis))
    a = cos(theta/2.0)
    b, c, d = -axis*sin(theta/2.0)
    aa, bb, cc, dd = a*a, b*b, c*c, d*d
    bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d
    return np.array([[aa+bb-cc-dd, 2*(bc+ad), 2*(bd-ac)],
                     [2*(bc-ad), aa+cc-bb-dd, 2*(cd+ab)],
                     [2*(bd+ac), 2*(cd-ab), aa+dd-bb-cc]])

def arrow(r0=zero,dr=zero,color='black',head_frac=0.1,head_ang=pi/12,head_count=7,name=''):
    r0=np.asarray(r0)
    dr=np.asarray(dr)
    a=np.cross(dr,[random() for i in range(3)]) #rotation axis perpendicular to dr
    [ra,rb]=[rotation_matrix(axis,angle) for (axis,angle) in [(a,head_ang),(dr,2*pi/head_count)]]
    piece=-head_frac*np.dot(ra,dr)
    steps=[r0,dr,piece,-piece]
    for i in range(head_count-1):
        piece=np.dot(rb,piece)
        steps += [piece,-piece]
    path,position=[],zero
    for step in steps:
        position = position + step
        path.append(position)
    lp=[list(p) for p in path]
    [x,y,z]=list(zip(*lp))
    trace = Scatter3d(x=x,y=y,z=z,mode='lines',line=Line(color=color,width=3),name=name)
    return(trace)

[u,v,w]=[np.asarray(l) for l in [[5,-1,3],[2,2,-1],[0,0,7]]]
parallelepiped=Data([arrow(r0=r0,dr=dr,color=color,name=name) for (r0,dr,color,name) in [
    (zero,u,'red','u'),(zero,v,'green','v'),(zero,w,'blue','w'),
    (w,u,'red','u'),(u,v,'green','v'),(v,w,'blue','w'),
    (v,u,'red','u'),(w,v,'green','v'),(u,w,'blue','w'),
    (v+w,u,'red','u'),(w+u,v,'green','v'),(u+v,w,'blue','w')
    ]])
py.iplot(Figure(data=parallelepiped))