content: # Copyright (c) Meta Platforms, Inc. and affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  from enum import Enum
  from typing import TYPE_CHECKING, Any, Callable, List, Optional, Protocol, Tuple, Union
  
  import torch
  
  from .functional import SE3 as _se3_impl
  from .functional import SO3 as _so3_impl
  from .functional.constants import DeviceType
  from .functional.lie_group import LieGroupFns
  
  if TYPE_CHECKING:
      from .lie_tensor import LieTensor, _LieTensorBase
  
  
  # The next two are similar to the ones in functional, except they return a LieTensor.
  class _RandFnType(Protocol):
      def __call__(
          self,
          *size: Any,
          generator: Optional[torch.Generator] = None,
          dtype: Optional[torch.dtype] = None,
          device: DeviceType = None,
          requires_grad: bool = False,
      ) -> "LieTensor":
          pass
  
  
  class _IdentityFnType(Protocol):
      def __call__(
          self,
          *size: int,
          dtype: Optional[torch.dtype] = None,
          device: DeviceType = None,
          requires_grad: bool = False,
      ) -> "LieTensor":
          pass
  
  
  Device = Union[torch.device, str, None]
  TensorType = Union[torch.Tensor, "_LieTensorBase"]
  _JFnReturnType = Tuple[List[torch.Tensor], TensorType]
  
  
  class _ltype(Enum):
      SE3 = 0
      SO3 = 1
  
  
  def _get_fn_lib(ltype: _ltype):
      return {
          _ltype.SE3: _se3_impl,
          _ltype.SO3: _so3_impl,
      }[ltype]
  
  
  def _eval_op(
      fn_lib: LieGroupFns,
      op_name: str,
      input0: torch.Tensor,
      jacobians: Optional[List[torch.Tensor]] = None,
  ) -> torch.Tensor:
      return getattr(fn_lib, op_name)(input0, jacobians=jacobians)
  
  
  class ltype:
      def __init__(self, lt: _ltype):
          self._lt = lt
          self._fn_lib = _get_fn_lib(lt)
  
      rand: _RandFnType
      randn: _RandFnType
      identity: _IdentityFnType
      _call_impl: Callable[[torch.Tensor], "LieTensor"]
  
      def __call__(self, tensor: torch.Tensor) -> "LieTensor":
          return self._call_impl(tensor)
  
      _create_lie_tensor: Callable[[torch.Tensor, "ltype"], "LieTensor"]
  
      def exp(self, tangent_vector: torch.Tensor) -> "LieTensor":
          return self._create_lie_tensor(
              _eval_op(self._fn_lib, "exp", tangent_vector), self
          )
  
      def jexp(self, tangent_vector: torch.Tensor) -> _JFnReturnType:
          jacs: List[torch.Tensor] = []
          op_res: TensorType = self._fn_lib.exp(tangent_vector, jacobians=jacs)
          return jacs, self._create_lie_tensor(op_res, self)
  
      def hat(self, tangent_vector: torch.Tensor) -> torch.Tensor:
          return _eval_op(self._fn_lib, "hat", tangent_vector)
  
      def vee(self, matrix: torch.Tensor) -> torch.Tensor:
          return _eval_op(self._fn_lib, "vee", matrix)
  
      def lift(self, matrix: torch.Tensor) -> torch.Tensor:
          return _eval_op(self._fn_lib, "lift", matrix)
  
      def project(self, matrix: torch.Tensor) -> torch.Tensor:
          return _eval_op(self._fn_lib, "project", matrix)
  
      def __str__(self) -> str:
          return self._lt.name
  
  
  SE3 = ltype(_ltype.SE3)
  SO3 = ltype(_ltype.SO3)
  