from sympy.core.basic import Basic
from sympy import (sympify, eye, sin, cos, rot_axis1, rot_axis2,
                   rot_axis3, ImmutableMatrix as Matrix, Symbol)
from sympy.core.cache import cacheit
import sympy.vector


class Orienter(Basic):
    """
    Super-class for all orienter classes.
    """

    def rotation_matrix(self):
        """
        The rotation matrix corresponding to this orienter
        instance.
        """
        return self._parent_orient


class AxisOrienter(Orienter):
    """
    Class to denote an axis orienter.
    """

    def __new__(cls, angle, axis):
        if not isinstance(axis, sympy.vector.Vector):
            raise TypeError("axis should be a Vector")
        angle = sympify(angle)

        obj = super(AxisOrienter, cls).__new__(cls, angle,
                                               axis)
        obj._angle = angle
        obj._axis = axis

        return obj

    def __init__(self, angle, axis):
        """
        Axis rotation is a rotation about an arbitrary axis by
        some angle. The angle is supplied as a SymPy expr scalar, and
        the axis is supplied as a Vector.

        Parameters
        ==========

        angle : Expr
            The angle by which the new system is to be rotated

        axis : Vector
            The axis around which the rotation has to be performed

        Examples
        ========

        >>> from sympy.vector import CoordSysCartesian
        >>> from sympy import symbols
        >>> q1 = symbols('q1')
        >>> N = CoordSysCartesian('N')
        >>> from sympy.vector import AxisOrienter
        >>> orienter = AxisOrienter(q1, N.i + 2 * N.j)
        >>> B = N.orient_new('B', (orienter, ))

        """
        #Dummy initializer for docstrings
        pass

    @cacheit
    def rotation_matrix(self, system):
        """
        The rotation matrix corresponding to this orienter
        instance.

        Parameters
        ==========

        system : CoordSysCartesian
            The coordinate system wrt which the rotation matrix
            is to be computed
        """

        axis = sympy.vector.express(self.axis, system).normalize()
        axis = axis.to_matrix(system)
        theta = self.angle
        parent_orient = ((eye(3) - axis * axis.T) * cos(theta) +
                Matrix([[0, -axis[2], axis[1]],
                        [axis[2], 0, -axis[0]],
                    [-axis[1], axis[0], 0]]) * sin(theta) +
                         axis * axis.T)
        parent_orient = parent_orient.T
        return parent_orient

    @property
    def angle(self):
        return self._angle

    @property
    def axis(self):
        return self._axis


class ThreeAngleOrienter(Orienter):
    """
    Super-class for Body and Space orienters.
    """

    def __new__(cls, angle1, angle2, angle3, rot_order):
        approved_orders = ('123', '231', '312', '132', '213',
                           '321', '121', '131', '212', '232',
                           '313', '323', '')
        original_rot_order = rot_order
        rot_order = str(rot_order).upper()
        if not (len(rot_order) == 3):
            raise TypeError('rot_order should be a str of length 3')
        rot_order = [i.replace('X', '1') for i in rot_order]
        rot_order = [i.replace('Y', '2') for i in rot_order]
        rot_order = [i.replace('Z', '3') for i in rot_order]
        rot_order = ''.join(rot_order)
        if not rot_order in approved_orders:
            raise TypeError('Invalid rot_type parameter')
        a1 = int(rot_order[0])
        a2 = int(rot_order[1])
        a3 = int(rot_order[2])
        angle1 = sympify(angle1)
        angle2 = sympify(angle2)
        angle3 = sympify(angle3)
        if cls._in_order:
            parent_orient = (_rot(a1, angle1) *
                             _rot(a2, angle2) *
                             _rot(a3, angle3))
        else:
            parent_orient = (_rot(a3, angle3) *
                             _rot(a2, angle2) *
                             _rot(a1, angle1))
        parent_orient = parent_orient.T

        obj = super(ThreeAngleOrienter, cls).__new__(
            cls, angle1, angle2, angle3, Symbol(original_rot_order))
        obj._angle1 = angle1
        obj._angle2 = angle2
        obj._angle3 = angle3