def jeffery_omega(L, K, n1, n2, n3, Omega, E):
    """
    Compute Jeffery angular velocity
    
    L: (lambda^2-1)/(lambda^2+1)
    K: (kappa^2-1)/(kappa^2+1)
    n1,n2,n3: vector triplet representing current orientation
    Omega: vorticity (lab frame) 
    E: strain matrix (lab frame)

    Returns (3,) ndarray with angular velocity of particle (body frame)
    See Appendix A in http://hdl.handle.net/2077/40830
    """

    omega1 = n1.dot(Omega) + (L-K)/(L*K-1.) * (n2.dot(E.dot(n3)))
    omega2 = n2.dot(Omega) + L * (n1.dot(E.dot(n3)))
    omega3 = n3.dot(Omega) - K * (n1.dot(E.dot(n2)))

    return np.array([omega1, omega2, omega3])


def jeffery_numerical(L, K, q0, Omega, E, max_t = None, dt = 1e-3):
    """
    Integrate one trajectory according to Jeffery's equations.
    
    L: (lambda^2-1)/(lambda^2+1) shape parameter 1
    K: (kappa^2-1)/(kappa^2+1) shape parameter 2  
    q0: quaternion representing initial orientation
    Omega: vorticity (lab frame) 
    E: strain matrix (lab frame)
    max_t: Max time of trajectory, defaults to 2 Jeffery periods based on L
    dt: Integration timestep
    
    See Appendix A in https://arxiv.org/abs/1705.06997 for quaternion convention.
    
    Returns (ts, qs, n2s, n3s) where
        ts is (N,1) ndarray with timestamps (starting at 0) for N steps
        qs is (N,4) ndarray with orientations (quaternions) for N steps
        n2s is (N,3) ndarray with n2 vector for N steps
        n3s is (N,3) ndarray with n3 vector for N steps        
    """
    
    if max_t is None:
        maxKL = max(abs(L),abs(K))
        jeffery_T = 4*np.pi/np.sqrt(1-maxKL*maxKL)
        max_t = 2*jeffery_T
    
    N = int(max_t/dt)

    ts = np.zeros((N,1))
    n2s = np.zeros((N,3))
    n3s = np.zeros((N,3))    
    qs = np.zeros((N,4))
    q = q0
    t=0
    for n in range(N):
        R = q.get_R()
        n1 = R[:,0]
        n2 = R[:,1]
        n3 = R[:,2]        
        
        ts[n] = n*dt
        n2s[n,:] = n2
        n3s[n,:] = n3
        qs[n,:] = q.q
        
        omega = jeffery_omega(L, K, n1, n2, n3, Omega, E)
        qdot = 0.5 * omega.dot(q.get_G())
        q = q + dt*qdot
        q.normalize()
    
    return ts, qs, n2s, n3s

def jeffery_axisymmetric_exact(L, q0, Omega, E, max_t = None, dt = 1e-1):
    """
    Generate one exact trajectory for axisymmetric particle ('Jeffery orbit')
    
    L: (lambda^2-1)/(lambda^2+1) shape parameter
    q0: quaternion representing initial orientation
    Omega: vorticity (lab frame) 
    E: strain matrix (lab frame)
    max_t: Max time of trajectory, defaults to 2 Jeffery periods based on L
    dt: Sample spacing
    
    See Appendix A in https://arxiv.org/abs/1705.06997 for quaternion convention.
    
    Returns (ts, qs, n2s, n3s) where
        ts is (N,1) ndarray with timestamps (starting at 0) for N steps
        n3s is (N,3) ndarray with n3 vector for N steps        
    """    
    if max_t is None:
        jeffery_T = 4*np.pi/np.sqrt(1-L*L)
        max_t = 2*jeffery_T

    N = int(max_t/dt)    
        
    levi_civita = np.zeros((3, 3, 3))
    levi_civita[0, 1, 2] = levi_civita[1, 2, 0] = levi_civita[2, 0, 1] = 1
    levi_civita[0, 2, 1] = levi_civita[2, 1, 0] = levi_civita[1, 0, 2] = -1        
    O = -np.einsum('ijk,k',levi_civita, Omega)
 
    B = O + L*E
    n30 = q0.get_R().dot(np.array([0,0,1]))
    
    ts = np.zeros((N,1))
    n3s = np.zeros( (N,3) )
    for n in range(N):

        t = dt*n

        M = scipy.linalg.expm(B*t)
        n3 = M.dot(n30)
        n3 = n3/np.linalg.norm(n3)

        ts[n] = t
        n3s[n,:] = n3
    
    return (ts, n3s)

"""
Explanation: Functions for angular velocity & integration
The particle is an ellipsoid. The reference state (corresponding to no rotation) is that the ellipsoid is axis-aligned and the axis lengths are (a_x, a_y, a_z). The shape parameters in the code below are
l = a_z/a_x
k = a_y/a_x
Its orientation is represented by the rotation (as a Quaternion) from the reference state. See Appendix A of https://arxiv.org/abs/1705.06997 for the quaternion convention.
End of explanation
"""


Omega = np.array([0,0,-.5])
E = np.array([
    [0,.5,0],
    [.5,0,0],
    [0,0,0]
])

"""
Explanation: Omega & E (strain) for simple shear flow
End of explanation
"""


angles = np.pi/2 * np.linspace(0.05,1,5)

## first test is axisymmetric along n3 (K=0)
ax = plt.subplot(1,2,1)
for angle in angles:
    q0 = Quaternion(axis=[0,1,0], angle=angle)

    l = 7
    k = 1
    L = (l**2-1)/(l**2+1)
    K = (k**2-1)/(k**2+1)

    (ts, qs, n2s, n3s) = jeffery_numerical(L, K, q0, Omega, E)
    ax.plot(n3s[:,0],n3s[:,1],ls='solid', color='C0')

    (ts, n3s) = jeffery_axisymmetric_exact(L,q0,Omega,E)
    ax.plot(n3s[:,0],n3s[:,1],ls=(0, (5, 10)),color='C1')

ax.set_xlim(-1.1,1.1)
ax.set_ylim(-1.1,1.1)
ax.set_aspect('equal')    

## second test is axisymmetric along n2 (L=0)
ax = plt.subplot(1,2,2)
for angle in angles:
    q0_tri = Quaternion(axis=[1,0,0], angle=-angle)
    q0_axi = Quaternion(axis=[1,0,0], angle=np.pi/2-angle)
    l = 1
    k = 7
    L = (l**2-1)/(l**2+1)
    K = (k**2-1)/(k**2+1)

    (ts, qs, n2s, n3s) = jeffery_numerical(L, K, q0_tri, Omega, E)
    ax.plot(n2s[:,0],n2s[:,1],ls='solid', color='C0')

    (ts, n3s) = jeffery_axisymmetric_exact(K,q0_axi,Omega,E)
    ax.plot(n3s[:,0],n3s[:,1],ls=(0, (5, 10)),color='C1')

ax.set_xlim(-1.1,1.1)
ax.set_ylim(-1.1,1.1)
ax.set_aspect('equal')    

plt.show()
    

"""
Explanation: Validate code against axisymmetric case (Jeffery orbits)
End of explanation
"""


rot1=Quaternion(axis=[0,0,1], angle=0.1*np.pi/2) # this sets psi
rot2=Quaternion(axis=[1,0,0], angle=np.pi/2-0.1) # this sets theta
q0 = rot1.mul(rot2)
max_t = 300

fig = plt.figure(figsize=(15,8))

l = 7
k = 1
L = (l**2-1)/(l**2+1)
K = (k**2-1)/(k**2+1)

ax = fig.add_subplot(1,2,1, projection='3d')

(ts, qs, n2s, n3s) = jeffery_numerical(L, K, q0, Omega, E, max_t = max_t)
ax.plot(n3s[:,0],n3s[:,1],n3s[:,2],ls='solid', color='C0')

ax.set_xlim(-1.1,1.1)
ax.set_ylim(-1.1,1.1)
ax.set_zlim(-1.1,1.1)
ax.set_aspect('equal') 
ax.set_title('l={:.2f} | k={:.2f}'.format(l,k))

ax = fig.add_subplot(1,2,2, projection='3d')
l = 7
k = 1.2
L = (l**2-1)/(l**2+1)
K = (k**2-1)/(k**2+1)

(ts, qs, n2s, n3s) = jeffery_numerical(L, K, q0, Omega, E, max_t = max_t)
ax.plot(n3s[:,0],n3s[:,1],n3s[:,2],ls='solid', color='C0')

ax.set_xlim(-1.1,1.1)
ax.set_ylim(-1.1,1.1)
ax.set_zlim(-1.1,1.1)
ax.set_aspect('equal') 
ax.set_title('l={:.2f} | k={:.2f}'.format(l,k))
plt.show()

"""
Explanation: Case 1: Axisymmetric (1,1,7) vs slightly asymmetric (1,1.2,7)
Side-by-side comparison between two slightly different particles started in the same initial condition.
A: initial condition in integrable region
See Fig 3.11 in Jonas' thesis for definitions of psi & theta.
These initial conditions are inside the integrable region, so the difference between symmetric and asymmetric particle is bounded.
End of explanation
"""


rot1=Quaternion(axis=[0,0,1], angle=0.95*np.pi/2) # this sets psi
rot2=Quaternion(axis=[1,0,0], angle=np.pi/2-0.1) # this sets theta
q0 = rot1.mul(rot2)
max_t = 300

fig = plt.figure(figsize=(15,8))

l = 7
k = 1
L = (l**2-1)/(l**2+1)
K = (k**2-1)/(k**2+1)

ax = fig.add_subplot(1,2,1, projection='3d')

(ts, qs, n2s, n3s) = jeffery_numerical(L, K, q0, Omega, E, max_t = max_t)
ax.plot(n3s[:,0],n3s[:,1],n3s[:,2],ls='solid', color='C0')
np.savetxt('symmetric-l7-k1.csv', np.hstack((ts,qs)), delimiter=',') # export!

ax.set_xlim(-1.1,1.1)
ax.set_ylim(-1.1,1.1)
ax.set_zlim(-1.1,1.1)
ax.set_aspect('equal') 
ax.set_title('l={:.2f} | k={:.2f}'.format(l,k))

ax = fig.add_subplot(1,2,2, projection='3d')
l = 7
k = 1.2
L = (l**2-1)/(l**2+1)
K = (k**2-1)/(k**2+1)

(ts, qs, n2s, n3s) = jeffery_numerical(L, K, q0, Omega, E, max_t = max_t)
ax.plot(n3s[:,0],n3s[:,1],n3s[:,2],ls='solid', color='C0')
np.savetxt('asymmetric-l7-k1-2.csv', np.hstack((ts,qs)), delimiter=',') # export!

ax.set_xlim(-1.1,1.1)
ax.set_ylim(-1.1,1.1)
ax.set_zlim(-1.1,1.1)
ax.set_aspect('equal') 
ax.set_title('l={:.2f} | k={:.2f}'.format(l,k))
plt.show()


"""
Explanation: B: initial condition in chaotic region
See Fig 3.11 in Jonas' thesis for definitions of psi & theta.
These initial conditions are inside the chaotic region, so the difference between symmetric and asymmetric particle is more pronounced.
End of explanation
"""


