#!/usr/bin/env python3
"""
rotate_pose.py

Small utility to rotate a 7-element pose (x,y,z,qx,qy,qz,qw or qw,qx,qy,qz)
around a given axis by a specified angle (degrees).

Usage examples:
  ./rotate_pose.py --pose 0.1 0.2 0.3 -0.6 0.1 -0.1 0.7 --axis z --angle 180 --order qxqyqzqw --frame world
  python3 rotate_pose.py --pose 0.179387 0.341708 -0.099286 -0.88140507 0.12851699 -0.17498077 0.41951188 --axis z --angle 180 --order qxqyqzqw --frame world

Supported frames:
  - world: apply rotation in world frame (q' = q_rot * q)
  - body: apply rotation in body/local frame (q' = q * q_rot)

Supported quaternion order strings: "qxqyqzqw" (default) or "wxyz" (qw,qx,qy,qz)

"""
import argparse
import math
from typing import List, Tuple


def normalize_quat(q: Tuple[float, float, float, float]) -> Tuple[float, float, float, float]:
    x, y, z, w = q
    n = math.sqrt(x * x + y * y + z * z + w * w)
    if n == 0:
        return (0.0, 0.0, 0.0, 1.0)
    return (x / n, y / n, z / n, w / n)


def axis_angle_to_quat(axis: str, angle_deg: float) -> Tuple[float, float, float, float]:
    a = math.radians(angle_deg)
    ca = math.cos(a / 2.0)
    sa = math.sin(a / 2.0)
    ax = axis.lower()
    if ax == 'x':
        return (sa, 0.0, 0.0, ca)
    if ax == 'y':
        return (0.0, sa, 0.0, ca)
    if ax == 'z':
        return (0.0, 0.0, sa, ca)
    raise ValueError('axis must be x,y or z')


def quat_mul(a: Tuple[float,float,float,float], b: Tuple[float,float,float,float]) -> Tuple[float,float,float,float]:
    ax, ay, az, aw = a
    bx, by, bz, bw = b
    x = aw*bx + ax*bw + ay*bz - az*by
    y = aw*by - ax*bz + ay*bw + az*bx
    z = aw*bz + ax*by - ay*bx + az*bw
    w = aw*bw - ax*bx - ay*by - az*bz
    return (x, y, z, w)


def rotate_position(pos: Tuple[float,float,float], axis: str, angle_deg: float) -> Tuple[float,float,float]:
    # For 180-degree rotations we can do simple sign flips but implement general rotation using quaternion
    qrot = axis_angle_to_quat(axis, angle_deg)
    # convert pos to quaternion (v,0)
    px, py, pz = pos
    p = (px, py, pz, 0.0)
    qrot_n = normalize_quat(qrot)
    qrot_conj = (-qrot_n[0], -qrot_n[1], -qrot_n[2], qrot_n[3])
    tmp = quat_mul(qrot_n, p)
    p_rot = quat_mul(tmp, qrot_conj)
    return (p_rot[0], p_rot[1], p_rot[2])


def parse_pose(vals: List[float], order: str) -> Tuple[Tuple[float,float,float], Tuple[float,float,float,float]]:
    if order == 'qxqyqzqw':
        if len(vals) != 7:
            raise ValueError('need 7 values for pose')
        return (vals[0], vals[1], vals[2]), (vals[3], vals[4], vals[5], vals[6])
    if order == 'wxyz':
        if len(vals) != 7:
            raise ValueError('need 7 values for pose')
        return (vals[0], vals[1], vals[2]), (vals[3], vals[4], vals[5], vals[6]) if False else (vals[4], vals[5], vals[6], vals[3])
    raise ValueError('unsupported order')


def format_pose(pos: Tuple[float,float,float], quat: Tuple[float,float,float,float], order: str) -> List[float]:
    if order == 'qxqyqzqw':
        return [pos[0], pos[1], pos[2], quat[0], quat[1], quat[2], quat[3]]
    if order == 'wxyz':
        # qw,qx,qy,qz
        return [pos[0], pos[1], pos[2], quat[3], quat[0], quat[1], quat[2]]
    raise ValueError('unsupported order')


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--pose', nargs=7, type=float, required=True, help='pose values')
    p.add_argument('--axis', choices=['x','y','z'], default='z')
    p.add_argument('--angle', type=float, default=180.0)
    p.add_argument('--order', choices=['qxqyqzqw','wxyz'], default='qxqyqzqw')
    p.add_argument('--frame', choices=['world','body'], default='world')
    args = p.parse_args()

    pos, quat = parse_pose(list(args.pose), args.order)
    # rotation quaternion
    qrot = axis_angle_to_quat(args.axis, args.angle)
    qrot = normalize_quat(qrot)

    # apply rotation to position in world frame: pos' = R * pos
    pos_rot = rotate_position(pos, args.axis, args.angle)

    # apply rotation to orientation
    if args.frame == 'world':
        qnew = quat_mul(qrot, quat)
    else:
        qnew = quat_mul(quat, qrot)
    qnew = normalize_quat(qnew)

    out = format_pose(pos_rot, qnew, args.order)
    # Print as numeric list without quotes
    print('rotated_pose:', [float('{:.8f}'.format(x)) for x in out])

if __name__ == '__main__':
    main()
