#!/usr/bin/env python
# coding: utf-8

# This notebook compares various approaches to creating splines for quaternions, as discussed in [this issue](https://github.com/moble/quaternionic/issues/41).  I have re-implemented `Squad` as @mgeier's [suggestion for `quaternionic`](https://github.com/moble/quaternionic/pull/42), because it was wrong in `splines`.
# 
# Also, I strenuously object to automatically "unflipping" rotors, but for the sake of comparison, this is what I mostly do below.

# # Code loading
# Highlight this cell and press "y" to convert to python code, then run it if
# you need to install `numpy`, `matplotlib`, `quaternionic`, and/or `splines`.

import sys
from pathlib import Path

executable = str(Path(sys.executable).resolve())

! {executable} -m pip install numpy matplotlib quaternionic splines
# In[1]:


import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import CubicSpline

import quaternionic
from splines.quaternion import CatmullRom, UnitQuaternion
from helper import angles2quat, animate_rotations, display_animation


# In[2]:


# Some helpful utilities for converting between `splines` and `quaternionic`

def uq(q):
    return UnitQuaternion.from_unit_xyzw(q.vector.tolist() + [q.w[()],])

def q(uq):
    try:
        iterator = iter(uq)
    except TypeError:
        return quaternionic.array(uq.wxyz)
    else:
        return quaternionic.array([uqi.wxyz for uqi in uq])

class QuaternionicSquad:
    def __init__(self, Rs, t):
        self.Rs = Rs
        self.grid = t
    def evaluate(self, t):
        try:
            iterator = iter(t)
        except TypeError:
            return uq(quaternionic.squad(self.Rs, self.grid, np.array([t])[0]))
        else:
            return list(map(uq, quaternionic.squad(self.Rs, self.grid, np.array(t))))


# In[3]:


def evaluate(spline, frames=200):
    times = np.linspace(
        spline.grid[0], spline.grid[-1], frames, endpoint=False)
    return spline.evaluate(times)


# In[4]:


def squad(R_in, t_in, t_out):
    from functools import partial
    from quaternionic import slerp

    roll = partial(np.roll, axis=0)

    if R_in.size == 0 or t_out.size == 0:
        return quaternionic.array(np.zeros((0, 4)))

    R_in = quaternionic.array(R_in)
    i_in_for_out = t_in.searchsorted(t_out, side='right')-1
    t_in_broadcast_shape = t_in.shape + (1,)*len(R_in.shape[1:])
    t_out_broadcast_shape = t_out.shape + (1,)*len(R_in.shape[1:])
    A = R_in * np.exp(
        (
            - np.log((~R_in) * roll(R_in, -1))
            + np.log((~roll(R_in, 1)) * R_in)
            * np.reshape((roll(t_in, -1) - t_in) / (t_in - roll(t_in, 1)), t_in_broadcast_shape)[..., 0]
         ) * (roll(t_in, -1) - t_in) / (2 * (roll(t_in, -1) - roll(t_in, 1)))
    )
    B = roll(R_in, -1) * np.exp(
        (
            np.log((~roll(R_in, -1)) * roll(R_in, -2))
            * np.reshape((roll(t_in, -1) - t_in) / (roll(t_in, -2) - roll(t_in, -1)), t_in_broadcast_shape)[..., 0]
            - np.log((~R_in) * roll(R_in, -1))
        ) * -(roll(t_in, -1) - t_in) / (2 * (roll(t_in, -2) - t_in))
    )
    A[0] = R_in[0]
    A[-1] = R_in[-1]
    B[-2] = R_in[-1]
    B[-1] = R_in[-1] * (~R_in[-2]) * R_in[-1]
    R_ip1 = roll(R_in, -1)
    R_ip1[-1] = R_in[-1]*(~R_in[-2])*R_in[-1]
    R_ip1 = quaternionic.array(R_ip1[i_in_for_out])
    t_inp1 = roll(t_in, -1)
    t_inp1[-1] = t_in[-1] + (t_in[-1] - t_in[-2])
    tau = np.reshape((t_out - t_in[i_in_for_out]) / ((t_inp1 - t_in)[i_in_for_out]), t_out_broadcast_shape)[..., 0]

    R_out = slerp(
        slerp(R_in[i_in_for_out], R_ip1, tau),
        slerp(A[i_in_for_out], B[i_in_for_out], tau),
        2*tau*(1-tau)
    )

    return R_out

class Squad:
    def __init__(self, Rs, t):
        self.Rs = quaternionic.unflip_rotors(Rs)
        self.grid = t
    def evaluate(self, t):
        try:
            iterator = iter(t)
        except TypeError:
            return uq(squad(self.Rs, self.grid, np.array([t])[0]))
        else:
            return list(map(uq, squad(self.Rs, self.grid, np.array(t))))


# In[5]:


rotations = [
    angles2quat(0, 0, 0),
    angles2quat(90, 0, -45),
    angles2quat(-45, 45, -90),
    angles2quat(135, -35, 90),
    angles2quat(90, 0, 0),
]

# `splines` treats the rotations as periodic
rotations_periodic = rotations + rotations[:1]

Rs = q(rotations_periodic) # quaternionic.array([np.array(q.wxyz) for q in rotations_periodic])

## I'm gonna make these times a little more non-uniform to accentuate things
#times = 0, 0.75, 1.2, 2, 3.5, 4
times = 0, 0.15, 1.2, 2, 3.9, 4
t = np.array(times)

#sq = Squad(rotations)
cr = CatmullRom(rotations, endconditions='closed')

#sq2 = Squad(rotations, alpha=0.5)
cr2 = CatmullRom(rotations, alpha=0.5, endconditions='closed')

sq3 = Squad(Rs, t)
cr3 = CatmullRom(rotations, times, endconditions='closed')


# # Flipped rotors

# The first thing I note is that `sq3` and `cr3` do not actually reproduce the input data:

# In[6]:


sq3.evaluate(t) == rotations_periodic


# In[7]:


all(cr3.evaluate(times) == rotations_periodic)


# The reason is that there is a sign flip in the middle.

# In[8]:


sq3.evaluate(times)


# In[9]:


rotations_periodic


# This looks a lot like what would happen if I applied [`unflip_rotors`](https://github.com/moble/quaternionic/blob/074b626d0c63aa78479ff04ed41638931ca6693a/quaternionic/interpolation.py#L28) to the input:

# In[10]:


quaternionic.unflip_rotors(Rs)


# Let's take a look at what my code does.

# In[11]:


qs = QuaternionicSquad(Rs, t)
qs2 = QuaternionicSquad(quaternionic.unflip_rotors(Rs), t)


# Here, `qs` reproduces the input data exactly

# In[12]:


qs.evaluate(times) == rotations_periodic


# whereas `qs2` reproduces the results of `sq3` exactly:

# In[13]:


qs2.evaluate(times) == sq3.evaluate(times)


# So it looks as though `splines` effectively unflips its rotors on input.
# 
# You can actually see this difference quite clearly in the animation comparing `quaternionic` (on the left) to `splines` (on the right):

# In[14]:


ani = animate_rotations({
    'quaternionic': evaluate(qs),
    'Squad': evaluate(sq3),
    'unflipped': evaluate(qs2),
    'CatmullRom': evaluate(cr3),
})
display_animation(ani, default_mode='loop')


# Towards the middle of this animation, `quaternionic.squad` is actually rotating in the "wrong" direction, compared to the others.  If I "unflip" the input before evaluating `quaternionic.squad`, I get more nearly consistent rotations.  As mentioned above, I believe it is a bad idea to "unflip" automatically, but for the sake of comparison, I'll stick with this in what follows.
# 
# It's interesting to note that in the very first steps of each of these animations, it looks `CatmullRom` has actually started off going the *wrong* direction.  This will show up below in the angular velocity.

# # Angular velocity

# Now let's interpolate everything to a *much* finer time step, and evaluate the angular velocities:

# In[15]:


T = np.linspace(times[0], times[-1], num=20_000)

SQ = q(sq3.evaluate(T))
CR = q(cr3.evaluate(T))
QS = quaternionic.squad(quaternionic.unflip_rotors(Rs), t, T)

ωSQ = SQ.to_angular_velocity(T)
ωCR = CR.to_angular_velocity(T)
ωQS = QS.to_angular_velocity(T)


# In[16]:


fig = plt.figure()
for t_i in times:
    plt.axvline(t_i, ls="dashed", c="black")
plt.plot(T, np.linalg.norm(ωSQ, axis=1), label="Squad")
plt.plot(T, np.linalg.norm(ωCR, axis=1), label="Catmull-Rom")
plt.plot(T, np.linalg.norm(ωQS, axis=1), label="quaternionic.squad ∘ unflip")
plt.xlabel("Time")
plt.ylabel(r"$|\vec{\omega}|$")
plt.legend(loc="upper center")
plt.savefig("angular_velocities.png");


# A few notes:
# 
# * At the very beginning, the `CatmullRom` curve has this quick little dip, which coincides with the period where it appears to be going the wrong way in the animation above.
# * `Squad` has a pretty bizarre little jump just after the first transition, which is obvious in the animation above.  This looks totally unphysical to me.
# 
# I also want to look a little more closely at that last transition.

# In[17]:


fig.get_axes()[0].set_xlim(3.8, 4.0)
fig


# Note that the `to_angular_velocity` method used above isn't really specialized for this kind of problem, so the very sharp wiggles near the transition are artifacts of that method, which change with time-step size, and could possibly go away with more careful treatment.  But the general discontinuities should be trustworthy.  All three curves are discontinuous, but `quaternionic.squad` curves are less discontinuous than `CatmullRom`, which is less discontinuous than `Squad`.

# # Torque

# One criterion for a "good" spline that has appeared in the literature is the minimization of "torque", which is just taken to be the derivative of the angular velocity.  There are two important classes of torque, one is the "impulsive" torque described by discontinuities like those seen above.  By this measure, `quaternionic.squad` is the best of the three splines examined here, while `Squad` is the worst.  The other important class of torque is continuous torque, which we examine now.
# 
# Obviously, these results are only valid for the particular example set of rotors chosen above, but they are suggestive of a broader trend.

# In[18]:


ω̇SQ = CubicSpline(T, ωSQ).derivative()(T)
ω̇CR = CubicSpline(T, ωCR).derivative()(T)
ω̇QS = CubicSpline(T, ωQS).derivative()(T)


# First, we'll look at the average magnitude of the torque over the entire time span:

# In[19]:


αSQ = CubicSpline(T, np.linalg.norm(ω̇SQ, axis=1)).integrate(T[0], T[-1]) / (T[-1]-T[0])
αCR = CubicSpline(T, np.linalg.norm(ω̇CR, axis=1)).integrate(T[0], T[-1]) / (T[-1]-T[0])
αQS = CubicSpline(T, np.linalg.norm(ω̇QS, axis=1)).integrate(T[0], T[-1]) / (T[-1]-T[0])

αSQ, αCR, αQS, αSQ/αQS, αCR/αQS


# (These numbers seem pretty well converged; as I vary the number of time steps over which these are evaluated, the results only change by ~0.1%.)
# 
# This shows that `quaternionic.squad` has the lowest average torque, followed by `Squad`, and then `CatmullRom`.
# 
# Now let's look at the torque as functions of time:

# In[20]:


fig = plt.figure()
for t_i in times:
    plt.axvline(t_i, ls="dashed", c="black")
plt.semilogy(T, np.linalg.norm(ω̇SQ, axis=1), label="Squad")
plt.semilogy(T, np.linalg.norm(ω̇CR, axis=1), label="Catmull-Rom")
plt.semilogy(T, np.linalg.norm(ω̇QS, axis=1), label="quaternionic.squad ∘ unflip")
plt.xlabel("Time")
plt.ylabel(r"$|\dot{\vec{\omega}}|$")
plt.ylim(ymax=650, ymin=1)
plt.legend(loc="upper right");
#plt.savefig("angular_velocities.png");


# In[21]:


fig.get_axes()[0].set_xlim(3.85, 4.0)
fig


# In[22]:


fig.get_axes()[0].set_xlim(0.0, 0.3)
fig


# As we would expect, that bizarre behavior after the first transition shows up as a huge bump in the `Squad` curve.  Other than that bump and the larger discontinuity in $|\vec{\omega}|$ near the final transition, `Squad` generally does better than the others, but this advantage seems to come at the cost of larger torques in certain regions — which, on balance, gives `Squad` a disadvantage compared to `quaternionic.squad`.

# In[ ]:





# In[ ]:





# # Random quaternions

# In[23]:


N = 5_000

np.random.seed(1234)
random_rotors = quaternionic.array(np.random.randn(N, 4)).normalized  # Simplest way to get uniform random unit quaternions
rotations2 = [uq(R) for R in quaternionic.unflip_rotors(random_rotors)]
rotations2_periodic = rotations2 + rotations2[:1]
Rs2 = q(rotations2_periodic)

minδt = 0.05
t2 = np.cumsum(minδt+(1-minδt)*np.random.rand(N+1))
times2 = t2.tolist()


# In[24]:


T2 = np.linspace(times2[0], times2[-1], num=100*N)

SQ2 = q(Squad(Rs2, t2).evaluate(T2))
CR2 = q(CatmullRom(rotations2, times2, endconditions='closed').evaluate(T2))
QS2 = quaternionic.squad(Rs2, t2, T2)

ωSQ2 = SQ2.to_angular_velocity(T2)
ωCR2 = CR2.to_angular_velocity(T2)
ωQS2 = QS2.to_angular_velocity(T2)

ω̇SQ2 = CubicSpline(T2, ωSQ2).derivative()(T2)
ω̇CR2 = CubicSpline(T2, ωCR2).derivative()(T2)
ω̇QS2 = CubicSpline(T2, ωQS2).derivative()(T2)

αSQ2 = CubicSpline(T2, np.linalg.norm(ω̇SQ2, axis=1)).integrate(T2[0], T2[-1]) / (T2[-1]-T2[0])
αCR2 = CubicSpline(T2, np.linalg.norm(ω̇CR2, axis=1)).integrate(T2[0], T2[-1]) / (T2[-1]-T2[0])
αQS2 = CubicSpline(T2, np.linalg.norm(ω̇QS2, axis=1)).integrate(T2[0], T2[-1]) / (T2[-1]-T2[0])

αSQ2, αCR2, αQS2, αSQ2/αQS2, αCR2/αQS2


# In[ ]: