# Number of rotations to generate.
ROTATIONS = 100
# Number of in between values to interpolate.
BETWEEN = 8

# Random seed.
SEED = 123

import bpy

from random import uniform, seed, randint
from math import pi

from mathutils import Matrix, Vector, Quaternion, Euler


def random_qt():
    x = Vector(
        (uniform(-1.0, 1.0), uniform(-1.0, 1.0), uniform(-1.0, 1.0))
    ).normalized()
    while True:
        y = Vector(
            (uniform(-1.0, 1.0), uniform(-1.0, 1.0), uniform(-1.0, 1.0))
        ).normalized()
        z = x.cross(y)
        if z.length_squared > 0.1:
            z.normalize()
            break
    x = z.cross(y).normalized()
    m = Matrix([x, y, z])

    return m.to_quaternion()


def eulers_compat(eulers):
    e_prev = Euler((0, 0, 0))
    eulers_compat = []
    for e in eulers:
        e_copy = e.copy()
        e_copy.make_compatible(e_prev)
        e_prev = e_copy
        eulers_compat.append(e_copy)
    return eulers_compat


def wrap_eulers(eulers):
    eulers_wrapped = []
    for e in eulers:
        e_copy = e.copy()
        for axis in range(3):
            e_copy[axis] += randint(-10, 10) * (pi * 2.0)
        eulers_wrapped.append(e_copy)
    return eulers_wrapped


def calc_error(eulers):
    error = 0.0
    e_prev = eulers[0]
    for e in eulers:
        error += abs(e[0] - e_prev[0]) + abs(e[1] - e_prev[1]) + abs(e[2] - e_prev[2])
        e_prev = e
    return error / (ROTATIONS * BETWEEN)


def plot_image(eulers):
    import matplotlib.pyplot as plt

    plt.style.use("_mpl-gallery")
    t = list(range(len(eulers)))
    rgb = ("#ff0000", "#00ff00", "#0000ff")

    fig, ax = plt.subplots()

    for axis in range(3):
        x = [e[axis] for e in eulers]
        ax.set_title("XYZ"[axis])
        ax.plot(t, x, linewidth=0.2, color=rgb[axis])

    ax.set(xticks=range(0, len(eulers), BETWEEN + 1), yticks=range(1, 8))

    plt.savefig("/tmp/myfig.png", dpi=500)


def main():
    seed(SEED)
    quats = [random_qt()]
    for i in range(ROTATIONS):
        q_prev = quats[-1]
        q = random_qt()
        for j in range(1, BETWEEN + 1):
            fac = j / (BETWEEN + 1)
            quats.append(q_prev.slerp(q, fac))
        quats.append(q)

    eulers = [q.to_euler() for q in quats]

    eulers = wrap_eulers(eulers)

    print("Interpolated eulers", calc_error(eulers_compat(eulers)))

    seed(SEED)
    eulers = [
        Euler(
            (
                uniform(-pi, pi) * randint(-10, 10),
                uniform(-pi, pi) * randint(-10, 10),
                uniform(-pi, pi) * randint(-10, 10),
            )
        )
        for i in range(ROTATIONS * BETWEEN)
    ]

    print("Random eulers", calc_error(eulers_compat(eulers)))

    # plot_image(eulers_nice)


if __name__ == "__main__":
    main()