from __future__ import annotations

import functools
from enum import IntEnum
from typing import overload, Union, Iterable, Tuple, Optional, Callable, List, Sequence, Dict, TypeVar, \
    TYPE_CHECKING
import numpy as np
import casadi as ca

from scipy import sparse as sp

if TYPE_CHECKING:
    from ..world_entity import Body

all_expressions = Union[Symbol_, Symbol, Expression, Point3, Vector3, RotationMatrix, TransformationMatrix, Quaternion]
all_expressions_float = Union[Symbol, Expression, Point3, Vector3, RotationMatrix, TransformationMatrix, float, Quaternion]
symbol_expr_float = Union[Symbol, Expression, float, int, IntEnum]
symbol_expr = Union[Symbol, Expression]
PreservedCasType = TypeVar('PreservedCasType', Point3, Vector3, TransformationMatrix, RotationMatrix, Quaternion, Expression)

pi: float

def _operation_type_error(arg1: object, operation: str, arg2: object) -> TypeError: ...

class ReferenceFrameMixin:
    reference_frame: Optional[Body]


class StackedCompiledFunction:
    compiled_f: CompiledFunction
    split_out_view: List[np.ndarray]
    symbol_parameters: List[Symbol]

    def __init__(self,
                 expressions: List[Expression],
                 parameters: Optional[Union[List[Symbol], List[List[Symbol]]]] = None,
                 additional_views: Optional[List[slice]] = None): ...

    def fast_call(self, *args: np.ndarray) -> np.ndarray: ...


class CompiledFunction:
    symbol_parameters: List[Symbol]
    compiled_casadi_function: ca.Function
    function_buffer: ca.FunctionBuffer
    function_evaluator: functools.partial
    out: Union[np.ndarray, sp.csc_matrix]
    sparse: bool

    def __init__(self,
                 expression: Symbol_,
                 parameters: Optional[Union[List[Symbol], List[List[Symbol]]]] = None,
                 sparse: bool = False): ...

    def __call__(self, **kwargs) -> np.ndarray: ...

    def fast_call(self, *args: np.ndarray) -> Union[np.ndarray, sp.csc_matrix]: ...


class Symbol_:
    s: Union[ca.SX, np.ndarray]
    np_data: Optional[np.ndarray]

    @property
    def shape(self) -> Tuple[int, int]: ...

    def __getitem__(self, item: Union[np.ndarray, Union[int, slice], Tuple[Union[int, slice], Union[int, slice]]]) -> Expression: ...

    def __setitem__(self, key: Union[Union[int, slice], Tuple[Union[int, slice], Union[int, slice]]], value: symbol_expr_float): ...

    def __len__(self) -> int: ...

    def free_symbols(self) -> List[ca.SX]: ...

    def is_constant(self) -> bool: ...

    def to_np(self) -> Union[float, np.ndarray]: ...

    def compile(self, parameters: Optional[Union[List[Symbol], List[List[Symbol]]]] = None, sparse: bool = False) \
            -> CompiledFunction: ...

    def __hash__(self) -> int: ...

    def pretty_str(self) -> str: ...


class Symbol(Symbol_):
    _registry: Dict[str, Symbol]
    name: str

    def __init__(self, name: str): ...

    @overload
    def __add__(self, other: Point3) -> Point3: ...
    @overload
    def __add__(self, other: Vector3) -> Vector3: ...
    @overload
    def __add__(self, other: Union[Symbol, Expression, float, Quaternion]) -> Expression: ...
    def __radd__(self, other: float) -> Expression: ...

    @overload
    def __sub__(self, other: Point3) -> Point3: ...
    @overload
    def __sub__(self, other: Vector3) -> Vector3: ...
    @overload
    def __sub__(self, other: RotationMatrix) -> RotationMatrix: ...
    @overload
    def __sub__(self, other: TransformationMatrix) -> TransformationMatrix: ...
    @overload
    def __sub__(self, other: Union[Symbol, Expression, float, Quaternion]) -> Expression: ...
    def __rsub__(self, other: float) -> Expression: ...

    @overload
    def __mul__(self, other: Point3) -> Point3: ...
    @overload
    def __mul__(self, other: Vector3) -> Vector3: ...
    @overload
    def __mul__(self, other: RotationMatrix) -> RotationMatrix: ...
    @overload
    def __mul__(self, other: TransformationMatrix) -> TransformationMatrix: ...
    @overload
    def __mul__(self, other: Union[Symbol, Expression, float, Quaternion]) -> Expression: ...
    def __rmul__(self, other: float) -> Expression: ...

    @overload
    def __truediv__(self, other: Point3) -> Point3: ...
    @overload
    def __truediv__(self, other: Vector3) -> Vector3: ...
    @overload
    def __truediv__(self, other: RotationMatrix) -> RotationMatrix: ...
    @overload
    def __truediv__(self, other: TransformationMatrix) -> TransformationMatrix: ...
    @overload
    def __truediv__(self, other: Union[Symbol, Expression, float, Quaternion]) -> Expression: ...
    def __rtruediv__(self, other: float) -> Expression: ...

    def __floordiv__(self, other: symbol_expr_float) -> Expression: ...
    def __rfloordiv__(self, other: symbol_expr_float) -> Expression: ...
    def __mod__(self, other: symbol_expr_float) -> Expression: ...
    def __rmod__(self, other: symbol_expr_float) -> Expression: ...
    def __divmod__(self, other: symbol_expr_float) -> Tuple[Expression, Expression]: ...
    def __rdivmod__(self, other: symbol_expr_float) -> Tuple[Expression, Expression]: ...

    def __lt__(self, other: all_expressions_float) -> Expression: ...
    def __le__(self, other: all_expressions_float) -> Expression: ...
    def __gt__(self, other: all_expressions_float) -> Expression: ...
    def __ge__(self, other: all_expressions_float) -> Expression: ...
    def __and__(self, other: all_expressions_float) -> Expression: ...
    def __or__(self, other: all_expressions_float) -> Expression: ...
    def __neg__(self) -> Expression: ...
    def __invert__(self) -> Expression: ...
    def __pow__(self, other: all_expressions_float) -> Expression: ...
    def __rpow__(self, other: all_expressions_float) -> Expression: ...
    def __eq__(self, other: object) -> bool: ...
    def __ne__(self, other: object) -> bool: ...


class Expression(Symbol_):
    def __init__(self, data: Optional[Union[Symbol,
                                            Expression,
                                            float,
                                            Vector3, Point3, TransformationMatrix, RotationMatrix, Quaternion,
                                            Iterable[symbol_expr_float],
                                            Iterable[Iterable[symbol_expr_float]],
                                            np.ndarray]] = None): ...

    @overload
    def __add__(self, other: Point3) -> Point3: ...
    @overload
    def __add__(self, other: Vector3) -> Vector3: ...
    @overload
    def __add__(self, other: Union[Symbol, Expression, float, TransformationMatrix, RotationMatrix, Quaternion]) -> Expression: ...
    def __radd__(self, other: float) -> Expression: ...

    @overload
    def __sub__(self, other: Point3) -> Point3: ...
    @overload
    def __sub__(self, other: Vector3) -> Vector3: ...
    @overload
    def __sub__(self, other: Union[Symbol, Expression, float, TransformationMatrix, RotationMatrix, Quaternion]) -> Expression: ...
    def __rsub__(self, other: float) -> Expression: ...

    @overload
    def __mul__(self, other: Point3) -> Point3: ...
    @overload
    def __mul__(self, other: Vector3) -> Vector3: ...
    @overload
    def __mul__(self, other: Union[Symbol, Expression, float, RotationMatrix, TransformationMatrix, Quaternion]) -> Expression: ...
    def __rmul__(self, other: float) -> Expression: ...

    @overload
    def __truediv__(self, other: Point3) -> Point3: ...
    @overload
    def __truediv__(self, other: Vector3) -> Vector3: ...
    @overload
    def __truediv__(self, other: Union[Symbol, Expression, float, RotationMatrix, TransformationMatrix, Quaternion]) -> Expression: ...
    def __rtruediv__(self, other: float) -> Expression: ...

    def __floordiv__(self, other: symbol_expr_float) -> Expression: ...
    def __rfloordiv__(self, other: symbol_expr_float) -> Expression: ...
    def __mod__(self, other: symbol_expr_float) -> Expression: ...
    def __rmod__(self, other: symbol_expr_float) -> Expression: ...
    def __divmod__(self, other: symbol_expr_float) -> Tuple[Expression, Expression]: ...
    def __rdivmod__(self, other: symbol_expr_float) -> Tuple[Expression, Expression]: ...

    def __eq__(self, other: symbol_expr_float) -> Expression: ...
    def __lt__(self, other: symbol_expr_float) -> Expression: ...
    def __le__(self, other: symbol_expr_float) -> Expression: ...
    def __gt__(self, other: symbol_expr_float) -> Expression: ...
    def __ge__(self, other: symbol_expr_float) -> Expression: ...
    def __and__(self, other: all_expressions_float) -> Expression: ...
    def __or__(self, other: all_expressions_float) -> Expression: ...
    def __neg__(self) -> Expression: ...
    def __invert__(self) -> Expression: ...
    def __pow__(self, other: symbol_expr_float) -> Expression: ...
    def __rpow__(self, other: symbol_expr_float) -> Expression: ...
    def __copy__(self) -> Expression: ...

    @property
    def T(self) -> Expression: ...

    def remove(self, rows: List[int], columns: List[int]): ...

    def split(self) -> List[Expression]: ...

    def dot(self, other: Expression) -> Expression: ...

    def reshape(self, new_shape: Tuple[int, int]) -> Expression: ...


class Point3(Symbol_, ReferenceFrameMixin):

    @property
    def x(self) -> Expression: ...
    @x.setter
    def x(self, value: symbol_expr_float): ...
    @property
    def y(self) -> Expression: ...
    @y.setter
    def y(self, value: symbol_expr_float): ...
    @property
    def z(self) -> Expression: ...
    @z.setter
    def z(self, value: symbol_expr_float): ...

    def __init__(self, x: symbol_expr_float = 0, y: symbol_expr_float = 0, z: symbol_expr_float = 0,
                 reference_frame: Optional[Body] = None): ...

    @classmethod
    def from_iterable(cls, data: Optional[Union[Expression, Point3, Vector3,
                                            ca.SX,
                                            np.ndarray,
                                            Iterable[symbol_expr_float]]] = None,
                 reference_frame: Optional[Body] = None) -> Point3: ...

    def norm(self) -> Expression: ...

    @overload
    def __add__(self, other: Union[Vector3, Symbol, Expression]) -> Point3: ...
    @overload
    def __add__(self, other: float) -> Point3: ...
    def __radd__(self, other: float) -> Point3: ...

    @overload
    def __sub__(self, other: Point3) -> Vector3: ...
    @overload
    def __sub__(self, other: Union[Symbol, Expression, Vector3]) -> Point3: ...
    @overload
    def __sub__(self, other: float) -> Point3: ...
    def __rsub__(self, other: float) -> Point3: ...

    @overload
    def __mul__(self, other: Symbol) -> Point3: ...
    @overload
    def __mul__(self, other: Expression) -> Point3: ...
    @overload
    def __mul__(self, other: float) -> Point3: ...
    def __rmul__(self, other: float) -> Point3: ...

    @overload
    def __matmul__(self, other: Point3) -> Expression: ...
    @overload
    def __matmul__(self, other: Vector3) -> Expression: ...

    @overload
    def __pow__(self, other: Symbol) -> Point3: ...
    @overload
    def __pow__(self, other: Expression) -> Point3: ...
    @overload
    def __pow__(self, other: float) -> Point3: ...
    def __rpow__(self, other: float) -> Point3: ...

    @overload
    def __truediv__(self, other: Symbol) -> Point3: ...
    @overload
    def __truediv__(self, other: Expression) -> Point3: ...
    @overload
    def __truediv__(self, other: float) -> Point3: ...
    def __rtruediv__(self, other: float) -> Point3: ...

    def __neg__(self) -> Point3: ...

    @overload
    def dot(self, other: Point3) -> Expression: ...
    @overload
    def dot(self, other: Vector3) -> Expression: ...


class Vector3(Symbol_, ReferenceFrameMixin):
    vis_frame: Optional[Body]

    @property
    def x(self) -> Expression: ...
    @x.setter
    def x(self, value: symbol_expr_float): ...
    @property
    def y(self) -> Expression: ...
    @y.setter
    def y(self, value: symbol_expr_float): ...
    @property
    def z(self) -> Expression: ...
    @z.setter
    def z(self, value: symbol_expr_float): ...

    def __init__(self, x: symbol_expr_float = 0, y: symbol_expr_float = 0, z: symbol_expr_float = 0,
                 reference_frame: Optional[Body] = None): ...

    @classmethod
    def from_iterable(cls, data: Optional[Union[Expression, Point3, Vector3,
                                            ca.SX,
                                            np.ndarray,
                                            Iterable[symbol_expr_float]]] = None,
                 reference_frame: Optional[Body] = None) -> Vector3: ...

    @classmethod
    def X(cls, reference_frame: Optional[Body] = None) -> Vector3: ...
    @classmethod
    def Y(cls, reference_frame: Optional[Body] = None) -> Vector3: ...
    @classmethod
    def Z(cls, reference_frame: Optional[Body] = None) -> Vector3: ...

    @classmethod
    def unit_vector(cls, x: symbol_expr_float = 0, y: symbol_expr_float = 0, z: symbol_expr_float = 0,
                    reference_frame: Optional[Body] = None) -> Vector3: ...

    def norm(self) -> Expression: ...

    def scale(self, a: symbol_expr_float, unsafe: bool = False): ...

    @overload
    def __add__(self, other: Point3) -> Point3: ...
    @overload
    def __add__(self, other: Vector3) -> Vector3: ...
    @overload
    def __add__(self, other: Symbol) -> Vector3: ...
    @overload
    def __add__(self, other: Expression) -> Vector3: ...
    @overload
    def __add__(self, other: float) -> Vector3: ...
    def __radd__(self, other: float) -> Vector3: ...

    @overload
    def __sub__(self, other: Vector3) -> Vector3: ...
    @overload
    def __sub__(self, other: Point3) -> Point3: ...
    @overload
    def __sub__(self, other: Symbol) -> Vector3: ...
    @overload
    def __sub__(self, other: Expression) -> Vector3: ...
    @overload
    def __sub__(self, other: float) -> Vector3: ...
    def __rsub__(self, other: float) -> Vector3: ...

    @overload
    def __mul__(self, other: Symbol) -> Vector3: ...
    @overload
    def __mul__(self, other: Expression) -> Vector3: ...
    @overload
    def __mul__(self, other: float) -> Vector3: ...
    def __rmul__(self, other: float) -> Vector3: ...

    @overload
    def __matmul__(self, other: Point3) -> Expression: ...
    @overload
    def __matmul__(self, other: Vector3) -> Expression: ...

    @overload
    def __truediv__(self, other: Symbol) -> Vector3: ...
    @overload
    def __truediv__(self, other: Expression) -> Vector3: ...
    @overload
    def __truediv__(self, other: float) -> Vector3: ...
    def __rtruediv__(self, other: float) -> Vector3: ...

    def __neg__(self) -> Vector3: ...

    @overload
    def __pow__(self, other: Symbol) -> Vector3: ...
    @overload
    def __pow__(self, other: Expression) -> Vector3: ...
    @overload
    def __pow__(self, other: float) -> Vector3: ...
    def __rpow__(self, other: float) -> Vector3: ...

    @overload
    def dot(self, other: Point3) -> Expression: ...
    @overload
    def dot(self, other: Vector3) -> Expression: ...

    def cross(self, other: Vector3) -> Vector3: ...


TrinaryFalse = 0
TrinaryUnknown = 0.5
TrinaryTrue = 1

BinaryTrue: Expression
BinaryFalse: Expression



class TransformationMatrix(Symbol_, ReferenceFrameMixin):
    child_frame: Optional[Body]

    def __init__(self, data: Optional[Union[TransformationMatrix,
                                            RotationMatrix,
                                            ca.SX,
                                            np.ndarray,
                                            Expression,
                                            Iterable[Iterable[symbol_expr_float]]]] = None,
                 reference_frame: Optional[Body] = None,
                 child_frame: Optional[Body] = None,
                 sanity_check: bool = True): ...

    @property
    def x(self) -> Expression: ...
    @x.setter
    def x(self, value: symbol_expr_float): ...
    @property
    def y(self) -> Expression: ...
    @y.setter
    def y(self, value: symbol_expr_float): ...
    @property
    def z(self) -> Expression: ...
    @z.setter
    def z(self, value: symbol_expr_float): ...

    @overload
    def dot(self, other: Point3) -> Point3: ...
    @overload
    def dot(self, other: Vector3) -> Vector3: ...
    @overload
    def dot(self, other: RotationMatrix) -> RotationMatrix: ...
    @overload
    def dot(self, other: TransformationMatrix) -> TransformationMatrix: ...

    @overload
    def __matmul__(self, other: Point3) -> Point3: ...
    @overload
    def __matmul__(self, other: Vector3) -> Vector3: ...
    @overload
    def __matmul__(self, other: RotationMatrix) -> RotationMatrix: ...
    @overload
    def __matmul__(self, other: TransformationMatrix) -> TransformationMatrix: ...

    def to_rotation(self) -> RotationMatrix: ...
    def to_quaternion(self) -> Quaternion: ...
    def to_position(self) -> Point3: ...
    def to_translation(self) -> TransformationMatrix: ...

    @classmethod
    def from_xyz_rpy(cls,
                     x: Optional[symbol_expr_float] = 0,
                     y: Optional[symbol_expr_float] = 0,
                     z: Optional[symbol_expr_float] = 0,
                     roll: Optional[symbol_expr_float] = 0,
                     pitch: Optional[symbol_expr_float] = 0,
                     yaw: Optional[symbol_expr_float] = 0,
                     reference_frame: Optional[Body] = None,
                     child_frame: Optional[Body] = None) -> TransformationMatrix: ...
    @classmethod
    def from_xyz_quat(cls,
                      pos_x: symbol_expr_float = 0, pos_y: symbol_expr_float = 0, pos_z: symbol_expr_float =00,
                      quat_w: symbol_expr_float = 0, quat_x: symbol_expr_float = 0,
                      quat_y: symbol_expr_float = 0, quat_z: symbol_expr_float = 1,
                      reference_frame: Optional[Body] = None, child_frame: Optional[Body] = None) \
            -> TransformationMatrix: ...
    @classmethod
    def from_point_rotation_matrix(cls,
                                   point: Optional[Point3] = None,
                                   rotation_matrix: Optional[RotationMatrix] = None,
                                   reference_frame: Optional[Body] = None,
                                   child_frame: Optional[Body] = None) -> TransformationMatrix: ...

    def inverse(self) -> TransformationMatrix: ...


class RotationMatrix(Symbol_, ReferenceFrameMixin):
    child_frame: Optional[Body]

    def __init__(self, data: Optional[Union[TransformationMatrix,
                                            RotationMatrix,
                                            Expression,
                                            Quaternion,
                                            ca.SX,
                                            np.ndarray,
                                            Iterable[Iterable[symbol_expr_float]]]] = None,
                 reference_frame: Optional[Body] = None,
                 child_frame: Optional[Body] = None,
                 sanity_check: bool = True): ...
    @classmethod
    def __quaternion_to_rotation_matrix(cls, q: Quaternion) -> RotationMatrix: ...

    @classmethod
    def from_axis_angle(cls, axis: Vector3, angle: symbol_expr_float, reference_frame: Optional[Body] = None) \
            -> RotationMatrix: ...
    @classmethod
    def from_quaternion(cls, q: Quaternion) -> RotationMatrix: ...
    @classmethod
    @overload
    def from_vectors(cls, x: Vector3, y: Vector3, reference_frame: Optional[Body] = None) -> RotationMatrix: ...
    @classmethod
    @overload
    def from_vectors(cls, x: Vector3, z: Vector3, reference_frame: Optional[Body] = None) -> RotationMatrix: ...
    @classmethod
    @overload
    def from_vectors(cls, y: Vector3, z: Vector3, reference_frame: Optional[Body] = None) -> RotationMatrix: ...
    @classmethod
    @overload
    def from_vectors(cls, x: Vector3, y: Vector3, z: Vector3,
                     reference_frame: Optional[Body] = None) -> RotationMatrix: ...

    @classmethod
    def from_rpy(cls,
                 roll: Optional[symbol_expr_float] = None,
                 pitch: Optional[symbol_expr_float] = None,
                 yaw: Optional[symbol_expr_float] = None,
                 reference_frame: Optional[Body] = None) -> RotationMatrix: ...

    def x_vector(self) -> Vector3: ...
    def y_vector(self) -> Vector3: ...
    def z_vector(self) -> Vector3: ...

    def to_axis_angle(self) -> Tuple[Vector3, Expression]: ...
    def to_angle(self, hint: Optional[Callable] = None) -> Expression: ...
    def to_rpy(self) -> Tuple[Expression, Expression, Expression]: ...
    def to_quaternion(self) -> Quaternion: ...

    @overload
    def dot(self, other: Vector3) -> Vector3: ...
    @overload
    def dot(self, other: RotationMatrix) -> RotationMatrix: ...
    @overload
    def dot(self, other: TransformationMatrix) -> TransformationMatrix: ...

    @overload
    def __matmul__(self, other: Vector3) -> Vector3: ...
    @overload
    def __matmul__(self, other: RotationMatrix) -> RotationMatrix: ...
    @overload
    def __matmul__(self, other: TransformationMatrix) -> TransformationMatrix: ...

    def normalize(self): ...
    def inverse(self) -> RotationMatrix: ...
    @property
    def T(self) -> RotationMatrix: ...


class Quaternion(Symbol_, ReferenceFrameMixin):
    @property
    def x(self) -> Expression: ...
    @x.setter
    def x(self, value: symbol_expr_float): ...
    @property
    def y(self) -> Expression: ...
    @y.setter
    def y(self, value: symbol_expr_float): ...
    @property
    def z(self) -> Expression: ...
    @z.setter
    def z(self, value: symbol_expr_float): ...
    @property
    def w(self) -> Expression: ...
    @w.setter
    def w(self, value: symbol_expr_float): ...

    def __init__(self, x: symbol_expr_float = 0.0, y: symbol_expr_float = 0.0,
                 z: symbol_expr_float = 0.0, w: symbol_expr_float = 1.0,
                 reference_frame: Optional[Body] = None): ...

    @classmethod
    def from_iterable(cls, data: Optional[Union[Expression,
                                            Quaternion,
                                            ca.SX,
                                            Tuple[symbol_expr_float,
                                                  symbol_expr_float,
                                                  symbol_expr_float,
                                                  symbol_expr_float]]] = None,
                 reference_frame: Optional[Body] = None) -> Quaternion: ...

    @classmethod
    def from_axis_angle(cls, axis: Vector3, angle: symbol_expr_float, reference_frame: Optional[Body] = None) \
            -> Quaternion: ...
    @classmethod
    def from_rpy(cls, roll: symbol_expr_float, pitch: symbol_expr_float, yaw: symbol_expr_float,
                 reference_frame: Optional[Body] = None) -> Quaternion: ...
    @classmethod
    def from_rotation_matrix(cls, r: Union[RotationMatrix, TransformationMatrix]) -> Quaternion: ...

    def conjugate(self) -> Quaternion: ...
    def multiply(self, q: Quaternion) -> Quaternion: ...
    def diff(self, q: Quaternion) -> Quaternion: ...
    def norm(self) -> Expression: ...
    def normalize(self): ...
    def to_axis_angle(self) -> Tuple[Vector3, Expression]: ...
    def to_rotation_matrix(self) -> RotationMatrix: ...
    def to_rpy(self) -> Tuple[Expression, Expression, Expression]: ...
    def dot(self, other: Quaternion) -> Expression: ...

    def __neg__(self) -> Quaternion: ...


def norm(v: Union[Vector3, Point3, Expression, Quaternion]) -> Expression: ...

@overload
def save_division(nominator: Vector3, denominator: symbol_expr_float, if_nan: Optional[Vector3] = None) -> Vector3: ...
@overload
def save_division(nominator: Point3, denominator: symbol_expr_float, if_nan: Optional[Point3] = None) -> Point3: ...
@overload
def save_division(nominator: symbol_expr_float, denominator: symbol_expr_float,
                  if_nan: Optional[symbol_expr_float] = None) -> Expression: ...

def diag(args: Union[List[symbol_expr_float], Expression]) -> Expression: ...

def hessian(expressions: Union[symbol_expr, List[symbol_expr]], symbols: Iterable[Symbol]) -> Expression: ...

def jacobian(expressions: Union[symbol_expr, List[symbol_expr]], symbols: Iterable[Symbol]) -> Expression: ...

def jacobian_dot(expressions: Expression, symbols: List[symbol_expr], symbols_dot: List[symbol_expr]) -> Expression: ...

def jacobian_ddot(expressions: Expression, symbols: List[symbol_expr], symbols_dot: List[symbol_expr],
                  symbols_ddot: List[symbol_expr]) -> Expression: ...

def equivalent(expression1: symbol_expr, expression2: symbol_expr) -> bool: ...

def var(variables_names: str) -> List[Symbol]: ...

def free_symbols(expression: all_expressions) -> List[ca.SX]: ...

def create_symbols(names: Union[List[str], int]) -> List[Symbol]: ...

def compile_and_execute(f: Callable[..., all_expressions],
                        params: Union[List[Union[float, np.ndarray]], np.ndarray]) \
        -> Union[float, np.ndarray]: ...

def zeros(rows: int, columns: int) -> Expression: ...

def ones(x: int, y: int) -> Expression: ...

def tri(dimension: int) -> Expression: ...

def eye(size: int) -> Expression: ...

@overload
def abs(x: Vector3) -> Vector3: ...
@overload
def abs(x: Point3) -> Point3: ...
@overload
def abs(x: symbol_expr_float) -> Expression: ...

@overload
def max(x: Point3) -> Expression: ...
@overload
def max(x: Vector3) -> Expression: ...
@overload
def max(x: symbol_expr_float, y: Optional[symbol_expr_float] = None) -> Expression: ...

@overload
def min(x: Point3) -> Expression: ...
@overload
def min(x: Vector3) -> Expression: ...
@overload
def min(x: symbol_expr_float, y: Optional[symbol_expr_float] = None) -> Expression: ...

def limit(x: symbol_expr_float,
          lower_limit: symbol_expr_float,
          upper_limit: symbol_expr_float) -> Expression: ...

def equal(x: symbol_expr_float, y: symbol_expr_float) -> Expression: ...

def not_equal(x: symbol_expr_float, y: symbol_expr_float) -> Expression: ...

def less_equal(x: symbol_expr_float, y: symbol_expr_float) -> Expression: ...
def greater_equal(x: symbol_expr_float, y: symbol_expr_float) -> Expression: ...
def less(x: symbol_expr_float, y: symbol_expr_float) -> Expression: ...
def greater(x: symbol_expr_float, y: symbol_expr_float) -> Expression: ...

def logic_and(*args: symbol_expr_float) -> symbol_expr_float: ...

def logic_and3(*args: symbol_expr_float) -> symbol_expr_float: ...

def logic_or(*args: symbol_expr_float, simplify: bool = True) -> symbol_expr_float: ...

def logic_or3(*args: symbol_expr_float) -> symbol_expr_float: ...

def logic_any(args: Expression) -> symbol_expr_float: ...

def logic_all(args: Expression) -> symbol_expr_float: ...

def logic_not(expr: symbol_expr_float) -> Expression: ...

def logic_not3(expr: symbol_expr_float) -> Expression: ...

@overload
def if_else(condition: symbol_expr_float, if_result: Vector3, else_result: Vector3) -> Vector3: ...
@overload
def if_else(condition: symbol_expr_float, if_result: Point3, else_result: Point3) -> Point3: ...
@overload
def if_else(condition: symbol_expr_float, if_result: symbol_expr_float,
            else_result: symbol_expr_float) -> Expression: ...

@overload
def if_greater(a: symbol_expr_float, b: symbol_expr_float, if_result: Vector3,
               else_result: Vector3) -> Vector3: ...
@overload
def if_greater(a: symbol_expr_float, b: symbol_expr_float, if_result: Point3,
               else_result: Point3) -> Point3: ...
@overload
def if_greater(a: symbol_expr_float, b: symbol_expr_float, if_result: symbol_expr_float,
               else_result: symbol_expr_float) -> Expression: ...

@overload
def if_less(a: symbol_expr_float, b: symbol_expr_float, if_result: Quaternion,
            else_result: Quaternion) -> Quaternion: ...
@overload
def if_less(a: symbol_expr_float, b: symbol_expr_float, if_result: Vector3,
            else_result: Vector3) -> Vector3: ...
@overload
def if_less(a: symbol_expr_float, b: symbol_expr_float, if_result: Point3,
            else_result: Point3) -> Point3: ...
@overload
def if_less(a: symbol_expr_float, b: symbol_expr_float, if_result: symbol_expr_float,
            else_result: symbol_expr_float) -> Expression: ...

@overload
def if_greater_zero(condition: symbol_expr_float,
                    if_result: Vector3,
                    else_result: Vector3) -> Vector3: ...
@overload
def if_greater_zero(condition: symbol_expr_float,
                    if_result: Point3,
                    else_result: Point3) -> Point3: ...
@overload
def if_greater_zero(condition: symbol_expr_float,
                    if_result: Quaternion,
                    else_result: Quaternion) -> Quaternion: ...
@overload
def if_greater_zero(condition: symbol_expr_float,
                    if_result: symbol_expr_float,
                    else_result: symbol_expr_float) -> Expression: ...

@overload
def if_greater_eq_zero(condition: symbol_expr_float, if_result: Vector3,
                       else_result: Vector3) -> Vector3: ...
@overload
def if_greater_eq_zero(condition: symbol_expr_float, if_result: Point3,
                       else_result: Point3) -> Point3: ...
@overload
def if_greater_eq_zero(condition: symbol_expr_float, if_result: symbol_expr_float,
                       else_result: symbol_expr_float) -> Expression: ...

@overload
def if_greater_eq(a: symbol_expr_float, b: symbol_expr_float, if_result: Vector3,
                  else_result: Vector3) -> Vector3: ...
@overload
def if_greater_eq(a: symbol_expr_float, b: symbol_expr_float, if_result: Point3,
                  else_result: Point3) -> Point3: ...
@overload
def if_greater_eq(a: symbol_expr_float, b: symbol_expr_float, if_result: symbol_expr_float,
                  else_result: symbol_expr_float) -> Expression: ...

@overload
def if_less_eq(a: symbol_expr_float, b: symbol_expr_float, if_result: Vector3,
               else_result: Vector3) -> Vector3: ...
@overload
def if_less_eq(a: symbol_expr_float, b: symbol_expr_float, if_result: Point3,
               else_result: Point3) -> Point3: ...
@overload
def if_less_eq(a: symbol_expr_float, b: symbol_expr_float, if_result: symbol_expr_float,
               else_result: symbol_expr_float) -> Expression: ...

@overload
def if_eq_zero(condition: symbol_expr_float,
               if_result: Vector3,
               else_result: Vector3) -> Vector3: ...
@overload
def if_eq_zero(condition: symbol_expr_float,
               if_result: Point3,
               else_result: Point3) -> Point3: ...
@overload
def if_eq_zero(condition: symbol_expr_float,
               if_result: symbol_expr_float,
               else_result: symbol_expr_float) -> Expression: ...

@overload
def if_eq(a: symbol_expr_float,
          b: symbol_expr_float,
          if_result: Vector3,
          else_result: Vector3) -> Vector3: ...
@overload
def if_eq(a: symbol_expr_float,
          b: symbol_expr_float,
          if_result: Point3,
          else_result: Point3) -> Point3: ...
@overload
def if_eq(a: symbol_expr_float,
          b: symbol_expr_float,
          if_result: symbol_expr_float,
          else_result: symbol_expr_float) -> Expression: ...

def if_eq_cases(a: symbol_expr_float,
                b_result_cases: Iterable[Tuple[symbol_expr_float, symbol_expr_float]],
                else_result: symbol_expr_float) -> Expression: ...

def if_eq_cases_grouped(a: symbol_expr_float,
                        b_result_cases: Iterable[Tuple[symbol_expr_float, symbol_expr_float]],
                        else_result: symbol_expr_float) -> Expression: ...

def if_cases(cases: Sequence[Tuple[symbol_expr_float, symbol_expr_float]],
             else_result: symbol_expr_float) -> Expression: ...

def if_less_eq_cases(a: symbol_expr_float,
                     b_result_cases: Sequence[Tuple[symbol_expr_float, symbol_expr_float]],
                     else_result: symbol_expr_float) -> Expression: ...

def cross(u: Union[Vector3, Expression], v: Union[Vector3, Expression]) -> Vector3: ...

def _to_sx(thing: object) -> ca.SX: ...

@overload
def scale(v: Vector3, a: symbol_expr_float) -> Vector3: ...
@overload
def scale(v: Point3, a: symbol_expr_float) -> Point3: ...
@overload
def scale(v: Expression, a: symbol_expr_float) -> Expression: ...

@overload
def dot(e1: TransformationMatrix, e2: Point3) -> Point3: ...
@overload
def dot(e1: TransformationMatrix, e2: Vector3) -> Vector3: ...
@overload
def dot(e1: RotationMatrix, e2: Point3) -> Point3: ...
@overload
def dot(e1: RotationMatrix, e2: Vector3) -> Vector3: ...
@overload
def dot(e1: RotationMatrix, e2: RotationMatrix) -> RotationMatrix: ...
@overload
def dot(e1: RotationMatrix, e2: TransformationMatrix) -> TransformationMatrix: ...
@overload
def dot(e1: TransformationMatrix, e2: RotationMatrix) -> TransformationMatrix: ...
@overload
def dot(e1: TransformationMatrix, e2: TransformationMatrix) -> TransformationMatrix: ...
@overload
def dot(e1: Quaternion, e2: Quaternion) -> Expression: ...
@overload
def dot(e1: Union[Vector3, Point3], e2: Union[Vector3, Point3]) -> Expression: ...
@overload
def dot(e1: Expression, e2: Expression) -> Expression: ...

def kron(m1: Expression, m2: Expression) -> Expression: ...

def trace(matrix: Union[Expression, RotationMatrix, TransformationMatrix]) -> Expression: ...

# def rotation_distance(a_R_b: Expression, a_R_c: Expression) -> Expression: ...

@overload
def vstack(list_of_matrices: List[Union[Point3, Vector3, Quaternion]]) -> Expression: ...
@overload
def vstack(list_of_matrices: List[TransformationMatrix]) -> Expression: ...
@overload
def vstack(list_of_matrices: List[Expression]) -> Expression: ...

@overload
def hstack(list_of_matrices: List[TransformationMatrix]) -> Expression: ...
@overload
def hstack(list_of_matrices: List[Expression]) -> Expression: ...

@overload
def diag_stack(list_of_matrices: List[TransformationMatrix]) -> Expression: ...
@overload
def diag_stack(list_of_matrices: List[Expression]) -> Expression: ...

def normalize_axis_angle(axis: Vector3, angle: symbol_expr_float) -> Tuple[Vector3, Expression]: ...

def axis_angle_from_rpy(roll: symbol_expr_float, pitch: symbol_expr_float, yaw: symbol_expr_float) \
        -> Tuple[Vector3, Expression]: ...

def cosine_distance(v0: symbol_expr_float, v1: symbol_expr_float) -> Expression: ...

@overload
def euclidean_distance(v1: symbol_expr_float, v2: symbol_expr_float) -> Expression: ...
@overload
def euclidean_distance(v1: Point3, v2: Point3) -> Expression: ...

def fmod(a: symbol_expr_float, b: symbol_expr_float) -> Expression: ...

def normalize_angle_positive(angle: symbol_expr_float) -> Expression: ...

def normalize_angle(angle: symbol_expr_float) -> Expression: ...

def shortest_angular_distance(from_angle: symbol_expr_float, to_angle: symbol_expr_float) -> Expression: ...

def quaternion_slerp(q1: Quaternion, q2: Quaternion, t: symbol_expr_float) -> Quaternion: ...

@overload
def slerp(v1: Vector3, v2: Vector3, t: symbol_expr_float) -> Vector3: ...
@overload
def slerp(v1: Expression, v2: Expression, t: symbol_expr_float) -> Expression: ...

def save_acos(angle: symbol_expr_float) -> Expression: ...

def entrywise_product(matrix1: Expression, matrix2: Expression) -> Expression: ...

def floor(x: symbol_expr_float) -> Expression: ...

def ceil(x: symbol_expr_float) -> Expression: ...

def round_up(x: symbol_expr_float, decimal_places: symbol_expr_float) -> Expression: ...

def round_down(x: symbol_expr_float, decimal_places: symbol_expr_float) -> Expression: ...

def sum(matrix: Expression) -> Expression: ...

def sum_row(matrix: Expression) -> Expression: ...

def sum_column(matrix: Expression) -> Expression: ...

def distance_point_to_line_segment(frame_P_current: Point3, frame_P_line_start: Point3, frame_P_line_end: Point3) \
        -> Tuple[Expression, Point3]: ...

def distance_point_to_line(frame_P_point: Point3, frame_P_line_point: Point3, frame_V_line_direction: Vector3) \
    -> Expression: ...

def distance_point_to_plane(frame_P_current: Point3, frame_V_v1: Vector3,
                                          frame_V_v2: Vector3) -> \
        Tuple[Expression, Point3]: ...

def distance_point_to_plane_signed(frame_P_current: Point3, frame_V_v1: Vector3,
                                          frame_V_v2: Vector3) -> \
        Tuple[Expression, Point3]: ...

def project_to_cone(frame_V_current: Vector3, frame_V_cone_axis: Vector3, cone_theta: Union[Symbol, float, Expression]) -> Vector3: ...

def project_to_plane(frame_V_plane_vector1: Vector3, frame_V_plane_vector2: Vector3, frame_P_point: Point3) -> Point3: ...

def angle_between_vector(v1: Vector3, v2: Vector3) -> Expression: ...

def velocity_limit_from_position_limit(acceleration_limit: Union[Symbol, float],
                                       position_limit: Union[Symbol, float],
                                       current_position: Union[Symbol, float],
                                       step_size: Union[Symbol, float],
                                       eps: float = 1e-5) -> Expression: ...

def rotational_error(r1: RotationMatrix, r2: RotationMatrix) -> Expression: ...

def to_str(expression: all_expressions) -> List[List[str]]: ...

def total_derivative(expr: Union[Symbol, Expression],
                     symbols: Iterable[Symbol],
                     symbols_dot: Iterable[Symbol]) \
        -> Expression: ...

def total_derivative2(expr: Union[Symbol, Expression],
                      symbols: Iterable[Symbol],
                      symbols_dot: Iterable[Symbol],
                      symbols_ddot: Iterable[Symbol]) -> Expression: ...

def quaternion_multiply(q1: Quaternion, q2: Quaternion) -> Quaternion: ...

def quaternion_conjugate(q: Quaternion) -> Quaternion: ...

def quaternion_diff(q1: Quaternion, q2: Quaternion) -> Quaternion: ...

def sign(x: symbol_expr_float) -> Expression: ...

def cos(x: symbol_expr_float) -> Expression: ...

def sin(x: symbol_expr_float) -> Expression: ...

def exp(x: symbol_expr_float) -> Expression: ...

def log(x: symbol_expr_float) -> Expression: ...

def tan(x: symbol_expr_float) -> Expression: ...

def sinh(x: symbol_expr_float) -> Expression: ...

def cosh(x: symbol_expr_float) -> Expression: ...

def sqrt(x: symbol_expr_float) -> Expression: ...

def acos(x: symbol_expr_float) -> Expression: ...

def atan2(x: symbol_expr_float, y: symbol_expr_float) -> Expression: ...

def solve_for(expression: Expression, target_value: float, start_value: float = 0.0001, max_tries: int = 10000,
              eps: float = 1e-10, max_step: float = 1) -> float: ...

def one_step_change(current_acceleration: symbol_expr_float, jerk_limit: symbol_expr_float, dt: float) \
        -> Expression: ...

def gauss(n: symbol_expr_float) -> Expression: ...


def r_gauss(integral: symbol_expr_float) -> Expression: ...


def substitute(expression: Union[Symbol, Expression], old_symbols: List[Symbol], new_symbols: List[Union[Symbol, Expression]]) -> Expression:
    ...

def matrix_inverse(a: Expression) -> Expression: ...

def gradient(ex: Expression, arg:Expression) -> Expression: ...

def is_true_symbol(expr: Expression) -> bool: ...

def is_true3(expr: Union[Symbol, Expression]) -> Expression: ...
def is_true3_symbol(expr: Expression) -> Expression: ...

def is_false3(expr: Union[Symbol, Expression]) -> Expression: ...
def is_false3_symbol(expr: Expression) -> Expression: ...

def is_unknown3(expr: Union[Symbol, Expression]) -> Expression: ...
def is_unknown3_symbol(expr: Expression) -> Expression: ...

def is_false_symbol(expr: Expression) -> bool: ...

def is_constant(expr: Expression) -> bool: ...

def det(expr: Expression) -> Expression: ...

def distance_projected_on_vector(point1: Point3, point2: Point3, vector: Vector3) -> Expression: ...

def distance_vector_projected_on_plane(point1: Point3, point2:Point3, normal_vector: Vector3) -> Vector3: ...

def replace_with_three_logic(expr: Expression) -> Expression: ...

def is_inf(expr: Expression) -> bool: ...

SpatialType = TypeVar('SpatialType', Point3, TransformationMatrix, Vector3, Quaternion, RotationMatrix)