#!/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, canonicalized
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)


# To show the phenomenon, we need a large angle,
# followed by multiple small angles
# with large changes in rotation axis in between.
# 
# See https://splines.readthedocs.io/en/0.3.0/euclidean/catmull-rom-properties.html#Wrong-Tangent-Vectors
# for a similar situation in the context of Euclidean splines.

# In[4]:


rotations = [
    angles2quat(0, 0, 0),
    angles2quat(90, 0, -45),
    angles2quat(-45, 45, -90),
    angles2quat(135, -35, 90),
    angles2quat(134.5, -34, 91),
    angles2quat(134.8, -33, 89),
    angles2quat(90, 0, 0),
]

rotations_periodic = list(canonicalized(rotations + rotations[:1]))

# centripetal parameterization
angles = np.array([
    a.rotation_to(b).angle
    for a, b in zip(rotations_periodic, rotations_periodic[1:])])
t = times = np.concatenate([[0], np.cumsum(np.sqrt(angles))])


# In[5]:


cr = CatmullRom(rotations, times, endconditions='closed')


# In[6]:


sq = Squad(rotations, times)


# In[7]:


Rs = q(rotations_periodic)

qs = QuaternionicSquad(Rs, t)


# In[8]:


ani = animate_rotations({
    'CatmullRom': evaluate(cr),
    'quaternionic': evaluate(qs),
    'Squad': evaluate(sq),
})
display_animation(ani, default_mode='reflect')


# The following function tests whether a spline segment violates
# an expected property of centripetal splines:
# A point on a centripetal spline never moves towards its preceding control point
# and it never moves away from the following control point.
# 
# In the Euclidan case, "moving towards" means "reducing the Euclidean distance",
# but in case of rotations, I guess it should mean "reducing the angle".

# In[9]:


def test_segment(spline, idx):
    time_slice = np.linspace(times[idx], times[idx + 1], 100)
    first = spline.evaluate(time_slice[0])
    last = spline.evaluate(time_slice[-1])
    from_first = 0
    to_last = 999
    for time in time_slice:
        current = spline.evaluate(time)
        new_from_first = first.rotation_to(current).angle
        if new_from_first < from_first:
            print(
                'moving towards beginning:',
                np.degrees(from_first - new_from_first),
                'degrees')
        from_first = new_from_first
        new_to_last = current.rotation_to(last).angle
        if to_last < new_to_last:
            print(
                'moving away from end:',
                np.degrees(new_to_last - to_last),
                'degrees')
        to_last = new_to_last


# In[10]:


test_segment(qs, 3)


# Those values are quite small,
# but they are too large to be explained by numerical errors.

# For comparison, `splines.quaternion.CatmullRom`
# does not show any violations:

# In[11]:


test_segment(cr, 3)


# OTOH, `splines.quaternion.Squad` is horribly broken,
# as can already be seen in the animation above.
# It shows many very strong violations:

# In[12]:


test_segment(sq, 3)


# It might have to do with the given rotations only being
# a few degrees apart, but the corresponding quadrangle points being
# 80 degrees apart.

# In[13]:


[[np.degrees(q1.rotation_to(q2).angle) for q1, q2 in zip(s, s[1:])]
 for s in sq.segments]


# In[ ]: