# Source code for spatialmath.pose2d

    #!/usr/bin/env python3
    # -*- coding: utf-8 -*-

    from collections import UserList
    import numpy as np
    import math

    from spatialmath.base import argcheck
    import spatialmath.base as tr
    from spatialmath import super_pose as sp
    import spatialmath.pose3d as p3

    # ============================== SO2 =====================================#

    class SO2(sp.SMPose):

        # SO2()  identity matrix
        # SO2(angle, unit)
        # SO2( obj )   # deep copy
        # SO2( np )  # make numpy object
        # SO2( nplist )  # make from list of numpy objects

        # constructor needs to take ndarray -> SO2, or list of ndarray -> SO2
        def __init__(self, arg=None, *, unit='rad', check=True):
            """
            Construct new SO(2) object

            :param unit: angular units 'deg' or 'rad' [default] if applicable
            :type unit: str, optional
            :param check: check for valid SO(2) elements if applicable, default to True
            :type check: bool
            :return: SO(2) rotation
            :rtype: SO2 instance

            - ``SO2()`` is an SO2 instance representing a null rotation -- the identity matrix.
            - ``SO2(theta)`` is an SO2 instance representing a rotation by ``theta`` radians.  If ``theta`` is array_like
              `[theta1, theta2, ... thetaN]` then an SO2 instance containing a sequence of N rotations.
            - ``SO2(theta, unit='deg')`` is an SO2 instance representing a rotation by ``theta`` degrees.  If ``theta`` is array_like
              `[theta1, theta2, ... thetaN]` then an SO2 instance containing a sequence of N rotations.
            - ``SO2(R)`` is an SO2 instance with rotation described by the SO(2) matrix R which is a 2x2 numpy array.  If ``check``
              is ``True`` check the matrix belongs to SO(2).
            - ``SO2([R1, R2, ... RN])`` is an SO2 instance containing a sequence of N rotations, each described by an SO(2) matrix
              Ri which is a 2x2 numpy array. If ``check`` is ``True`` then check each matrix belongs to SO(2).
            - ``SO2([X1, X2, ... XN])`` is an SO2 instance containing a sequence of N rotations, where each Xi is an SO2 instance.

            """
            super().__init__()  # activate the UserList semantics

            if arg is None:
                # empty constructor
                if type(self) is SO2:
                    self.data = [np.eye(2)]
            elif argcheck.isvector(arg):
                # SO2(value)
                # SO2(list of values)
                self.data = [tr.rot2(x, unit) for x in argcheck.getvector(arg)]

            elif isinstance(arg, np.ndarray) and arg.shape == (2,2):
                self.data = [arg]
            else:
                super()._arghandler(arg, check=check)

        @classmethod
        def Rand(cls, *, range=[0, 2 * math.pi], unit='rad', N=1):

            r"""
            Construct new SO(2) with random rotation

            :param range: rotation range, defaults to :math:`[0, 2\pi)`.
            :type range: 2-element array-like, optional
            :param unit: angular units as 'deg or 'rad' [default]
            :type unit: str, optional
            :param N: number of random rotations, defaults to 1
            :type N: int
            :return: SO(2) rotation matrix
            :rtype: SO2 instance

            - ``SO2.Rand()`` is a random SO(2) rotation.
            - ``SO2.Rand([-90, 90], unit='deg')`` is a random SO(2) rotation between
              -90 and +90 degrees.
            - ``SO2.Rand(N)`` is a sequence of N random rotations.

            Rotations are uniform over the specified interval.

            """
            rand = np.random.uniform(low=range[0], high=range[1], size=N)  # random values in the range
            return cls([tr.rot2(x) for x in argcheck.getunit(rand, unit)])

        @classmethod
        def Exp(cls, S, check=True):
            """
            Construct new SO(2) rotation matrix from so(2) Lie algebra

            :param S: element of Lie algebra so(2)
            :type S: numpy ndarray
            :param check: check that passed matrix is valid so(2), default True
            :type check: bool
            :return: SO(2) rotation matrix
            :rtype: SO2 instance

            - ``SO2.Exp(S)`` is an SO(2) rotation defined by its Lie algebra
              which is a 2x2 so(2) matrix (skew symmetric)

            :seealso: :func:`spatialmath.base.transforms2d.trexp`, :func:`spatialmath.base.transformsNd.skew`
            """
            if argcheck.ismatrix(S, (-1,2)) and not so2:
                return cls([tr.trexp2(s, check=check) for s in S])
            else:
                return cls(tr.trexp2(S, check=check), check=False)

        @staticmethod
        def isvalid(x):
            """
            Test if matrix is valid SO(2)

            :param x: matrix to test
            :type x: numpy.ndarray
            :return: True if the matrix is a valid element of SO(2), ie. it is a 2x2
                orthonormal matrix with determinant of +1.
            :rtype: bool

            :seealso: :func:`~spatialmath.base.transform3d.isrot`
            """
            return tr.isrot2(x, check=True)

        def inv(self):
            """
            Inverse of SO(2)

            :return: inverse rotation
            :rtype: SO2 instance

            - ``x.inv()`` is the inverse of `x`.

            Notes:

                - for elements of SO(2) this is the transpose.
                - if `x` contains a sequence, returns an `SO2` with a sequence of inverses
            """
            if len(self) == 1:
                return SO2(self.A.T)
            else:
                return SO2([x.T for x in self.A])

        @property
        def R(self):
            """
            SO(2) or SE(2) as rotation matrix

            :return: rotational component
            :rtype: numpy.ndarray, shape=(2,2)

            ``x.R`` returns the rotation matrix, when `x` is `SO2` or `SE2`. If `len(x)` is:

            - 1, return an ndarray with shape=(2,2)
            - N>1, return ndarray with shape=(N,2,2)
            """
            return self.A[:2, :2]

        def theta(self, units='rad'):
            """
            SO(2) as a rotation angle

            :param unit: angular units 'deg' or 'rad' [default]
            :type unit: str, optional
            :return: rotation angle
            :rtype: float or list

            ``x.theta`` is the rotation angle such that `x` is `SO2(x.theta)`.

            """
            if units == 'deg':
                conv = 180.0 / math.pi
            else:
                conv = 1.0

            if len(self) == 1:
                return conv * math.atan2(self.A[1,0], self.A[0,0])
            else:
                return [conv * math.atan2(x.A[1,0], x.A[0,0]) for x in self]

        def SE2(self):
            """
            Create SE(2) from SO(2)

            :return: SE(2) with same rotation but zero translation
            :rtype: SE2 instance

            """
            return SE2(tr.rt2tr(self.A, [0, 0]))

    # ============================== SE2 =====================================#

    class SE2(SO2):
        # constructor needs to take ndarray -> SO2, or list of ndarray -> SO2
        def __init__(self, x=None, y=None, theta=None, *, unit='rad', check=True):
            """
            Construct new SE(2) object

            :param unit: angular units 'deg' or 'rad' [default] if applicable
            :type unit: str, optional
            :param check: check for valid SE(2) elements if applicable, default to True
            :type check: bool
            :return: homogeneous rigid-body transformation matrix
            :rtype: SE2 instance

            - ``SE2()`` is an SE2 instance representing a null motion -- the identity matrix
            - ``SE2(x, y)`` is an SE2 instance representing a pure translation of (``x``, ``y``)
            - ``SE2(t)`` is an SE2 instance representing a pure translation of (``x``, ``y``) where``t``=[x,y] is a 2-element array_like
            - ``SE2(x, y, theta)`` is an SE2 instance representing a translation of (``x``, ``y``) and a rotation of ``theta`` radians
            - ``SE2(x, y, theta, unit='deg')`` is an SE2 instance representing a translation of (``x``, ``y``) and a rotation of ``theta`` degrees
            - ``SE2(t)`` is an SE2 instance representing a translation of (``x``, ``y``) and a rotation of ``theta`` where ``t``=[x,y,theta] is a 3-element array_like
            - ``SE2(T)`` is an SE2 instance with rigid-body motion described by the SE(2) matrix T which is a 3x3 numpy array.  If ``check``
              is ``True`` check the matrix belongs to SE(2).
            - ``SE2([T1, T2, ... TN])`` is an SE2 instance containing a sequence of N rigid-body motions, each described by an SE(2) matrix
              Ti which is a 3x3 numpy array. If ``check`` is ``True`` then check each matrix belongs to SE(2).
            - ``SE2([X1, X2, ... XN])`` is an SE2 instance containing a sequence of N rigid-body motions, where each Xi is an SE2 instance.

            """
            super().__init__()  # activate the UserList semantics

            if x is None and y is None and theta is None:
                # SE2()
                # empty constructor
                self.data = [np.eye(3)]

            elif x is not None:
                if y is not None and theta is None:
                    # SE2(x, y)
                    self.data = [tr.transl2(x, y)]
                elif y is not None and theta is not None:
                    # SE2(x, y, theta)
                    self.data = [tr.trot2(theta, t=[x, y], unit=unit)]
                elif y is None and theta is None:
                    if argcheck.isvector(x, 2):
                        # SE2([x,y])
                        self.data = [tr.transl2(x)]
                    elif argcheck.isvector(x, 3):
                        # SE2([x,y,theta])
                        self.data = [tr.trot2(x[2], t=x[:2], unit=unit)]
                    else:
                        super()._arghandler(x, check=check)
            else:
                raise ValueError('bad arguments to constructor')

        @classmethod
        def Rand(cls, *, xrange=[-1, 1], yrange=[-1, 1], trange=[0, 2 * math.pi], unit='rad', N=1):
            r"""
            Construct a new random SE(2)

            :param xrange: x-axis range [min,max], defaults to [-1, 1]
            :type xrange: 2-element sequence, optional
            :param yrange: y-axis range [min,max], defaults to [-1, 1]
            :type yrange: 2-element sequence, optional
            :param trange: theta range [min,max], defaults to :math:`[0, 2\pi)`
            :type yrange: 2-element sequence, optional
            :param N: number of random rotations, defaults to 1
            :type N: int
            :return: homogeneous rigid-body transformation matrix
            :rtype: SE2 instance

            Return an SE2 instance with random rotation and translation.

            - ``SE2.Rand()`` is a random SE(2) rotation.
            - ``SE2.Rand(N)`` is an SE2 object containing a sequence of N random
              poses.

            Example, create random ten vehicles in the xy-plane::

                >>> x = SE3.Rand(N=10, xrange=[-2,2], yrange=[-2,2])
                >>> len(x)
                10

            """
            x = np.random.uniform(low=xrange[0], high=xrange[1], size=N)  # random values in the range
            y = np.random.uniform(low=yrange[0], high=yrange[1], size=N)  # random values in the range
            theta = np.random.uniform(low=trange[0], high=trange[1], size=N)  # random values in the range
            return cls([tr.trot2(t, t=[x, y]) for (t, x, y) in zip(x, y, argcheck.getunit(theta, unit))])

        @classmethod
        def Exp(cls, S, check=True, se2=True):
            """
            Construct a new SE(2) from se(2) Lie algebra

            :param S: element of Lie algebra se(2)
            :type S: numpy ndarray
            :param check: check that passed matrix is valid se(2), default True
            :type check: bool
            :param se2: input is an se(2) matrix (default True)
            :type se2: bool
            :return: homogeneous transform matrix
            :rtype: SE2 instance

            - ``SE2.Exp(S)`` is an SE(2) rotation defined by its Lie algebra
              which is a 3x3 se(2) matrix (skew symmetric)
            - ``SE2.Exp(t)`` is an SE(2) rotation defined by a 3-element twist
              vector array_like (the unique elements of the se(2) skew-symmetric matrix)
            - ``SE2.Exp(T)`` is a sequence of SE(2) rigid-body motions defined by an Nx3 matrix of twist vectors, one per row.

            Note:

            - an input 3x3 matrix is ambiguous, it could be the first or third case above. In this case the argument ``se2`` is the decider.

            :seealso: :func:`spatialmath.base.transforms2d.trexp`, :func:`spatialmath.base.transformsNd.skew`
            """
            if isinstance(S, np.ndarray) and S.shape[1] == 3 and not se2:
                return cls([tr.trexp2(s) for s in S])
            else:
                return cls(tr.trexp2(S), check=False)

        @staticmethod
        def isvalid(x):
            """
            Test if matrix is valid SE(2)

            :param x: matrix to test
            :type x: numpy.ndarray
            :return: true if the matrix is a valid element of SE(2), ie. it is a
                     3x3 homogeneous rigid-body transformation matrix.
            :rtype: bool

            :seealso: :func:`~spatialmath.base.transform2d.ishom`
            """
            return tr.ishom2(x, check=True)

        @property
        def t(self):
            """
            Translational component of SE(2)

            :param self: SE(2)
            :type self: SE2 instance
            :return: translational component
            :rtype: numpy.ndarray

            ``x.t`` is the translational vector component.  If ``len(x)`` is:

            - 1, return an ndarray with shape=(2,)
            - N>1, return an ndarray with shape=(N,2)
            """
            if len(self) == 1:
                return self.A[:2, 2]
            else:
                return np.array([x[:2, 2] for x in self.A])

        def xyt(self):
            r"""
            SE(2) as a configuration vector

            :return: An array :math:`[x, y, \theta]`
            :rtype: numpy.ndarray

            ``x.xyt`` is the rigidbody motion in minimal form as a translation and rotation expressed
            in vector form as :math:`[x, y, \theta]`.  If ``len(x)`` is:

            - 1, return an ndarray with shape=(3,)
            - N>1, return an ndarray with shape=(N,3)
            """
            if len(self) == 1:
                return np.r_[self.t, self.theta()]
            else:
                return [np.r_[x.t, x.theta()] for x in self]

        def inv(self):
            r"""
            Inverse of SE(2)

            :param self: pose
            :type self: SE2 instance
            :return: inverse
            :rtype: SE2

            Notes:

                - for elements of SE(2) this takes into account the matrix structure :math:`T^{-1} = \left[ \begin{array}{cc} R & t \\ 0 & 1 \end{array} \right], T^{-1} = \left[ \begin{array}{cc} R^T & -R^T t \\ 0 & 1 \end{array} \right]`
                - if `x` contains a sequence, returns an `SE2` with a sequence of inverses

            """
            if len(self) == 1:
                return SE2(tr.rt2tr(self.R.T, -self.R.T @ self.t))
            else:
                return SE2([tr.rt2tr(x.R.T, -x.R.T @ x.t) for x in self])

        def SE3(self, z=0):
            """
            Create SE(3) from SE(2)

            :param z: default z coordinate, defaults to 0
            :type z: float
            :return: SE(2) with same rotation but zero translation
            :rtype: SE2 instance

            "Lifts" 2D rigid-body motion to 3D, rotation in the xy-plane (about the z-axis) and
            z-coordinate is settable.

            """
            def lift3(x):
                y = np.eye(4)
                y[:2,:2] = x.A[:2,:2]
                y[:2,3] = x.A[:2,2]
                y[2,3] = z
                return y
            return p3.SE3([lift3(x) for x in self])

    if __name__ == '__main__':  # pragma: no cover

        import pathlib
        import os.path

        exec(open(os.path.join(pathlib.Path(__file__).parent.absolute(), "test_pose2d.py")).read())