"""
Rotation conversion utilities.

Supports three rotation representations:
- axis_angle: [3] - native SMPL-X format, compact but discontinuous at pi
- rotation_matrix: [3, 3] - native WiLoR format, interpretable but redundant
- rotation_6d: [6] - first two columns of rotation matrix, continuous (recommended for NNs)

Based on:
- Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019
"""

from __future__ import annotations

from typing import Literal

import numpy as np

RotationFormat = Literal["axis_angle", "rotation_matrix", "rotation_6d"]


def axis_angle_to_matrix(axis_angle: np.ndarray) -> np.ndarray:
    """
    Convert axis-angle rotation to rotation matrix using Rodrigues formula.

    Args:
        axis_angle: [..., 3] axis-angle vectors

    Returns:
        [..., 3, 3] rotation matrices
    """
    batch_shape = axis_angle.shape[:-1]
    aa = axis_angle.reshape(-1, 3)

    # Compute angle (magnitude of axis-angle vector)
    angle = np.linalg.norm(aa, axis=-1, keepdims=True)

    # Handle near-zero angles: use [0, 0, 1] as default axis for identity rotations
    # This avoids division by zero and returns identity matrix for zero rotations
    small_angle_mask = angle.flatten() < 1e-8

    # Safe division: use angle or 1.0 where angle is small
    safe_angle = np.where(angle < 1e-8, 1.0, angle)
    axis = aa / safe_angle

    # For small angles, set a default axis (doesn't matter since angle ~= 0)
    axis[small_angle_mask] = [0.0, 0.0, 1.0]

    # Rodrigues formula components
    cos_angle = np.cos(angle)
    sin_angle = np.sin(angle)

    # Skew-symmetric matrix K
    K = np.zeros((aa.shape[0], 3, 3), dtype=aa.dtype)
    K[:, 0, 1] = -axis[:, 2]
    K[:, 0, 2] = axis[:, 1]
    K[:, 1, 0] = axis[:, 2]
    K[:, 1, 2] = -axis[:, 0]
    K[:, 2, 0] = -axis[:, 1]
    K[:, 2, 1] = axis[:, 0]

    # Rotation matrix: R = I + sin(θ)K + (1-cos(θ))K²
    I = np.eye(3, dtype=aa.dtype)
    R = I + sin_angle[..., None] * K + (1 - cos_angle[..., None]) * np.matmul(K, K)

    return R.reshape(*batch_shape, 3, 3)


def matrix_to_axis_angle(matrix: np.ndarray) -> np.ndarray:
    """
    Convert rotation matrix to axis-angle.

    Args:
        matrix: [..., 3, 3] rotation matrices

    Returns:
        [..., 3] axis-angle vectors
    """
    batch_shape = matrix.shape[:-2]
    mat = matrix.reshape(-1, 3, 3)

    # Compute angle from trace: trace(R) = 1 + 2*cos(θ)
    trace = np.trace(mat, axis1=-2, axis2=-1)
    cos_angle = (trace - 1) / 2
    cos_angle = np.clip(cos_angle, -1, 1)
    angle = np.arccos(cos_angle)

    # Compute axis from skew-symmetric part
    # R - R^T = 2*sin(θ)*K where K is skew-symmetric with axis components
    skew = mat - np.transpose(mat, (0, 2, 1))
    axis = np.stack([skew[:, 2, 1], skew[:, 0, 2], skew[:, 1, 0]], axis=-1)

    # Normalize axis with proper handling of near-zero angles (identity matrices)
    axis_norm = np.linalg.norm(axis, axis=-1, keepdims=True)
    small_norm_mask = axis_norm.flatten() < 1e-8

    # Safe division: use axis_norm or 1.0 where norm is small
    safe_norm = np.where(axis_norm < 1e-8, 1.0, axis_norm)
    axis = axis / safe_norm

    # For identity matrices (small angle), set default axis [0, 0, 1]
    # The resulting axis-angle will be [0, 0, 0] since angle ~= 0
    axis[small_norm_mask] = [0.0, 0.0, 1.0]

    # Scale by angle
    aa = axis * angle[..., None]

    return aa.reshape(*batch_shape, 3)


def matrix_to_6d(matrix: np.ndarray) -> np.ndarray:
    """
    Convert rotation matrix to 6D representation.

    The 6D representation is the first two columns of the rotation matrix,
    which provides a continuous representation suitable for neural networks.

    Args:
        matrix: [..., 3, 3] rotation matrices

    Returns:
        [..., 6] 6D rotation vectors
    """
    # Extract first two columns
    col1 = matrix[..., :, 0]  # [..., 3]
    col2 = matrix[..., :, 1]  # [..., 3]

    return np.concatenate([col1, col2], axis=-1)


def rotation_6d_to_matrix(r6d: np.ndarray) -> np.ndarray:
    """
    Convert 6D rotation to rotation matrix using QR decomposition.

    Uses QR decomposition for numerical stability instead of single-pass
    Gram-Schmidt orthogonalization.

    Args:
        r6d: [..., 6] 6D rotation vectors

    Returns:
        [..., 3, 3] rotation matrices
    """
    from scipy.linalg import qr

    batch_shape = r6d.shape[:-1]
    r6d_flat = r6d.reshape(-1, 6)
    n_batch = r6d_flat.shape[0]

    # Split into two 3D vectors
    a1 = r6d_flat[:, :3]
    a2 = r6d_flat[:, 3:6]

    # Build matrices for QR decomposition (3x2 matrix per sample)
    # We want to orthogonalize [a1, a2] to get [b1, b2]
    result = np.zeros((n_batch, 3, 3), dtype=r6d.dtype)

    for i in range(n_batch):
        # Create 3x2 matrix with a1, a2 as columns
        A = np.column_stack([a1[i], a2[i]])

        # QR decomposition gives orthonormal columns
        Q, R = qr(A, mode="economic")

        # Ensure proper orientation (positive diagonal in R)
        # This maintains the direction of a1 and orientation of a2
        signs = np.sign(np.diag(R))
        signs = np.where(signs == 0, 1, signs)  # Handle zero diagonal
        Q = Q * signs

        # b1 and b2 are the orthonormalized columns
        b1 = Q[:, 0]
        b2 = Q[:, 1]

        # b3 = b1 × b2 (cross product for right-handed coordinate system)
        b3 = np.cross(b1, b2)

        # Stack into rotation matrix (columns are b1, b2, b3)
        result[i] = np.column_stack([b1, b2, b3])

    return result.reshape(*batch_shape, 3, 3)


def axis_angle_to_6d(axis_angle: np.ndarray) -> np.ndarray:
    """Convert axis-angle to 6D representation."""
    return matrix_to_6d(axis_angle_to_matrix(axis_angle))


def rotation_6d_to_axis_angle(r6d: np.ndarray) -> np.ndarray:
    """Convert 6D representation to axis-angle."""
    return matrix_to_axis_angle(rotation_6d_to_matrix(r6d))


class RotationConverter:
    """
    Convert between rotation formats.

    Example:
        >>> converter = RotationConverter(target_format="rotation_6d")
        >>> r6d = converter.from_axis_angle(axis_angle_data)
        >>> r6d = converter.from_matrix(rotation_matrices)
    """

    def __init__(self, target_format: RotationFormat = "rotation_6d"):
        """
        Initialize converter.

        Args:
            target_format: Output format ('axis_angle', 'rotation_matrix', 'rotation_6d')
        """
        self.target_format = target_format

    def from_axis_angle(self, data: np.ndarray) -> np.ndarray:
        """Convert from axis-angle to target format."""
        if self.target_format == "axis_angle":
            return data
        elif self.target_format == "rotation_matrix":
            return axis_angle_to_matrix(data)
        elif self.target_format == "rotation_6d":
            return axis_angle_to_6d(data)
        else:
            raise ValueError(f"Unknown target format: {self.target_format}")

    def from_matrix(self, data: np.ndarray) -> np.ndarray:
        """Convert from rotation matrix to target format."""
        if self.target_format == "rotation_matrix":
            return data
        elif self.target_format == "axis_angle":
            return matrix_to_axis_angle(data)
        elif self.target_format == "rotation_6d":
            return matrix_to_6d(data)
        else:
            raise ValueError(f"Unknown target format: {self.target_format}")

    def from_6d(self, data: np.ndarray) -> np.ndarray:
        """Convert from 6D representation to target format."""
        if self.target_format == "rotation_6d":
            return data
        elif self.target_format == "rotation_matrix":
            return rotation_6d_to_matrix(data)
        elif self.target_format == "axis_angle":
            return rotation_6d_to_axis_angle(data)
        else:
            raise ValueError(f"Unknown target format: {self.target_format}")

    def convert(self, data: np.ndarray, source_format: RotationFormat) -> np.ndarray:
        """
        Convert from any source format to target format.

        Args:
            data: Input rotation data
            source_format: Format of input data

        Returns:
            Data in target format
        """
        if source_format == "axis_angle":
            return self.from_axis_angle(data)
        elif source_format == "rotation_matrix":
            return self.from_matrix(data)
        elif source_format == "rotation_6d":
            return self.from_6d(data)
        else:
            raise ValueError(f"Unknown source format: {source_format}")
