#!/usr/bin/env python3
"""
Convert poses to quaternions (qx, qy, qz, qw) from either:
- Euler angles (rx, ry, rz)
- Rotation vector (Rodrigues) (rx, ry, rz) where |r| is angle in rad

Supports:
- CSV input/output with configurable column names
- Unit options: radians/degrees for angles, mm->m for position
- Euler order: xyz (default) or zyx

Examples:
  python src/scripts/euler_to_quaternion.py \
        --in poses.csv --out poses_quat.csv --mode euler \
    --rx rx --ry ry --rz rz --order xyz \
    --x x --y y --z z --mm-to-m

  # Single value
    python src/scripts/euler_to_quaternion.py --mode euler --single 0.1 0.2 -0.3 --order zyx

    # Rotation vector (e.g., from many robots):
    python src/scripts/euler_to_quaternion.py --mode rotvec --single -1.433 0.114 -0.430
"""
from __future__ import annotations

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


def euler_to_quaternion(rx: float, ry: float, rz: float, order: str = "xyz", degrees: bool = False) -> Tuple[float, float, float, float]:
    """Convert Euler angles to quaternion [x, y, z, w].

    Args:
        rx, ry, rz: Euler angles. If degrees=True, values are in degrees; else radians.
        order: Rotation order. Supported: 'xyz' (default), 'zyx'. Intrinsic rotations.
        degrees: Whether the input angles are in degrees.

    Returns:
        (qx, qy, qz, qw)
    """
    if degrees:
        rx, ry, rz = math.radians(rx), math.radians(ry), math.radians(rz)

    def q_from_axis_angle(ax: str, angle: float) -> Tuple[float, float, float, float]:
        s = math.sin(angle / 2.0)
        c = math.cos(angle / 2.0)
        if ax == 'x':
            return (s, 0.0, 0.0, c)
        if ax == 'y':
            return (0.0, s, 0.0, c)
        if ax == 'z':
            return (0.0, 0.0, s, c)
        raise ValueError('Invalid axis')

    def q_mul(q1: Tuple[float, float, float, float], q2: Tuple[float, float, float, float]) -> Tuple[float, float, float, float]:
        x1, y1, z1, w1 = q1
        x2, y2, z2, w2 = q2
        x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
        y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
        z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
        w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
        return (x, y, z, w)

    if order == 'xyz':
        qx = q_from_axis_angle('x', rx)
        qy = q_from_axis_angle('y', ry)
        qz = q_from_axis_angle('z', rz)
        q = q_mul(q_mul(qx, qy), qz)
    elif order == 'zyx':
        qz = q_from_axis_angle('z', rz)
        qy = q_from_axis_angle('y', ry)
        qx = q_from_axis_angle('x', rx)
        q = q_mul(q_mul(qz, qy), qx)
    else:
        raise ValueError("Unsupported Euler order. Use 'xyz' or 'zyx'.")

    # Normalize
    norm = math.sqrt(sum(v * v for v in q))
    return tuple(v / (norm + 1e-12) for v in q)  # type: ignore


def rotvec_to_quaternion(rx: float, ry: float, rz: float) -> Tuple[float, float, float, float]:
    """Convert rotation vector (Rodrigues) to quaternion [x,y,z,w].
    The vector direction is the rotation axis and its norm is the rotation angle in radians.
    """
    angle = math.sqrt(rx * rx + ry * ry + rz * rz)
    if angle < 1e-12:
        return (0.0, 0.0, 0.0, 1.0)
    ax = rx / angle
    ay = ry / angle
    az = rz / angle
    s = math.sin(angle / 2.0)
    c = math.cos(angle / 2.0)
    q = (ax * s, ay * s, az * s, c)
    norm = math.sqrt(sum(v * v for v in q))
    return tuple(v / (norm + 1e-12) for v in q)  # type: ignore


def convert_csv(in_path: str, out_path: str,
                rx_col: str = 'rx', ry_col: str = 'ry', rz_col: str = 'rz',
                x_col: str | None = None, y_col: str | None = None, z_col: str | None = None,
                order: str = 'xyz', degrees: bool = False, mm_to_m: bool = False,
                mode: str = 'euler') -> None:
    """Read CSV, append quaternion columns (qx,qy,qz,qw) and write to output.

    If x/y/z columns are provided, optionally convert mm->m.
    """
    with open(in_path, 'r', newline='') as f_in:
        reader = csv.DictReader(f_in)
        fieldnames: List[str] = list(reader.fieldnames or [])
        # Ensure quaternion columns at the end
        for c in ['qx', 'qy', 'qz', 'qw']:
            if c not in fieldnames:
                fieldnames.append(c)
        # If converting position units, overwrite or add x_m/y_m/z_m
        if x_col and y_col and z_col and mm_to_m:
            for c in ['x_m', 'y_m', 'z_m']:
                if c not in fieldnames:
                    fieldnames.append(c)

        with open(out_path, 'w', newline='') as f_out:
            writer = csv.DictWriter(f_out, fieldnames=fieldnames)
            writer.writeheader()
            for row in reader:
                try:
                    rx = float(row[rx_col])
                    ry = float(row[ry_col])
                    rz = float(row[rz_col])
                except KeyError as e:
                    raise KeyError(f"Missing rotation columns in CSV: {e}")
                if mode == 'euler':
                    qx, qy, qz, qw = euler_to_quaternion(rx, ry, rz, order=order, degrees=degrees)
                elif mode == 'rotvec':
                    qx, qy, qz, qw = rotvec_to_quaternion(rx, ry, rz)
                else:
                    raise ValueError("mode must be 'euler' or 'rotvec'")
                row.update({'qx': f"{qx:.8f}", 'qy': f"{qy:.8f}", 'qz': f"{qz:.8f}", 'qw': f"{qw:.8f}"})

                if x_col and y_col and z_col and (x_col in row and y_col in row and z_col in row):
                    try:
                        x = float(row[x_col])
                        y = float(row[y_col])
                        z = float(row[z_col])
                        if mm_to_m:
                            row['x_m'] = f"{x / 1000.0:.6f}"
                            row['y_m'] = f"{y / 1000.0:.6f}"
                            row['z_m'] = f"{z / 1000.0:.6f}"
                    except ValueError:
                        pass

                writer.writerow(row)


def main():
    parser = argparse.ArgumentParser(description='Convert pose angles to quaternion (qx,qy,qz,qw) from Euler or rotation vector.')
    parser.add_argument('--mode', type=str, default='euler', choices=['euler', 'rotvec'], help='Input angle format')
    parser.add_argument('--in', dest='in_path', type=str, help='Input CSV file path')
    parser.add_argument('--out', dest='out_path', type=str, help='Output CSV file path')
    parser.add_argument('--rx', type=str, default='rx', help='Column name for rx')
    parser.add_argument('--ry', type=str, default='ry', help='Column name for ry')
    parser.add_argument('--rz', type=str, default='rz', help='Column name for rz')
    parser.add_argument('--x', type=str, default=None, help='Column name for x (position)')
    parser.add_argument('--y', type=str, default=None, help='Column name for y (position)')
    parser.add_argument('--z', type=str, default=None, help='Column name for z (position)')
    parser.add_argument('--order', type=str, default='xyz', choices=['xyz', 'zyx'], help='Euler order')
    parser.add_argument('--degrees', action='store_true', help='Input angles are in degrees (default radians)')
    parser.add_argument('--mm-to-m', action='store_true', help='Convert position units from mm to m when x/y/z provided')

    # single conversion mode
    parser.add_argument('--single', nargs=3, type=float, metavar=('RX', 'RY', 'RZ'),
                        help='Convert a single Euler triplet to quaternion (prints to stdout)')

    args = parser.parse_args()

    if args.single is not None:
        rx, ry, rz = args.single
        if args.mode == 'euler':
            q = euler_to_quaternion(rx, ry, rz, order=args.order, degrees=args.degrees)
        else:
            q = rotvec_to_quaternion(rx, ry, rz)
        print(f"qx,qy,qz,qw = {q[0]:.8f}, {q[1]:.8f}, {q[2]:.8f}, {q[3]:.8f}")
        return

    if not args.in_path or not args.out_path:
        parser.error('CSV mode requires --in and --out')

    convert_csv(
        in_path=args.in_path,
        out_path=args.out_path,
        rx_col=args.rx,
        ry_col=args.ry,
        rz_col=args.rz,
        x_col=args.x,
        y_col=args.y,
        z_col=args.z,
        order=args.order,
        degrees=args.degrees,
        mm_to_m=args.mm_to_m,
        mode=args.mode,
    )
    print(f"Written: {args.out_path}")


if __name__ == '__main__':
    main()
