import sympy as sp


class SymQuaternion:
    def __init__(self, angle=None, axis=None, w=None, x=None, y=None, z=None):
        """
        Initialize a quaternion in multiple ways:

        1. From angle-axis: SymQuaternion(angle=theta, axis=[x, y, z])
        2. From components: SymQuaternion(w=w, x=x, y=y, z=z)
           axis can be list/tuple or sp.Matrix(3,1) or sp.Matrix(3,)

        If no arguments provided, creates an identity quaternion.
        """
        if angle is not None and axis is not None:
            # Create from angle-axis representation
            theta = sp.sympify(angle)

            # Extract axis components from Matrix or iterable
            axis_vec = self._extract_vector_components(axis, "axis")

            # Normalize axis
            norm = sp.sqrt(sum(v**2 for v in axis_vec))
            if norm != 0:
                axis_vec = [v / norm for v in axis_vec]

            self.w = sp.cos(theta / 2)
            self.x = axis_vec[0] * sp.sin(theta / 2)
            self.y = axis_vec[1] * sp.sin(theta / 2)
            self.z = axis_vec[2] * sp.sin(theta / 2)
        else:
            # Create from components or identity
            self.w = sp.sympify(w if w is not None else 1)
            self.x = sp.sympify(x if x is not None else 0)
            self.y = sp.sympify(y if y is not None else 0)
            self.z = sp.sympify(z if z is not None else 0)

    def _extract_vector_components(self, vector, name="vector"):
        """
        Extract vector components from either sp.Matrix or iterable.
        Returns: list of 3 components
        """
        if isinstance(vector, sp.Matrix):
            # Handle sp.Matrix
            if vector.shape == (3, 1) or vector.shape == (3,):
                return [sp.sympify(vector[i]) for i in range(3)]
            elif vector.shape == (1, 3):
                return [sp.sympify(vector[0, i]) for i in range(3)]
            else:
                raise ValueError(
                    f"{name} must be 3x1, 1x3, or (3,) matrix, got shape {vector.shape}"
                )
        elif isinstance(vector, (list, tuple)):
            # Handle list/tuple
            if len(vector) != 3:
                raise ValueError(f"{name} must have 3 components, got {len(vector)}")
            return [sp.sympify(v) for v in vector]
        else:
            raise TypeError(
                f"{name} must be sp.Matrix, list, or tuple, got {type(vector)}"
            )

    def __repr__(self):
        return f"SymQuaternion(w={self.w}, x={self.x}, y={self.y}, z={self.z})"

    def __str__(self):
        return f"{self.w} + ({self.x})i + ({self.y})j + ({self.z})k"

    def vector_product(self, other):
        """Cross product of vector parts (returns vector as list)"""
        return [
            self.y * other.z - self.z * other.y,
            self.z * other.x - self.x * other.z,
            self.x * other.y - self.y * other.x,
        ]

    def dot_product(self, other):
        """Dot product of vector parts"""
        return self.x * other.x + self.y * other.y + self.z * other.z

    def __mul__(self, other):
        """Quaternion multiplication"""
        if isinstance(other, SymQuaternion):
            # Hamilton product
            w = self.w * other.w - self.dot_product(other)
            cross = self.vector_product(other)
            x = self.w * other.x + other.w * self.x + cross[0]
            y = self.w * other.y + other.w * self.y + cross[1]
            z = self.w * other.z + other.w * self.z + cross[2]
            return SymQuaternion(w=w, x=x, y=y, z=z)
        elif isinstance(other, (int, float, sp.Basic)):
            # Scalar multiplication
            scalar = sp.sympify(other)
            return SymQuaternion(
                w=self.w * scalar,
                x=self.x * scalar,
                y=self.y * scalar,
                z=self.z * scalar,
            )
        else:
            raise TypeError(f"Unsupported type for multiplication: {type(other)}")

    def __rmul__(self, other):
        """Reverse multiplication for scalars"""
        if isinstance(other, (int, float, sp.Basic)):
            return self * other
        return NotImplemented

    def conjugate(self):
        """Return conjugate quaternion"""
        return SymQuaternion(w=self.w, x=-self.x, y=-self.y, z=-self.z)

    def norm(self):
        """Return norm (magnitude squared)"""
        return self.w**2 + self.x**2 + self.y**2 + self.z**2

    def inv(self):
        """Return inverse quaternion"""
        n = self.norm()
        if n == 0:
            raise ValueError("Cannot invert quaternion with zero norm")
        return SymQuaternion(w=self.w / n, x=-self.x / n, y=-self.y / n, z=-self.z / n)

    def normalize(self):
        """Return normalized (unit) quaternion"""
        n = sp.sqrt(self.norm())
        if n == 0:
            raise ValueError("Cannot normalize zero quaternion")
        return SymQuaternion(w=self.w / n, x=self.x / n, y=self.y / n, z=self.z / n)

    def rotate_vector(self, vector):
        """
        Rotate a vector using the quaternion.
        vector: sp.Matrix(3,1) of 3 symbolic expressions
        Returns: rotated vector as sp.Matrix(3,1)
        """
        # Validate input
        if not isinstance(vector, sp.Matrix):
            raise TypeError(f"vector must be sp.Matrix, got {type(vector)}")

        if not (vector.shape == (3, 1) or vector.shape == (3,)):
            raise ValueError(
                f"vector must be 3x1 or (3,) matrix, got shape {vector.shape}"
            )

        # Extract vector components
        v_components = [sp.sympify(vector[i]) for i in range(3)]

        # Create pure quaternion from vector
        v_quat = SymQuaternion(
            w=0, x=v_components[0], y=v_components[1], z=v_components[2]
        )

        # Perform rotation: q * v * q⁻¹
        result = self * v_quat * self.inv()

        # Return vector part as Matrix(3,1)
        return sp.Matrix([result.x, result.y, result.z])

    def rotate_basis(self, vector):
        """
        Rotate basis (inverse rotation of rotate_vector).
        vector: sp.Matrix(3,1) of 3 symbolic expressions
        Returns: rotated basis vector as sp.Matrix(3,1)
        """
        # Validate input
        if not isinstance(vector, sp.Matrix):
            raise TypeError(f"vector must be sp.Matrix, got {type(vector)}")

        if not (vector.shape == (3, 1) or vector.shape == (3,)):
            raise ValueError(
                f"vector must be 3x1 or (3,) matrix, got shape {vector.shape}"
            )

        # Extract vector components
        v_components = [sp.sympify(vector[i]) for i in range(3)]

        # Create pure quaternion from vector
        v_quat = SymQuaternion(
            w=0, x=v_components[0], y=v_components[1], z=v_components[2]
        )

        # Perform inverse rotation: q⁻¹ * v * q
        result = self.inv() * v_quat * self

        # Return vector part as Matrix(3,1)
        return sp.Matrix([result.x, result.y, result.z])

    def to_rotation_matrix(self):
        """
        Convert quaternion to 3x3 rotation matrix
        Returns: sympy.Matrix of shape (3, 3)
        """
        w, x, y, z = self.w, self.x, self.y, self.z

        # Normalize if not already normalized
        n = self.norm()
        if n != 1:
            w, x, y, z = w / sp.sqrt(n), x / sp.sqrt(n), y / sp.sqrt(n), z / sp.sqrt(n)

        # Using the same formula structure as NumQuaternion.get_matrix()
        return sp.Matrix(
            [
                [1 - 2 * y**2 - 2 * z**2, 2 * x * y - 2 * z * w, 2 * x * z + 2 * y * w],
                [2 * x * y + 2 * z * w, 1 - 2 * x**2 - 2 * z**2, 2 * y * z - 2 * x * w],
                [2 * x * z - 2 * y * w, 2 * y * z + 2 * x * w, 1 - 2 * x**2 - 2 * y**2],
            ]
        )

    def simplify(self):
        """Simplify all components"""
        return SymQuaternion(
            w=sp.simplify(self.w),
            x=sp.simplify(self.x),
            y=sp.simplify(self.y),
            z=sp.simplify(self.z),
        )

    def expand(self):
        """Expand all components"""
        return SymQuaternion(
            w=sp.expand(self.w),
            x=sp.expand(self.x),
            y=sp.expand(self.y),
            z=sp.expand(self.z),
        )

    def eval(self, subs_dict=None, **kwargs):
        """
        Evaluate/substitute values for symbols
        subs_dict: dictionary of substitutions
        **kwargs: keyword arguments for substitutions
        """
        if subs_dict is None:
            subs_dict = {}
        subs_dict.update(kwargs)

        return SymQuaternion(
            w=self.w.subs(subs_dict),
            x=self.x.subs(subs_dict),
            y=self.y.subs(subs_dict),
            z=self.z.subs(subs_dict),
        )

    def to_angle_axis(self, return_matrix=False):
        """
        Convert to angle-axis representation.
        return_matrix: if True, axis is returned as sp.Matrix(3,1)
        Returns: (angle, axis) where axis is list or sp.Matrix
        """
        # Normalize the quaternion
        q = self.normalize()

        # Calculate angle (w = cos(θ/2))
        angle = 2 * sp.acos(q.w)

        # Calculate axis from vector part
        sin_half_angle = sp.sin(angle / 2)
        if sin_half_angle != 0:
            axis = [q.x / sin_half_angle, q.y / sin_half_angle, q.z / sin_half_angle]
        else:
            # For identity rotation, axis is arbitrary
            axis = [1, 0, 0]

        if return_matrix:
            axis = sp.Matrix(axis)

        return angle, axis

    @property
    def scalar(self):
        """Alias for w component"""
        return self.w

    @property
    def vector(self):
        """Vector part as list"""
        return [self.x, self.y, self.z]

    @property
    def vector_as_matrix(self):
        """Vector part as sp.Matrix(3,1)"""
        return sp.Matrix([self.x, self.y, self.z])

    @property
    def as_matrix(self):
        """Full quaternion as sp.Matrix(4,1)"""
        return sp.Matrix([self.w, self.x, self.y, self.z])

    @classmethod
    def from_rotation_matrix(cls, matrix):
        """
        Create quaternion from 3x3 rotation matrix
        matrix: sympy.Matrix of shape (3, 3)
        """
        m = matrix
        trace = m[0, 0] + m[1, 1] + m[2, 2]

        if trace > 0:
            s = 2 * sp.sqrt(trace + 1)
            w = s / 4
            x = (m[2, 1] - m[1, 2]) / s
            y = (m[0, 2] - m[2, 0]) / s
            z = (m[1, 0] - m[0, 1]) / s
        elif m[0, 0] > m[1, 1] and m[0, 0] > m[2, 2]:
            s = 2 * sp.sqrt(1 + m[0, 0] - m[1, 1] - m[2, 2])
            w = (m[2, 1] - m[1, 2]) / s
            x = s / 4
            y = (m[0, 1] + m[1, 0]) / s
            z = (m[0, 2] + m[2, 0]) / s
        elif m[1, 1] > m[2, 2]:
            s = 2 * sp.sqrt(1 + m[1, 1] - m[0, 0] - m[2, 2])
            w = (m[0, 2] - m[2, 0]) / s
            x = (m[0, 1] + m[1, 0]) / s
            y = s / 4
            z = (m[1, 2] + m[2, 1]) / s
        else:
            s = 2 * sp.sqrt(1 + m[2, 2] - m[0, 0] - m[1, 1])
            w = (m[1, 0] - m[0, 1]) / s
            x = (m[0, 2] + m[2, 0]) / s
            y = (m[1, 2] + m[2, 1]) / s
            z = s / 4

        return cls(w, x, y, z)
