QUATERNION_RODRIGUES_ALTMANN_COMPOSITE_AXIS_ANGLE_ROTATION

'''

USE A MINIMALIST FLIGHT SIMULATOR TO TEST AXIS-ANGLE ROTATION AND ALGEBRA

# 07_10_23
# ...
# 05_01_24

# MOUSE & KEYS
#
UP KEY - INCREASE VELOCITY
DOWN KEY - BRAKE
MOUSE LEFT, RIGHT - TURN
MOUSE UP, DOWN - PITCH
LEFT, RIGHT KEYS - ROLL
PAUSE, BREAK - PAUSE
SPACE -  VELOCITY = 0
ENTER, ESC - EXIT

'''

from math import acos, degrees, sin, cos, sqrt
from numpy import arange, array, concatenate, empty, identity, zeros
from OpenGL.GL import *
from OpenGL.GLU import *
import pygame
from pygame.locals import *
import sys


def main():

    # clock = pygame.time.Clock()

    screenwidth = 1920
    screenheight = 1080
    sw2 = screenwidth/2.0
    sh2 = screenheight/2.0

    pygame.init()
    pygame.display.set_caption('ESC or RETURN exits')
    gl_init(screenwidth, screenheight, FULL_SCREEN=1)

    make_hoops1()
    make_hoops2()
    make_ground()

    pygame.mouse.set_visible(0)
    pygame.mouse.set_pos([sw2, sh2])

    PAUSE = 0
    JOYSTICK = 0

    # supply a identity matrix to a function
    I3 = identity(3)
    I4 = identity(4)
    # supply a matrix/vector to a function to be used and thrown away
    Z3 = zeros((3,3))
    Z4 = zeros((4,4))
    vec3 = empty((3))

    pos = array([0.0, 170.0, 600.0])

    velocity = 0.0
    Phi = 0.0
    Theta = 0.0
    Psi = 0.0

    R = array([1.0, 0.0, 0.0]) # RIGHT (E1)
    U = array([0.0, 1.0, 0.0]) # UP (E2)
    B = array([0.0, 0.0, 1.0]) # BACK (E3)

    glTranslate(0.0, 0.0, -5.0)

    if JOYSTICK:
        joystick = pygame.joystick.Joystick(0)
        joystick.init()

    while 1:

        KEYS = pygame.key.get_pressed()

        if KEYS[K_UP]:
            velocity += 0.08
            if velocity > 8.0:
                velocity = 8.0

        if KEYS[K_DOWN]:
            velocity -= 0.04
            if velocity < 0.0:
                velocity = 0.0

        for event in pygame.event.get():

            if event.type == QUIT:
                pygame.quit()
                return

            elif event.type == KEYDOWN:
                if event.key in [K_ESCAPE, K_RETURN]:
                    pygame.quit()
                    return
                elif event.key == K_PAUSE:
                    PAUSE ^= 1
                elif event.key == K_SPACE:
                    velocity = 0.0

        if PAUSE: continue

        glClear(GL_COLOR_BUFFER_BIT|GL_DEPTH_BUFFER_BIT)

        mousex, mousey = pygame.mouse.get_pos()

        if JOYSTICK:
            # YAW (LEFT, RIGHT)
            Phi  = joystick.get_axis(2) / 80.0
            # PITCH (UP, DOWN)
            Theta = -joystick.get_axis(1) / 80.0
            # ROLL (CW, CCW)
            Psi = -joystick.get_axis(0) / 100.0
        else:
            # YAW (LEFT, RIGHT)
            Phi = (mousex - sw2) / 1.0e5
            # PITCH (UP, DOWN)
            Theta = (mousey - sh2) / 5.0e4
            # ROLL (CW, CCW)
            Psi = 0.0
            if KEYS[K_LEFT]:
                Psi = 0.01
            if KEYS[K_RIGHT]:
                Psi = -0.01

        # COMPOSE A SINGLE ANGLE-AXIS ROTATION FROM BODY FRAME AXIS ROTATIONS
        #
        Gamma, E = composite_rotation(-B, Psi, U, Phi, vec3)
        Gamma, E = composite_rotation(E, Gamma, R, Theta, vec3)

        if 1:
            C = rotation_matrix_4x4(Gamma, E, I4, Z4)
            glMultMatrixd(C)
        else:
            glRotatef(degrees(Gamma), *E)

        R = rotate_r(R, E, -Gamma, vec3)
        U = rotate_r(U, E, -Gamma, vec3)
        B = rotate_r(B, E, -Gamma, vec3)

        glPushMatrix()

        pos += -B * velocity
        glTranslatef(*-pos)

        glColor3f(0.0, 0.0, 1.0)
        glCallList(g.hatch_ground)
        glColor3f(0.7, 0.0, 0.0)
        glCallList(g.square_hoops1)
        glColor3f(0.0, 0.7, 0.7)
        glCallList(g.square_hoops2)

        glPopMatrix()

        pygame.display.flip()
        pygame.time.wait(10)
        # clock.tick(100)


def normalize(v):
    norm = sqrt(v[0]**2 + v[1]**2 + v[2]**2)
    if norm > 0.0:
        return v / norm
    return v


def magnitude(v):
    return sqrt(v[0]**2 + v[1]**2 + v[2]**2)


def npcross(u, v):
    x = u[1]*v[2] - u[2]*v[1]
    y = u[2]*v[0] - u[0]*v[2]
    z = u[0]*v[1] - u[1]*v[0]
    return array([x,y,z])


def rotate_r(r, e, 𝚹, cross_er):

    """

 r is the position vector to rotate
 e is the unit direction vector about which to rotate
  𝚹 is the rotation angle in radians
 cross_er = empty((3))

 Does not check if e is normalized

 Rotation is right-hand direction

 #------------------------------------------------------------------

 Rodrigues finite rotation formula for numpy vector format

 Rotation Transforms for Computer Graphics, John Vince (2011a), p.139
 Classical Mechanics, Goldstein et al. p. 162 ("the rotation formula")
 Computer Graphics Software Construction, John R. Rankin (1989), p.279
 Quaternions for Computer Graphics, John Vince (2011b) pp. 80-, 123(!)
 (Demonstrates the equivalence of Quaternion rotation and Rodrigues rotation)

 In matrix form:
 Analytical Mechanics of Space Systems by Schaub & Junkins eq.3.82 (p.105)
 in video: Axis-Angles, Euler Parameters, Quaternions Matlab Examples
   by Shane Ross (RossDynamicsLab) starting at 25 minutes
 Graphics Gems (1990), p.503-5 (Patrick-Gilles Maillot)
 Do We Really Need Quaternions? by Diana Gruber (2000)
 Demonstrates the equivalence of the Quaternion matrix and the
    direction cosine matrix from Graphics Gems (1990),
    and Schaub & Junkins p.102, which can be generated by the
    Rodrigues matrix formulation:

    Crodrigues0 = I*cos𝞍 - sin𝞍*Ex + (1.0-cos𝞍)*E*E.T
    Crodrigues1 = I*cos𝞍 - sin𝞍*Ex + (1.0-cos𝞍)*[email protected]
    where:

        |  0.0  -E[2]  E[1] |
   Ex = |  E[2]  0.0  -E[0] |
        | -E[1]  E[0]  0.0  |

 Rotations, Quaternions, and Double Groups
 by Simon L. Altmann (1986 book), p.159

 Hamilton, Rodrigues, and the Quaternion Scandal
 by Simon L. Altmann (1989 paper), p.302-3
 "...Rodrigues's couples are, therefore, quaternions..."

    """

    ## 26.6 microseconds
    # np.cross(e, r)

    ## 2.35 microseconds
    # x = e[1]*r[2] - e[2]*r[1]
    # y = e[2]*r[0] - e[0]*r[2]
    # z = e[0]*r[1] - e[1]*r[0]
    # cross_er = array([x, y, z])

    ## 1.42 microseconds
    e0, e1, e2 = e[0], e[1], e[2]
    r0, r1, r2 = r[0], r[1], r[2]
    cross_er[0] = e1*r2 - e2*r1
    cross_er[1] = e2*r0 - e0*r2
    cross_er[2] = e0*r1 - e1*r0

    ## 1.16 microseconds
    # dot_er = np.dot(e, r)

    ## 829 nanoseconds
    # dot_er = e[0]*r[0] + e[1]*r[1] + e[2]*r[2]

    ## 348 nanoseconds
    dot_er = e0*r0 + e1*r1 + e2*r2

    cos𝚹 = cos(𝚹)

    return r*cos𝚹 + e*dot_er*(1.0-cos𝚹) + cross_er*sin(𝚹)

    # Altmann (1986) p.163, (same as above)
    # return r*cos𝚹  + sin(𝚹)*cross_er + (1.0-cos𝚹)*dot_er*e  # (4)
    # Vince (2011b) p.123 (same as above)
    # (1-cos𝚹)*dot_er*e + cos𝚹*r + sin𝚹*cross_er


def composite_rotation(e1, Alpha, e2, Beta, cross_e1e2):

    """
COMPOSE A SINGLE AXIS-ANGLE ROTATION FROM TWO AXIS-ANGLE ROTATIONS

Olinde Rodrigues, 1840

Rotations, Quaternions, and Double Groups
by Simon L. Altmann (1986 book), p.159

Hamilton, Rodrigues, and the Quaternion Scandal
by Simon L. Altmann (1989 paper), p.302

Analytical Mechanics of Space Systems (Fourth Edition, 2018)
by Schaub & Junkins, p.106 (3.85, 3.86)

Quaternions for Computer Graphics (2011) by John Vince, pp.90-91 (7.1, 7.2)

    """

    sinAlpha = sin(Alpha / 2.0)
    sinBeta =  sin(Beta / 2.0)
    cosAlpha = cos(Alpha / 2.0)
    cosBeta =  cos(Beta / 2.0)

    e10, e11, e12 = e1[0], e1[1], e1[2]
    e20, e21, e22 = e2[0], e2[1], e2[2]

    cross_e1e2[0] = e11*e22 - e12*e21
    cross_e1e2[1] = e12*e20 - e10*e22
    cross_e1e2[2] = e10*e21 - e11*e20

    dot_e1e2 = e10*e20 + e11*e21 + e12*e22

    sinAlpha_sinBeta = sinAlpha*sinBeta

    Gamma = 2.0 * acos(cosAlpha*cosBeta - sinAlpha_sinBeta * dot_e1e2)

    if Gamma == 0.0:
        return 0.0, cross_e1e2 #array([1.0, 0.0, 0.0])

    E = (sinAlpha*cosBeta * e1 +
         cosAlpha*sinBeta * e2 +
         sinAlpha_sinBeta * cross_e1e2) / sin(Gamma/2.0)

    return Gamma, E


def relative_rotation(Alpha, e1, Beta, e2):

    """

    FIND THE RELATIVE ORIENTATION FROM TWO "PRINCIPAL ROTATION ELEMENT SETS"

    Analytical Mechanics of Space Systems (Fourth Edition, 2018)
    by Schaub & Junkins p.106 (3.87, 3.88)

    """

    sinAlpha = sin(Alpha / 2.0)
    sinBeta =  sin(Beta / 2.0)
    cosAlpha = cos(Alpha / 2.0)
    cosBeta =  cos(Beta / 2.0)

    # 14 times faster than np.cross, 3 times faster than np.dot
    #todo if you pass cross_e1e2 to the function, it's 17 times faster
    e10, e11, e12 = e1[0], e1[1], e1[2]
    e20, e21, e22 = e2[0], e2[1], e2[2]
    x = e11*e22 - e12*e21
    y = e12*e20 - e10*e22
    z = e10*e21 - e11*e20
    cross_e1e2 = np.array([x, y, z])
    dot_e1e2 = e10*e20 + e11*e21 + e12*e22

    sinAlpha_sinBeta = sinAlpha*sinBeta

    Gamma = 2.0 * acos(cosAlpha*cosBeta + sinAlpha_sinBeta * dot_e1e2)

    if Gamma == 0.0:
        return 0.0, np.array([1.0, 0.0, 0.0])

    E = (cosBeta*sinAlpha * e1 -
           cosAlpha*sinBeta * e2 +
             sinAlpha_sinBeta * cross_e1e2) / sin(Gamma/2.0)

    return Gamma, E


def rotation_matrix_3x3(Phi, E, I3, Ex):

    '''
    CONSTRUCT 3x3 ROTATION MATRIX WITH RODRIGUES EQUATION IN MATRIX FORM
    (Wikipedia with a sign change)

    Phi = angle to rotate in radians
    E = unit axis vector
    I = np.identity(3)
    Ex = np.zeros((3,3))

    # MAKE 4x4 MATRIX
    C = rotation_matrix_3x3(Gamma, E, I3, Z3)
    C = hstack( (C, array([[0.],[0.],[0.]])) )
    C = vstack( (C, array([[0.],[0.],[0.],[1.]]).T) )

        |  0.0  -E[2]  E[1] |
   Ex = |  E[2]  0.0  -E[0] |
        | -E[1]  E[0]  0.0  |

    '''

    Ex[0,1] = -E[2]
    Ex[0,2] =  E[1]

    Ex[1,0] =  E[2]
    Ex[1,2] = -E[0]

    Ex[2,0] = -E[1]
    Ex[2,1] =  E[0]

    return I3 - sin(Phi)*Ex + (1.0-cos(Phi)) * Ex@Ex


def rotation_matrix_4x4(Phi, E, I4, Ex):

    '''
    CONSTRUCT 4x4 ROTATION MATRIX WITH RODRIGUES IN MATRIX NOTATION
    (Wikipedia with a sign change)

    Phi = angle to rotate in radians
    E = unit axis vector
    I = np.identity(4)
    Ex = np.zeros((4,4))

    The cross product matrix or "tilde matrix operator" (Schaub, Junkins)

         |  0.0  -E[2]  E[1] |
    Ex = |  E[2]  0.0  -E[0] |
         | -E[1]  E[0]  0.0  |

    '''

    Ex[0,1] = -E[2]
    Ex[0,2] =  E[1]

    Ex[1,0] =  E[2]
    Ex[1,2] = -E[0]

    Ex[2,0] = -E[1]
    Ex[2,1] =  E[0]

    if 1:

        return I4 - sin(Phi)*Ex + (1.0-cos(Phi)) * Ex@Ex

    else:

        # Eulers' representation (?)
        E = concatenate((E, [0.0])) # because 4x4
        EE_T = E*E.T # = np.outer(E,E)
        return cos(Phi) * (I4 - EE_T) - sin(Phi)*Ex + EE_T


def make_ground():

    g.hatch_ground = glGenLists(1)
    glNewList(g.hatch_ground, GL_COMPILE)

    gr = 0.0
    s = 200.0
    xmin = -15000
    xmax =  15000
    zmin = -15000
    zmax =  15000

    glLineWidth(2.0)

    glBegin(GL_LINES)
    for x in arange(xmin, xmax, s):
        for z in arange(zmin, zmax, s):
            glVertex3f(x,   gr, z)
            glVertex3f(x+s, gr, z)
            glVertex3f(x,   gr, z)
            glVertex3f(x,   gr, z+s)
    glEnd()

    glEndList()


def make_hoops1():

    g.square_hoops1 = glGenLists(1)
    glNewList(g.square_hoops1, GL_COMPILE)

    x = 0.0
    y = 300.0
    size = 50.0
    angle = 0.0
    glLineWidth(3.0)

    for z in arange(-15000., 5250., 600.):
        y += 100*sin(-z/500.0)
        x += 170*sin(-z/500.0)
        glBegin(GL_QUADS)
        glVertex3f(x-size, y-size, z)
        glVertex3f(x+size, y-size, z)
        glVertex3f(x+size, y+size, z)
        glVertex3f(x-size, y+size, z)
        glEnd()

    glEndList()


def make_hoops2():

    g.square_hoops2 = glGenLists(1)
    glNewList(g.square_hoops2, GL_COMPILE)

    z = 0.0
    y = 300.0
    size = 50.0
    angle = 0.0
    glLineWidth(3.0)

    for x in arange(-15090., 5250., 600.):
        y += 100*sin(-x/500.0)
        z += 170*sin(-x/500.0)
        glBegin(GL_QUADS)
        glVertex3f(x, y-size, z-size)
        glVertex3f(x, y-size, z+size)
        glVertex3f(x, y+size, z+size)
        glVertex3f(x, y+size, z-size)
        glEnd()

    glEndList()


def gl_init(screenwidth, screenheight, FULL_SCREEN):

    g.screen = pygame.display.set_mode((screenwidth, screenheight),
                    DOUBLEBUF|OPENGL|FULL_SCREEN*FULLSCREEN)

    # THIS IS A PROBLEM ON LENOVO
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)

    glEnable(GL_BLEND)
    # glDisable(GL_BLEND)

    glEnable(GL_LINE_SMOOTH)
    glEnable(GL_POLYGON_SMOOTH)

    glEnable(GL_DEPTH_TEST)
    glDepthFunc(GL_LESS)
    glPolygonMode(GL_FRONT_AND_BACK, GL_LINE)

    glEnable(GL_FOG)
    fogColor = (0.0, 0.0, 0.0, 1.0)
    glFogi(GL_FOG_MODE, GL_LINEAR)
    glFogfv(GL_FOG_COLOR, fogColor)
    glFogf(GL_FOG_DENSITY, 0.3)
    glHint(GL_FOG_HINT, GL_NICEST)
    glFogf(GL_FOG_START, 1000.0)
    glFogf(GL_FOG_END, 5000.0)

    glClearColor(0.0, 0.0, 0.0, 1.0)
    glClear(GL_COLOR_BUFFER_BIT|GL_DEPTH_BUFFER_BIT)

    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    gluPerspective(60.0, screenwidth/screenheight, 1.0, 10000.0)

    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()


class g: ...

if __name__ == '__main__':
    sys.exit(main())