""" 
This modules contains functions to create and transform rotation matrices
and homogeneous tranformation matrices.

Vector arguments are what numpy refers to as ``array_like`` and can be a list,
tuple, numpy array, numpy row vector or numpy column vector.

Versions: 
    
    1. Luis Fernando Lara Tobar and Peter Corke, 2008
    2. Josh Carrigg Hodson, Aditya Dua, Chee Ho Chan, 2017
    3. Peter Corke, 2020
"""

import sys
import math
import numpy as np
import numpy.matlib as matlib
from spatialmath.base.vectors import *
from spatialmath.base import transforms2d as t2d
from spatialmath.base import transforms3d as t3d
from spatialmath.base import argcheck

    
_eps = np.finfo(np.float64).eps


# ---------------------------------------------------------------------------------------#
def r2t(R, check=False):
    """
    Convert SO(n) to SE(n)

    :param R: rotation matrix
    :param check: check if rotation matrix is valid (default False, no check)
    :return: homogeneous transformation matrix
    :rtype: numpy.ndarray, shape=(3,3) or (4,4)

    ``T = r2t(R)`` is an SE(2) or SE(3) homogeneous transform equivalent to an
    SO(2) or SO(3) orthonormal rotation matrix ``R`` with a zero translational
    component
    
    - if ``R`` is 2x2 then ``T`` is 3x3: SO(2) -> SE(2)
    - if ``R`` is 3x3 then ``T`` is 4x4: SO(3) -> SE(3)
    
    :seealso: t2r, rt2tr
    """
    
    assert isinstance(R, np.ndarray)
    dim = R.shape
    assert dim[0] == dim[1], 'Matrix must be square'
    
    if check and np.abs(np.linalg.det(R) - 1) < 100*_eps:
        raise ValueError('Invalid rotation matrix ')
    
    T  = np.pad( R, (0,1), mode='constant' )
    T[-1,-1] = 1.0
        
    return T


# ---------------------------------------------------------------------------------------#
def t2r(T, check=False):
    """
    Convert SE(n) to SO(n)

    :param T: homogeneous transformation matrix
    :param check: check if rotation matrix is valid (default False, no check)
    :return: rotation matrix
    :rtype: numpy.ndarray, shape=(2,2) or (3,3)


    ``R = T2R(T)`` is the orthonormal rotation matrix component of homogeneous
    transformation matrix ``T``
    
    - if ``T`` is 3x3 then ``R`` is 2x2: SE(2) -> SO(2)
    - if ``T`` is 4x4 then ``R`` is 3x3: SE(3) -> SO(3)
    
    Any translational component of T is lost.
    
    :seealso: r2t, tr2rt
    """
    assert isinstance(T, np.ndarray)
    dim = T.shape
    assert dim[0] == dim[1], 'Matrix must be square'
    
    if dim[0] == 3:
        R = T[:2,:2]
    elif dim[0] == 4:
        R = T[:3,:3]
    else:
        raise ValueError('Value must be a rotation matrix')
    
    if check and isR(R):
            raise ValueError('Invalid rotation matrix')

    return R

# ---------------------------------------------------------------------------------------#
def tr2rt(T, check=False):
    """
    Convert SE(3) to SO(3) and translation

    :param T: homogeneous transform matrix
    :param check: check if rotation matrix is valid (default False, no check)
    :return: Rotation matrix and translation vector
    :rtype: tuple: numpy.ndarray, shape=(2,2) or (3,3); numpy.ndarray, shape=(2,) or (3,)

    (R,t) = tr2rt(T) splits a homogeneous transformation matrix (NxN) into an orthonormal
    rotation matrix R (MxM) and a translation vector T (Mx1), where N=M+1.
        
    - if ``T`` is 3x3 - in SE(2) - then ``R`` is 2x2 and ``t`` is 2x1.
    - if ``T`` is 4x4 - in SE(3) - then ``R`` is 3x3 and ``t`` is 3x1.
    
    :seealso: rt2tr, tr2r
    """
    dim = T.shape
    assert dim[0] == dim[1], 'Matrix must be square'
    
    if dim[0] == 3:
        R = t2r(T, check)
        t = T[:2,2]
    elif dim[0] == 4:
        R = t2r(T, check)
        t = T[:3,3]
    else:
        raise ValueError('T must be an SE2 or SE3 homogeneous transformation matrix')
    
    return [R, t]

# ---------------------------------------------------------------------------------------#
def rt2tr(R, t, check=False):
    """
    Convert SO(3) and translation to SE(3)

    :param R: rotation matrix
    :param t: translation vector
    :param check: check if rotation matrix is valid (default False, no check)
    :return: homogeneous transform
    :rtype: numpy.ndarray, shape=(3,3) or (4,4)

    ``T = rt2tr(R, t)`` is a homogeneous transformation matrix (N+1xN+1) formed from an
    orthonormal rotation matrix ``R`` (NxN) and a translation vector ``t``
    (Nx1).
        
    - If ``R`` is 2x2 and ``t`` is 2x1, then ``T`` is 3x3
    - If ``R`` is 3x3 and ``t`` is 3x1, then ``T`` is 4x4
    
    :seealso: rt2m, tr2rt, r2t
    """
    t = argcheck.getvector(t, dim=None, out='array')
    if R.shape[0] != t.shape[0]:
        raise ValueError("R and t must have the same number of rows")
    if check and np.abs(np.linalg.det(R) - 1) < 100*_eps:
            raise ValueError('Invalid rotation matrix')
            
    if R.shape == (2, 2):
        T = np.eye(3)
        T[:2,:2] = R
        T[:2,2] = t
    elif R.shape == (3, 3):
        T = np.eye(4)
        T[:3,:3] = R
        T[:3,3] = t
    else:
        raise ValueError('R must be an SO2 or SO3 rotation matrix')
    
    return T

# ---------------------------------------------------------------------------------------#
def rt2m(R, t, check=False):
    """
    Pack rotation and translation to matrix

    :param R: rotation matrix
    :param t: translation vector
    :param check: check if rotation matrix is valid (default False, no check)
    :return: homogeneous transform
    :rtype: numpy.ndarray, shape=(3,3) or (4,4)

    ``T = rt2m(R, t)`` is a matrix (N+1xN+1) formed from a matrix ``R`` (NxN) and a vector ``t``
    (Nx1).  The bottom row is all zeros.
        
    - If ``R`` is 2x2 and ``t`` is 2x1, then ``T`` is 3x3
    - If ``R`` is 3x3 and ``t`` is 3x1, then ``T`` is 4x4
    
    :seealso: rt2tr, tr2rt, r2t
    """
    t = argcheck.getvector(t, dim=None, out='array')
    if R.shape[0] != t.shape[0]:
        raise ValueError("R and t must have the same number of rows")
    if check and np.abs(np.linalg.det(R) - 1) < 100*_eps:
            raise ValueError('Invalid rotation matrix')
            
    if R.shape == (2, 2):
        T = np.zeros((3,3))
        T[:2,:2] = R
        T[:2,2] = t
    elif R.shape == (3, 3):
        T = np.zeros((4,4))
        T[:3,:3] = R
        T[:3,3] = t
    else:
        raise ValueError('R must be an SO2 or SO3 rotation matrix')
    
    return T

#======================= predicates

def isR(R, tol = 10):
    r"""
    Test if matrix belongs to SO(n)
    
    :param R: matrix to test
    :type R: numpy.ndarray
    :param tol: tolerance in units of eps
    :type tol: float
    :return: whether matrix is a proper orthonormal rotation matrix
    :rtype: bool
    
    Checks orthogonality, ie. :math:`{\bf R} {\bf R}^T = {\bf I}` and :math:`\det({\bf R}) > 0`.
    For the first test we check that the norm of the residual is less than ``tol * eps``.
    
    :seealso: isrot2, isrot
    """
    return np.linalg.norm( R@R.T - np.eye(R.shape[0]) ) < tol*_eps \
        and np.linalg.det( R@R.T ) > 0



def isskew(S, tol = 10):
    r"""
    Test if matrix belongs to so(n)
    
    :param S: matrix to test
    :type S: numpy.ndarray
    :param tol: tolerance in units of eps
    :type tol: float
    :return: whether matrix is a proper skew-symmetric matrix
    :rtype: bool
    
    Checks skew-symmetry, ie. :math:`{\bf S} + {\bf S}^T = {\bf 0}`.
    We check that the norm of the residual is less than ``tol * eps``.
    
    :seealso: isskewa
    """
    return np.linalg.norm(S + S.T) < tol*_eps

def isskewa(S, tol = 10):
    r"""
    Test if matrix belongs to se(n)
    
    :param S: matrix to test
    :type S: numpy.ndarray
    :param tol: tolerance in units of eps
    :type tol: float
    :return: whether matrix is a proper skew-symmetric matrix
    :rtype: bool
    
    Check if matrix is augmented skew-symmetric, ie. the top left (n-1xn-1) partition ``S`` is
    skew-symmetric :math:`{\bf S} + {\bf S}^T = {\bf 0}`, and the bottom row is zero
    We check that the norm of the residual is less than ``tol * eps``.
    
    :seealso: isskew
    """
    return np.linalg.norm(S[0:-1,0:-1] + S[0:-1,0:-1].T) < tol*_eps \
        and np.all(S[-1,:] == 0)

def iseye(S, tol = 10):
    """
    Test if matrix is identity
    
    :param S: matrix to test
    :type S: numpy.ndarray
    :param tol: tolerance in units of eps
    :type tol: float
    :return: whether matrix is a proper skew-symmetric matrix
    :rtype: bool
    
    Check if matrix is an identity matrix. We test that the trace tom row is zero
    We check that the norm of the residual is less than ``tol * eps``.
    
    :seealso: isskew, isskewa
    """
    s = S.shape
    if len(s) != 2 or s[0] != s[1]:
        return False  # not a square matrix
    return norm(S - np.eye(s[0])) < tol * _eps
        
    
    
#========================= angle sequences



# ---------------------------------------------------------------------------------------#
def skew(v):
    r"""
    Create skew-symmetric metrix from vector

    :param v: 1- or 3-vector
    :type v: array_like
    :return: skew-symmetric matrix in so(2) or so(3)
    :rtype: numpy.ndarray, shape=(2,2) or (3,3)
    :raises: ValueError
    
    ``skew(V)`` is a skew-symmetric matrix formed from the elements of ``V``.

    - ``len(V)``  is 1 then ``S`` = :math:`\left[ \begin{array}{cc} 0 & -v \\ v & 0 \end{array} \right]`
    - ``len(V)`` is 3 then ``S`` = :math:`\left[ \begin{array}{ccc} 0 & -v_z & v_y \\ v_z & 0 & -v_x \\ -v_y & v_x & 0\end{array} \right]`

    Notes:
        
    - This is the inverse of the function ``vex()``.
    - These are the generator matrices for the Lie algebras so(2) and so(3).
    
    :seealso: vex, skewa
    """
    v = argcheck.getvector(v, None, 'sequence')
    if len(v) == 1:
        s = np.array([
                [0, -v[0]], 
                [v[0], 0]   ])
    elif len(v) == 3:
        s = np.array([
                [ 0,   -v[2], v[1]], 
                [ v[2], 0,   -v[0]], 
                [-v[1], v[0], 0]    ])
    else:
        raise AttributeError("argument must be a 1- or 3-vector")
        
    return s

# ---------------------------------------------------------------------------------------#
def vex(s):
    r"""
    Convert skew-symmetric matrix to vector

    :param s: skew-symmetric matrix
    :type s: numpy.ndarray, shape=(2,2) or (3,3)
    :return: vector of unique values
    :rtype: numpy.ndarray, shape=(1,) or (3,)
    :raises: ValueError

    ``vex(S)`` is the vector which has the corresponding skew-symmetric matrix ``S``.
    
    - ``S`` is 2x2 - so(2) case - where ``S`` :math:`= \left[ \begin{array}{cc} 0 & -v \\ v & 0 \end{array} \right]` then return :math:`[v]`
    - ``S`` is 3x3 - so(3) case -  where ``S`` :math:`= \left[ \begin{array}{ccc} 0 & -v_z & v_y \\ v_z & 0 & -v_x \\ -v_y & v_x & 0\end{array} \right]` then return :math:`[v_x, v_y, v_z]`.
    
    Notes:
        
    - This is the inverse of the function ``skew()``.
    - Only rudimentary checking (zero diagonal) is done to ensure that the matrix
      is actually skew-symmetric.
    - The function takes the mean of the two elements that correspond to each unique
      element of the matrix.
      
    :seealso: skew, vexa
    """
    if s.shape == (3,3):
        return 0.5 * np.array([s[2, 1] - s[1, 2], s[0, 2] - s[2, 0], s[1, 0] - s[0, 1]])
    elif s.shape == (2, 2):
        return 0.5 * np.array([s[1, 0] - s[0, 1]])
    else:
        raise ValueError("Argument must be 2x2 or 3x3 matrix")

# ---------------------------------------------------------------------------------------#
def skewa(v):
    r"""
    Create augmented skew-symmetric metrix from vector

    :param v: 3- or 6-vector
    :type v: array_like
    :return: augmented skew-symmetric matrix in se(2) or se(3)
    :rtype: numpy.ndarray, shape=(3,3) or (4,4)
    :raises: ValueError
    
    ``skewa(V)`` is an augmented skew-symmetric matrix formed from the elements of ``V``.

    - ``len(V)`` is 3 then S = :math:`\left[ \begin{array}{ccc} 0 & -v_3 & v_1 \\ v_3 & 0 & v_2 \\ 0 & 0 & 0 \end{array} \right]`
    - ``len(V)`` is 6 then S = :math:`\left[ \begin{array}{cccc} 0 & -v_6 & v_5 & v_1 \\ v_6 & 0 & -v_4 & v_2 \\ -v_5 & v_4 & 0 & v_3 \\ 0 & 0 & 0 & 0 \end{array} \right]`

    Notes:
        
    - This is the inverse of the function ``vexa()``.
    - These are the generator matrices for the Lie algebras se(2) and se(3).
    - Map twist vectors in 2D and 3D space to se(2) and se(3).
    
    :seealso: vexa, skew
    """
    
    v = argcheck.getvector(v, None, 'sequence')
    if len(v) == 3:
        omega = np.zeros((3,3))
        omega[:2,:2] = skew(v[2])
        omega[:2,2] = v[0:2]
        return omega
    elif len(v) == 6:
        omega = np.zeros((4,4))
        omega[:3,:3] = skew(v[3:6])
        omega[:3,3] = v[0:3]
        return omega
    else:
        raise AttributeError("expecting a 3- or 6-vector")

def vexa(Omega):
    r"""
    Convert skew-symmetric matrix to vector

    :param s: augmented skew-symmetric matrix
    :type s: numpy.ndarray, shape=(3,3) or (4,4)
    :return: vector of unique values
    :rtype: numpy.ndarray, shape=(3,) or (6,)
    :raises: ValueError

    ``vex(S)`` is the vector which has the corresponding skew-symmetric matrix ``S``.
    
    - ``S`` is 3x3 - se(2) case - where ``S`` :math:`= \left[ \begin{array}{ccc} 0 & -v_3 & v_1 \\ v_3 & 0 & v_2 \\ 0 & 0 & 0 \end{array} \right]` then return :math:`[v_1, v_2, v_3]`.
    - ``S`` is 4x4 - se(3) case -  where ``S`` :math:`= \left[ \begin{array}{cccc} 0 & -v_6 & v_5 & v_1 \\ v_6 & 0 & -v_4 & v_2 \\ -v_5 & v_4 & 0 & v_3 \\ 0 & 0 & 0 & 0 \end{array} \right]` then return :math:`[v_1, v_2, v_3, v_4, v_5, v_6]`.
    
    
    Notes:
        
    - This is the inverse of the function ``skewa``.
    - Only rudimentary checking (zero diagonal) is done to ensure that the matrix
      is actually skew-symmetric.
    - The function takes the mean of the two elements that correspond to each unique
      element of the matrix.
      
    :seealso: skewa, vex
    """
    if Omega.shape == (4,4):
        return np.hstack( (t3d.transl(Omega), vex(t2r(Omega))) )
    elif Omega.shape == (3,3):
        return np.hstack( (t2d.transl2(Omega), vex(t2r(Omega))) )
    else:
        raise AttributeError("expecting a 3x3 or 4x4 matrix")
        


    
def _rodrigues(w, theta):
    """
    Rodrigues' formula for rotation
    
    :param w: rotation vector
    :type w: array_like
    :param theta: rotation angle
    :type theta: float or None
    """
    w = argcheck.getvector(w)
    if iszerovec(w):
        # for a zero so(n) return unit matrix, theta not relevant
        if len(w) == 1:
            return np.eye(2)
        else:
            return np.eye(3)
    if theta is None:
        theta = norm(w)
        w = unitvec(w)

    skw = skew(w)
    return np.eye(skw.shape[0]) + math.sin(theta) * skw + (1.0 - math.cos(theta)) * skw @ skw

def h2e(v):
    """
    Convert from homogeneous to Euclidean form
    
    :param v: homogeneous vector or matrix
    :type v: array_like
    :return: Euclidean vector
    :rtype: numpy.ndarray

    - If ``v`` is an array, shape=(N,), return an array shape=(N-1,) where the elements have
      all been scaled by the last element of ``v``.
    - If ``v`` is a matrix, shape=(N,M), return a matrix shape=(N-1,N), where each column has
      been scaled by its last element.
      
    :seealso: e2h
    """
    if argcheck.isvector(v):
        # dealing with shape (N,) array
        v = argcheck.getvector(v)
        return v[0:-1] / v[-1]
    elif isinstance(v, np.ndarray) and len(v.shape) == 2:
        # dealing with matrix
        return v[:-1,:] / matlib.repmat(v[-1,:], v.shape[0]-1, 1)
    
def e2h(v):
    """
    Convert from Euclidean to homogeneous form
    
    :param v: Euclidean vector or matrix
    :type v: array_like
    :return: homogeneous vector
    :rtype: numpy.ndarray

    - If ``v`` is an array, shape=(N,), return an array shape=(N+1,) where a value of 1 has
      been appended
    - If ``v`` is a matrix, shape=(N,M), return a matrix shape=(N+1,N), where each column has
      been appended with a value of 1, ie. a row of ones has been appended to the matrix.
      
    :seealso: e2h
    """
    if argcheck.isvector(v):
        # dealing with shape (N,) array
        v = argcheck.getvector(v)
        return np.r_[v, 1]
    elif isinstance(v, np.ndarray) and len(v.shape) == 2:
        # dealing with matrix  
        return np.vstack([v, np.ones((1,v.shape[1]))])
        

if __name__ == '__main__':  # pragma: no cover
    import pathlib
    import os.path
    
    exec(open(os.path.join(pathlib.Path(__file__).parent.absolute(), "test_transforms.py")).read() )
    
    



