from .Rotation import Rotation as R

# Intrinsic ZYX euler angles are yaw pitch roll
# https://en.wikipedia.org/wiki/Euler_angles#Conventions
# Intrinsic euler angles are defined using capital letters in scipy
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.transform.Rotation.as_euler.html#


class IntrinsicZYXEuler:
    def __init__(self, rot):
        euler = rot.as_euler("ZYX", degrees=True)
        self.yaw, self.pitch, self.roll = euler[0] if not rot.single else euler

    def __str__(self):
        return f"Yaw = {self.yaw}, Pitch = {self.pitch}, Roll = {self.roll}"

    def __repr__(self):
        return f"IntrinsicZYXEuler(Yaw = {self.yaw}, Pitch = {self.pitch}, Roll = {self.roll})"


def make_yaw_offset(offset_angle):
    return R.from_euler(
        seq="ZYX",
        angles=[-offset_angle, 0, 0],
        degrees=True,
    )


def normalize_angle(angle):
    return (angle + 180) % 360 - 180

class TrunkAngle:
    def __init__(self):
        """
        Initializes the TrunkAngle class.

        """
        self.trunk_Yawoffset_q = None

        # sensor to segment alignment quaternion, inv denotes conjugate.
        self.BS_q_s1_inv = None
        self.BS_q_l1_inv = None
        self.flex = 0
        self.sway = 0
        self.rot = 0



    def calibrate(self, s1_quat, l1_quat):
        """
        Performs sensor to segment calibration.

        Args:
            l1_quat (Rotation): Rotation representing the l1 orientation.
            s1_quat (Rotation): Rotation representing the s1 orientation.

        """

        s1_yaw = IntrinsicZYXEuler(s1_quat).yaw
        l1_yaw = IntrinsicZYXEuler(l1_quat).yaw

        # This code below calculates how much the yaw differs between the
        # various segments and the pelvis. This is used to correct for yaw
        # differences between the segment and the pelvis.

        trunk_offset_angle = l1_yaw  - s1_yaw
        self.trunk_Yawoffset_q = make_yaw_offset(trunk_offset_angle)

        # Get body segment quaternions relative to the target.
        # Pelvis does not include the offset since it is the base.
        # We dont include the offset for the foot since it does not rotate.
        self.BS_q_s1_inv = self._initialize_quat_inv(s1_quat)
        self.BS_q_l1_inv = self._initialize_quat_inv(l1_quat)


    def _initialize_quat_inv(self, this_quat):
        """
        Initializes the conjugate quaternion for sensor to segment alignment.

        Args:
            this_quat (Rotation): Rotation representing the sensor orientation.

        Returns:
            Rotation: Conjugate Rotation representing the alignment target.
        """
        R_this_init_Yaw = IntrinsicZYXEuler(this_quat).yaw

        R_this_init_Yaw = (R_this_init_Yaw + 180) % 360 - 180

        GB_q0_target = R.from_euler(
            seq="ZYX", angles=[R_this_init_Yaw, 0, 0], degrees=True
        )

        this_quat = this_quat.inv()
        return this_quat * GB_q0_target

    def _calculate_GB_quat(self, GS_quat, bs_inv_quat, Yawoffset_quat=None):
        # This method calculates the quaternion relative to the body segment.
        GB_quat = GS_quat * bs_inv_quat
        if Yawoffset_quat != None:
            GB_quat = Yawoffset_quat * GB_quat
        return GB_quat

    def _base_angle_calc(self, s1_quat, l1_quat):
        GB_trunk_q = self._calculate_GB_quat(s1_quat, self.BS_q_s1_inv)

        GB_thigh_q = self._calculate_GB_quat(
            l1_quat, self.BS_q_l1_inv, self.trunk_Yawoffset_q
        )

        # This code calculates the hip angles by getting the relative rotation between
        # the pelvis and the thigh. Other segments are calculated in a similar way.
        B_q_trunk_angles = GB_trunk_q.inv() * GB_thigh_q

        # We use the zyx euler angle order since lower case is the extrinsic or global
        # rotation order. These body joint angles are global rotations. We grab the last
        # element of the euler angle array since we want the rotation around the x axis.
        # This is the same for the other angles too.
        Trunk_angles = B_q_trunk_angles.as_euler("zyx", degrees=True)
        return Trunk_angles
    
    def update(self, s1_quat, l1_quat):
        trunk_angles = self._base_angle_calc(s1_quat, l1_quat)
        self.flex = trunk_angles[2]
        self.sway = trunk_angles[1]
        self.rot = normalize_angle(trunk_angles[0] + 180)
