# -*- coding: utf-8 -*-
"""
Created on Thu Feb 13 15:07:12 2025

@author: Lee Etchels

A collection of key functions and code snippets used for the associated study.
Full codes required original raw data files with different formatting etc,
and were otherwise not especially readable or reusable without direct 
assistance.

Hopefully there are no errors from collating them like this but it would be
best to verify before use.
"""

import numpy as np
from pyquaternion import Quaternion
from math import degrees, atan, acos, atan2

def rot_vec_by_three_angles(v0, r1, ax1, r2, ax2, r3, ax3):
    """
    Rotate a vector by euler sequence.

    Parameters
    ----------
    v0 : list
        Containing three elements representing the initial vector [x,y,z].
    r1 : numeric
        Angle in degrees for first rotation in sequence.
    ax1 : list
        Containing three elements representing the first global coordinate
        axis vector from the Euler sequence [x,y,z].
    r2 : numeric
        Angle in degrees for second rotation in sequence.
    ax2 : list
        Containing three elements representing the second global coordinate
        axis vector from the Euler sequence [x,y,z].
    r3 : numeric
        Angle in degrees for third rotation in sequence.
    ax3 : list
        Containing three elements representing the second global coordinate
        axis vector from the Euler sequence [x,y,z].

    Returns
    -------
    v3 : list
        Initial vector rotated by the three angles in the sequence specified.
    x3 : list
        The first global coordinate axis vector from the Euler sequence
        rotated by the three angles in the sequence specified.
    y3 : list
        The second global coordinate axis vector from the Euler sequence
        rotated by the three angles in the sequence specified.
    z3 : list
        The third global coordinate axis vector from the Euler sequence
        rotated by the three angles in the sequence specified.

    """
    x0, y0, z0 = ax1, ax2, ax3
    q1 = Quaternion(axis=x0, degrees=r1)
    x1, y1, z1 = q1.rotate(x0), q1.rotate(y0), q1.rotate(z0)
    v1 = q1.rotate(v0)
    q2 = Quaternion(axis=y1, degrees=r2)
    x2, y2, z2 = q2.rotate(x1), q2.rotate(y1), q2.rotate(z1)
    v2 = q2.rotate(v1)
    q3 = Quaternion(axis=z2, degrees=r3)
    v3 = q3.rotate(v2)
    x3, y3, z3 = q3.rotate(x2), q3.rotate(y2), q3.rotate(z2)
    return v3, x3, y3, z3

def angle_between_vectors(v1, v2):
    """
    Calculate angle between two vectors.
    
    Includes catching error caused by rounding etc leading to a0 being
    slightly > 1 or < -1. Large corrections should be printed to the terminal
    for user evaluation.

    Parameters
    ----------
    v1 : list or np.array
        First vector.
    v2 : list or np.array
        Second vector.

    Returns
    -------
    a : float
        angle in degrees between the two vectors.

    """
    if type(v1) is list:
        v1 = np.array(v1, dtype=np.float32)
    if type(v2) is list:
        v2 = np.array(v2, dtype=np.float32)
    a0 = np.dot(v1, v2) / (np.linalg.norm(v1)
                           * np.linalg.norm(v2))
    if a0 > 1:
        if a0 > 1.01:
            print('a0 > 1 changed to 1, a0 = ' + str(a0))
        a0 = 1
    if a0 < -1:
        if a0 < -1.01:
            print('a0 < -1 changed to -1, a0 = ' + str(a0))
        a0 = -1
    a = np.degrees(np.arccos(a0))
    return a

def vec_from_rirv(ri, rv):
    """
    Calculate cup normal vector from RI and RV.

    Parameters
    ----------
    ri : float
        Radiographic inclination in degrees.
    rv : float
        Radiographic version in degrees.

    Returns
    -------
    imp_vec_r : list
        Cup face normal vector representing a cup implanted at the given
        RI/RV at the right hip.
    imp_vec_l : list
        Cup face normal vector representing a cup implanted at the given
        RI/RV at the left hip.

    """
    imp_vec_r = np.array([0, 0, -1], dtype=np.float32)
    csys = [np.array([1, 0, 0], dtype=np.float32),
            np.array([0, 1, 0], dtype=np.float32),
            np.array([0, 0, 1], dtype=np.float32)]
    q1 = Quaternion(axis=csys[1], degrees=(ri)*-1)
    imp_vec_r = q1.rotate(imp_vec_r)
    csys[0], csys[2] = q1.rotate(csys[0]), q1.rotate(csys[2])
    q2 = Quaternion(axis=csys[0], degrees=rv)
    imp_vec_r = q2.rotate(imp_vec_r)
    imp_vec_l = imp_vec_r * np.array([-1, 1, 1], dtype=np.float32)
    csys[1], csys[2] = q2.rotate(csys[1]), q2.rotate(csys[2])
    return imp_vec_r, imp_vec_l

def ri_rv_av__ov_from_vec(v):
    """
    Calculate the cup orientation from the cup normal vector using
    radiographic inclination, radiographic version, anatomic version,
    and operative version. Use vector for cup at right hip and then mirror if
    required at the left hip.

    Parameters
    ----------
    v : list
        Cup normal vector.

    Returns
    -------
    ri : float
        Cup radiographic inclination.
    rv : float
        Cup radiographic version.
    av : float
        Cup anatomic version.
    ov : float
        Cup operative version.

    """
    x, y, z = v
    ri = degrees(atan(x/-z))
    av = degrees(atan(y/x))
    ov = degrees(atan2(y, -1 * z))
    a = [x, 0, z]
    b = [x, y, z]
    dot_ab = (a[0] * b[0]) + (a[1] * b[1]) + (a[2] * b[2])
    mod_a = np.linalg.norm(a)
    mod_b = np.linalg.norm(b)
    rv = degrees(acos(dot_ab / (mod_a * mod_b)))
    return ri, rv, av, ov

def projected_vector2plane(v, plane):
    """
    Project a vector into a plane.

    Parameters
    ----------
    v : list or np.array
        Vector to be projected into the plane.
    plane : list or np.array
        Vector representing the plane into which the vector should be
        projected.

    Returns
    -------
    v_in_p : np.array
        Initial vector projected into the provided plane.

    """
    if type(v) is list:
        v = np.array(v, dtype=np.float32)
    if type(plane) is list:
        plane = np.array(plane, dtype=np.float32)

    v_in_p = v - ((np.dot(v, plane) / np.sqrt(sum(plane**2))**2) * plane)
    return v_in_p

# %% Main
if __name__ == "__main__":
    print('Testing functions')
    print(('*** For a cup implanted at RI/RV 30/32 the right side cup normal'
           + ' vector with a neutral pelvis should be:'))
    print('    [0.424, 0.530, -0.734]')
    cupvec_r = vec_from_rirv(30, 32)[0]
    print('    cupvec_r = ' + str(list(cupvec_r.round(3))))
    print(('*** For an Euler sequence of YZX, and the included rotations, the'
           + ' dynamic cup normal vector should be:'))
    r1, r2, r3 = -2.689992, 18.435335, 0.56569
    ax1, ax2, ax3 = [0, 1, 0], [0, 0, 1], [1, 0, 0] # Euler sequence YZX
    print('    [0.266, 0.644, -0.717]')
    dyn_cupvec = rot_vec_by_three_angles(
        cupvec_r, r1, ax1, r2, ax2, r3, ax3)[0]
    print('    dyn_cupvec = ' + str(list(dyn_cupvec.round(3))))
    print('*** For this dynamic cup vector the RI/RV should be:')
    print('    20.368°, 40.066°')
    orients = ri_rv_av__ov_from_vec(dyn_cupvec)
    print('    dyn_ri_rv = ' + str(list(np.array(orients[0:2]).round(3))))
    print(('*** Radiographic inclination is the angle between the cup normal,'
           + ' projected into the frontal plane, and the SI axis.'))
    dyn_cupvec_proj = projected_vector2plane(dyn_cupvec, [0, 1, 0])
    ri_from_projection = angle_between_vectors(dyn_cupvec_proj, [0, 0, -1])
    print('    Recalculating RI from the projected vector and SI axis:')
    print('    ri_from_projection = ' + str(ri_from_projection.round(3)) + '°')