# Many methods borrow heavily or entirely from transforms3d https://github.com/matthew-brett/transforms3d
# eventually some of these may be upstreamed, but credit to transforms3d
# authors for implementing the many of the formulations we use here.

import itertools

import numpy as np

"""
Rotations
=========

Note: these have caused many subtle bugs in the past.
Be careful while updating these methods and while using them in clever ways.

See MuJoCo documentation here: http://mujoco.org/book/modeling.html#COrientation

Conventions
-----------
    - All functions accept batches as well as individual rotations
    - All rotation conventions match respective MuJoCo defaults
    - All angles are in radians
    - Matricies follow LR convention
    - Euler Angles are all relative with 'xyz' axes ordering
    - See specific representation for more information

Representations
---------------

Euler
    There are many euler angle frames -- here we will strive to use the default
        in MuJoCo, which is eulerseq='xyz'.
    This frame is a relative rotating frame, about x, y, and z axes in order.
        Relative rotating means that after we rotate about x, then we use the
        new (rotated) y, and the same for z.

Quaternions
    These are defined in terms of rotation (angle) about a unit vector (x, y, z)
    We use the following <q0, q1, q2, q3> convention:
            q0 = cos(angle / 2)
            q1 = sin(angle / 2) * x
            q2 = sin(angle / 2) * y
            q3 = sin(angle / 2) * z
        This is also sometimes called qw, qx, qy, qz.
    Note that quaternions are ambiguous, because we can represent a rotation by
        angle about vector <x, y, z> and -angle about vector <-x, -y, -z>.
        To choose between these, we pick "first nonzero positive", where we
        make the first nonzero element of the quaternion positive.
    This can result in mismatches if you're converting an quaternion that is not
        "first nonzero positive" to a different representation and back.

Axis Angle
    (Not currently implemented)
    These are very straightforward.  Rotation is angle about a unit vector.

XY Axes
    (Not currently implemented)
    We are given x axis and y axis, and z axis is cross product of x and y.

Z Axis
    This is NOT RECOMMENDED.  Defines a unit vector for the Z axis,
        but rotation about this axis is not well defined.
    Instead pick a fixed reference direction for another axis (e.g. X)
        and calculate the other (e.g. Y = Z cross-product X),
        then use XY Axes rotation instead.

SO3
    (Not currently implemented)
    While not supported by MuJoCo, this representation has a lot of nice features.
    We expect to add support for these in the future.

TODO / Missing
--------------
    - Rotation integration or derivatives (e.g. velocity conversions)
    - More representations (SO3, etc)
    - Random sampling (e.g. sample uniform random rotation)
    - Performance benchmarks/measurements
    - (Maybe) define everything as to/from matricies, for simplicity
"""

# For testing whether a number is close to zero
_FLOAT_EPS = np.finfo(np.float64).eps
_EPS4 = _FLOAT_EPS * 4.0


def euler2mat(euler):
    """ Convert Euler Angles to Rotation Matrix.  See rotation.py for notes """
    euler = np.asarray(euler, dtype=np.float64)
    assert euler.shape[-1] == 3, "Invalid shaped euler {}".format(euler)

    ai, aj, ak = -euler[..., 2], -euler[..., 1], -euler[..., 0]
    si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak)
    ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak)
    cc, cs = ci * ck, ci * sk
    sc, ss = si * ck, si * sk

    mat = np.empty(euler.shape[:-1] + (3, 3), dtype=np.float64)
    mat[..., 2, 2] = cj * ck
    mat[..., 2, 1] = sj * sc - cs
    mat[..., 2, 0] = sj * cc + ss
    mat[..., 1, 2] = cj * sk
    mat[..., 1, 1] = sj * ss + cc
    mat[..., 1, 0] = sj * cs - sc
    mat[..., 0, 2] = -sj
    mat[..., 0, 1] = cj * si
    mat[..., 0, 0] = cj * ci
    return mat


def euler2quat(euler):
    """ Convert Euler Angles to Quaternions.  See rotation.py for notes """
    euler = np.asarray(euler, dtype=np.float64)
    assert euler.shape[-1] == 3, "Invalid shape euler {}".format(euler)

    ai, aj, ak = euler[..., 2] / 2, -euler[..., 1] / 2, euler[..., 0] / 2
    si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak)
    ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak)
    cc, cs = ci * ck, ci * sk
    sc, ss = si * ck, si * sk

    quat = np.empty(euler.shape[:-1] + (4,), dtype=np.float64)
    quat[..., 0] = cj * cc + sj * ss
    quat[..., 3] = cj * sc - sj * cs
    quat[..., 2] = -(cj * ss + sj * cc)
    quat[..., 1] = cj * cs - sj * sc
    return quat


def mat2euler(mat):
    """ Convert Rotation Matrix to Euler Angles.  See rotation.py for notes """
    mat = np.asarray(mat, dtype=np.float64)
    assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)

    cy = np.sqrt(mat[..., 2, 2] * mat[..., 2, 2] + mat[..., 1, 2] * mat[..., 1, 2])
    condition = cy > _EPS4
    euler = np.empty(mat.shape[:-1], dtype=np.float64)
    euler[..., 2] = np.where(
        condition,
        -np.arctan2(mat[..., 0, 1], mat[..., 0, 0]),
        -np.arctan2(-mat[..., 1, 0], mat[..., 1, 1]),
    )
    euler[..., 1] = np.where(
        condition, -np.arctan2(-mat[..., 0, 2], cy), -np.arctan2(-mat[..., 0, 2], cy)
    )
    euler[..., 0] = np.where(
        condition, -np.arctan2(mat[..., 1, 2], mat[..., 2, 2]), 0.0
    )
    return euler


def mat2quat(mat):
    """ Convert Rotation Matrix to Quaternion.  See rotation.py for notes """
    mat = np.asarray(mat, dtype=np.float64)
    assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)

    Qxx, Qyx, Qzx = mat[..., 0, 0], mat[..., 0, 1], mat[..., 0, 2]
    Qxy, Qyy, Qzy = mat[..., 1, 0], mat[..., 1, 1], mat[..., 1, 2]
    Qxz, Qyz, Qzz = mat[..., 2, 0], mat[..., 2, 1], mat[..., 2, 2]
    # Fill only lower half of symmetric matrix
    K = np.zeros(mat.shape[:-2] + (4, 4), dtype=np.float64)
    K[..., 0, 0] = Qxx - Qyy - Qzz
    K[..., 1, 0] = Qyx + Qxy
    K[..., 1, 1] = Qyy - Qxx - Qzz
    K[..., 2, 0] = Qzx + Qxz
    K[..., 2, 1] = Qzy + Qyz
    K[..., 2, 2] = Qzz - Qxx - Qyy
    K[..., 3, 0] = Qyz - Qzy
    K[..., 3, 1] = Qzx - Qxz
    K[..., 3, 2] = Qxy - Qyx
    K[..., 3, 3] = Qxx + Qyy + Qzz
    K /= 3.0
    # TODO: vectorize this -- probably could be made faster
    q = np.empty(K.shape[:-2] + (4,))
    it = np.nditer(q[..., 0], flags=["multi_index"])
    while not it.finished:
        # Use Hermitian eigenvectors, values for speed
        vals, vecs = np.linalg.eigh(K[it.multi_index])
        # Select largest eigenvector, reorder to w,x,y,z quaternion
        q[it.multi_index] = vecs[[3, 0, 1, 2], np.argmax(vals)]
        # Prefer quaternion with positive w
        # (q * -1 corresponds to same rotation as q)
        if q[it.multi_index][0] < 0:
            q[it.multi_index] *= -1
        it.iternext()
    return q


def quat2euler(quat):
    """ Convert Quaternion to Euler Angles.  See rotation.py for notes """
    return mat2euler(quat2mat(quat))


def subtract_euler(e1, e2):
    assert e1.shape == e2.shape
    assert e1.shape[-1] == 3
    q1 = euler2quat(e1)
    q2 = euler2quat(e2)
    q_diff = quat_mul(q1, quat_conjugate(q2))
    return quat2euler(q_diff)


def quat2mat(quat):
    """ Convert Quaternion to Euler Angles.  See rotation.py for notes """
    quat = np.asarray(quat, dtype=np.float64)
    assert quat.shape[-1] == 4, "Invalid shape quat {}".format(quat)

    w, x, y, z = quat[..., 0], quat[..., 1], quat[..., 2], quat[..., 3]
    Nq = np.sum(quat * quat, axis=-1)
    s = 2.0 / Nq
    X, Y, Z = x * s, y * s, z * s
    wX, wY, wZ = w * X, w * Y, w * Z
    xX, xY, xZ = x * X, x * Y, x * Z
    yY, yZ, zZ = y * Y, y * Z, z * Z

    mat = np.empty(quat.shape[:-1] + (3, 3), dtype=np.float64)
    mat[..., 0, 0] = 1.0 - (yY + zZ)
    mat[..., 0, 1] = xY - wZ
    mat[..., 0, 2] = xZ + wY
    mat[..., 1, 0] = xY + wZ
    mat[..., 1, 1] = 1.0 - (xX + zZ)
    mat[..., 1, 2] = yZ - wX
    mat[..., 2, 0] = xZ - wY
    mat[..., 2, 1] = yZ + wX
    mat[..., 2, 2] = 1.0 - (xX + yY)
    return np.where((Nq > _FLOAT_EPS)[..., np.newaxis, np.newaxis], mat, np.eye(3))


def quat_conjugate(q):
    inv_q = -q
    inv_q[..., 0] *= -1
    return inv_q


def quat_mul(q0, q1):
    assert q0.shape == q1.shape
    assert q0.shape[-1] == 4
    assert q1.shape[-1] == 4

    w0 = q0[..., 0]
    x0 = q0[..., 1]
    y0 = q0[..., 2]
    z0 = q0[..., 3]

    w1 = q1[..., 0]
    x1 = q1[..., 1]
    y1 = q1[..., 2]
    z1 = q1[..., 3]

    w = w0 * w1 - x0 * x1 - y0 * y1 - z0 * z1
    x = w0 * x1 + x0 * w1 + y0 * z1 - z0 * y1
    y = w0 * y1 + y0 * w1 + z0 * x1 - x0 * z1
    z = w0 * z1 + z0 * w1 + x0 * y1 - y0 * x1
    q = np.array([w, x, y, z])
    if q.ndim == 2:
        q = q.swapaxes(0, 1)
    assert q.shape == q0.shape
    return q


def quat_rot_vec(q, v0):
    q_v0 = np.array([0, v0[0], v0[1], v0[2]])
    q_v = quat_mul(q, quat_mul(q_v0, quat_conjugate(q)))
    v = q_v[1:]
    return v


def quat_identity():
    return np.array([1, 0, 0, 0])


def quat_difference(q, p):
    return quat_normalize(quat_mul(q, quat_conjugate(p)))


def quat_magnitude(q):
    w = q[..., 0]
    assert np.all(w >= 0)
    return 2 * np.arccos(np.clip(w, -1.0, 1.0))


def quat_normalize(q):
    assert q.shape[-1] == 4
    sign = np.sign(q[..., [0]])
    # Sign takes value of 0 whenever the input is 0, but we actually don't want to do that
    sign[sign == 0] = 1
    return q * sign  # use quat with w >= 0


def quat_average(quats, weights=None):
    """Weighted average of a list of quaternions."""
    n_quats = len(quats)
    weights = np.array([1.0 / n_quats] * n_quats if weights is None else weights)
    assert np.all(weights >= 0.0)
    assert len(weights) == len(quats)
    weights = weights / np.sum(weights)

    # Average of quaternion:
    # https://math.stackexchange.com/questions/1204228
    # /how-would-i-apply-an-exponential-moving-average-to-quaternions
    outer_prods = [w * np.outer(q, q) for w, q in zip(weights, quats)]
    summed_outer_prod = np.sum(outer_prods, axis=0)
    assert summed_outer_prod.shape == (4, 4)

    evals, evecs = np.linalg.eig(summed_outer_prod)
    evals, evecs = np.real(evals), np.real(evecs)
    biggest_i = np.argmax(np.real(evals))
    return quat_normalize(evecs[:, biggest_i])


def quat2axisangle(quat):
    theta = 0
    axis = np.array([0, 0, 1])
    sin_theta = np.linalg.norm(quat[1:])

    if sin_theta > 0.0001:
        theta = 2 * np.arcsin(sin_theta)
        theta *= 1 if quat[0] >= 0 else -1
        axis = quat[1:] / sin_theta

    return axis, theta


def euler2point_euler(euler):
    _euler = euler.copy()
    if len(_euler.shape) < 2:
        _euler = np.expand_dims(_euler, 0)
    assert _euler.shape[1] == 3
    _euler_sin = np.sin(_euler)
    _euler_cos = np.cos(_euler)
    return np.concatenate([_euler_sin, _euler_cos], axis=-1)


def point_euler2euler(euler):
    _euler = euler.copy()
    if len(_euler.shape) < 2:
        _euler = np.expand_dims(_euler, 0)
    assert _euler.shape[1] == 6
    angle = np.arctan(_euler[..., :3] / _euler[..., 3:])
    angle[_euler[..., 3:] < 0] += np.pi
    return angle


def quat2point_quat(quat):
    # Should be in qw, qx, qy, qz
    _quat = quat.copy()
    if len(_quat.shape) < 2:
        _quat = np.expand_dims(_quat, 0)
    assert _quat.shape[1] == 4
    angle = np.arccos(_quat[:, [0]]) * 2
    xyz = _quat[:, 1:]
    xyz[np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5] = (xyz / np.sin(angle / 2))[
        np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5
    ]
    return np.concatenate([np.sin(angle), np.cos(angle), xyz], axis=-1)


def point_quat2quat(quat):
    _quat = quat.copy()
    if len(_quat.shape) < 2:
        _quat = np.expand_dims(_quat, 0)
    assert _quat.shape[1] == 5
    angle = np.arctan(_quat[:, [0]] / _quat[:, [1]])
    qw = np.cos(angle / 2)

    qxyz = _quat[:, 2:]
    qxyz[np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5] = (qxyz * np.sin(angle / 2))[
        np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5
    ]
    return np.concatenate([qw, qxyz], axis=-1)


def normalize_angles(angles, low=-np.pi, high=np.pi):
    """Puts angles in [low, high] range."""
    angles = angles.copy()
    if angles.size > 0:
        angles = np.mod(angles - low, high - low) + low
        assert low - 1e-6 <= angles.min() and angles.max() <= high + 1e-6
    return angles


def round_to_straight_angles(angles):
    """Returns closest angle modulo 90 degrees """
    angles = np.round(angles / (np.pi / 2)) * (np.pi / 2)
    return normalize_angles(angles)


def round_to_straight_quat(quat):
    angles = quat2euler(quat)
    rounded_angles = round_to_straight_angles(angles)
    return euler2quat(rounded_angles)


def get_parallel_rotations():
    mult90 = [0, np.pi / 2, -np.pi / 2, np.pi]
    parallel_rotations = []
    for euler in itertools.product(mult90, repeat=3):
        canonical = mat2euler(euler2mat(euler))
        canonical = np.round(canonical / (np.pi / 2))
        if canonical[0] == -2:
            canonical[0] = 2
        if canonical[2] == -2:
            canonical[2] = 2
        canonical *= np.pi / 2
        if all([(canonical != rot).any() for rot in parallel_rotations]):
            parallel_rotations += [canonical]
    assert len(parallel_rotations) == 24
    return parallel_rotations


def get_parallel_rotations_180():
    mult180 = [0, np.pi]
    parallel_rotations = []
    for euler in itertools.product(mult180, repeat=3):
        canonical = mat2euler(euler2mat(euler))
        canonical = np.round(canonical / (np.pi / 2))
        if canonical[0] == -2:
            canonical[0] = 2
        if canonical[2] == -2:
            canonical[2] = 2
        canonical *= np.pi / 2
        if all([(canonical != rot).any() for rot in parallel_rotations]):
            parallel_rotations += [canonical]
    assert len(parallel_rotations) == 4
    return parallel_rotations


def quat_from_angle_and_axis(angle, axis):
    assert axis.shape[-1] == 3
    axis /= np.linalg.norm(axis, axis=-1, keepdims=True)
    angle = np.reshape(angle, axis[..., :1].shape)
    w = np.cos(angle / 2.0)
    v = np.sin(angle / 2.0) * axis
    quat = np.concatenate([w, v], axis=-1)
    quat /= np.linalg.norm(quat, axis=-1, keepdims=True)

    assert np.array_equal(quat.shape[:-1], axis.shape[:-1])
    return quat


def uniform_quat(random):
    """ Returns a quaternion uniformly at random. Choosing a random axis/angle or even uniformly
    random Euler angles will result in a biased angle rather than a spherically symmetric one.
    See https://en.wikipedia.org/wiki/Rotation_matrix#Uniform_random_rotation_matrices for details.
    """
    w = random.randn(4)
    return quat_normalize(w / np.linalg.norm(w))


def apply_euler_rotations(base_quat, rotation_angles):
    """Apply a sequence of euler angle rotations on to the base quaternion
    """
    new_rot_mat = np.eye(3)
    for rot_angle in rotation_angles:
        new_rot_mat = np.matmul(euler2mat(rot_angle * np.pi / 2.0), new_rot_mat)

    new_rot_mat = np.matmul(quat2mat(base_quat), new_rot_mat)
    new_quat = mat2quat(new_rot_mat)
    return new_quat


def any_orthogonal(vec):
    """ Return any (unit length) vector orthogonal to vec, in a numerically stable way """
    promising_axis = np.eye(3)[np.abs(vec).argmin()]
    non_unit_len = np.cross(vec, promising_axis)

    return non_unit_len / np.linalg.norm(non_unit_len)


def vectors2quat(v_from, v_to):
    """ Define a quaternion rotating along the shortest arc from v_from to v_to """
    q = np.zeros(4)
    dot = np.dot(v_from, v_to)
    v11 = np.dot(v_from, v_from)
    v22 = np.dot(v_to, v_to)
    q[0] = np.sqrt(v11 * v22) + dot
    q[1:4] = np.cross(v_from, v_to)

    if np.linalg.norm(q) < 1e-6:
        # The norm of q is zero if v_from == -v_to, in such case we need to rotate 180 degrees
        # along some not well defined vector orthogonal to both v_from and v_to
        orthogonal = any_orthogonal(v_from)

        q[0] = 0.0  # this is cos(alpha/2) which means rotation 180 deg
        q[1:4] = orthogonal

    return quat_normalize(q / np.linalg.norm(q))


def rot_z_aligned(cube_quat, quat_threshold, include_flip=True):
    """
    Determines if the cube is within quat_threshold of a z-aligned orientation, which means that
    one of the rotatable faces of the **face** cube is on top.
    This means that either the euler angles
    are some pure rotation around z, or they are 180 degree rotation around the x-axis plus some
    rotation around the z.
    """

    cube_angles = quat2euler(cube_quat)
    target_angle = np.eye(3)[-1] * cube_angles

    # if include_flip: True, allow z-axis rotation with top or bottom face as top
    # else, allow z-axis rotation only with top face as top
    if include_flip:
        x_flip = np.asarray([np.pi, 0, 0])
        aligned_angles = [target_angle, target_angle + x_flip]
    else:
        aligned_angles = [target_angle]

    for aligned_angle in aligned_angles:
        aligned_quat = euler2quat(aligned_angle)
        quat_diff = quat_difference(cube_quat, aligned_quat)
        quat_dist = quat_magnitude(quat_diff)
        if quat_dist < quat_threshold:
            return True
    return False


def rot_xyz_aligned(cube_quat, quat_threshold):
    """
    Determines if the cube is within quat_threshold of a xyz-aligned orientation, which means that
    one of the rotatable faces of the **full** cube is on top.

    This means that one of the axes of local coordinate system of the cube is pointing straight up.
    """
    z_up = np.array([0, 0, 1]).reshape(3, 1)
    mtx = quat2mat(cube_quat)
    # Axis that is the closest (by dotproduct) to z-up
    axis_nr = np.abs((z_up.T @ mtx)).argmax()

    # Axis of the cube pointing the closest to the top
    axis = mtx[:, axis_nr]
    axis = axis * np.sign(axis @ z_up)

    # Quaternion representing the rotation from "axis" that is almost up to
    # the actual "up" direction
    difference_quat = vectors2quat(axis, z_up[:, 0])

    return quat_magnitude(difference_quat) < quat_threshold


def random_unity2(random_state):
    """
    Generates a random 3D unit vector (direction) with a uniform spherical distribution
    Algo from http://stackoverflow.com/questions/5408276/python-uniform-spherical-distribution
    """
    phi = random_state.uniform(0, np.pi * 2)
    costheta = random_state.uniform(-1, 1)

    theta = np.arccos(costheta)
    x = np.sin(theta) * np.cos(phi)
    y = np.sin(theta) * np.sin(phi)
    z = np.cos(theta)
    return np.array([x, y, z])
