#!/usr/bin/env python
# coding: utf-8
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

import quaternionic
from splines.quaternion import Squad, 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]:


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(rotations, times)
cr3 = CatmullRom(rotations, times, endconditions='closed')


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

# In[5]:


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


# In[6]:


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


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

# In[7]:


sq3.evaluate(times)


# In[8]:


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[9]:


quaternionic.unflip_rotors(Rs)


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

# In[10]:


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


# Here, `qs` reproduces the input data exactly

# In[11]:


qs.evaluate(times) == rotations_periodic


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

# In[12]:


all(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[13]:


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


# Towards the middle of this animation, the two are actually rotating in opposite directions!
# 
# If I unflip the input before evaluating `quaternionic.squad`, I get more nearly consistent rotations:

# In[14]:


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


# 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 be borne out below in the 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(Rs, t, T)
QS2 = 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)
ωQS2 = QS2.to_angular_velocity(T)


# In[16]:


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")
plt.semilogy(T, np.linalg.norm(ωQS2, axis=1), label="quaternionic.squad ∘ unflip")
plt.xlabel("Time")
plt.ylabel(r"$|\vec{\omega}|$")
plt.legend(loc="upper right")
plt.savefig("angular_velocities.png");


# A few notes:
# 
# * I had to switch to a log scale on the vertical axis to show what's going on with `Squad` while still being able to see details about the other curves.
# * In general, the `Squad` curve is all over the place — discontinuous across boundaries, and having much larger variations than most of the other curves.
# * 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 animations.
# * Around the middle, `CatmullRom` looks a bit smoother than `quaternionic.squad`, but this appears to be consistent with the fact that we've made less work for `CatmullRom` by "unflipping" the input.  In fact, if we manually "unflip" the input to `quaternionic.squad`, we get another curve that looks smoother like `CatmullRom`.
# 
# 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


# Again, the `Squad` result is crazy, but all four curves are at least a little discontinuous.  Still, the `quaternionic.squad` curves are less discontinuous than `CatmullRom`.
# 
# Note that the `to_angular_velocity` method used above isn't really specialized for this kind of problem, so the very sharp wiggles in the red, green, and orange curves may be artifacts of that method, which could possibly go away with more careful treatment.

# In[ ]: