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.
  import cProfile
  import io
  import pstats
  import time
  from collections import defaultdict
  from typing import Any, Callable, Dict, List, Optional, Type
  
  import numpy as np
  import torch
  import torch.nn as nn
  
  import theseus as th
  
  
  def build_mlp(
      input_dim: int,
      hidden_dim: int,
      output_dim: int,
      hidden_depth: int,
      act_name: str = "relu",
  ) -> torch.nn.Sequential:
      # assign types
      act: Type[nn.Module]
      mods: List[nn.Module]
  
      # find activation
      if act_name == "relu":
          act = nn.ReLU
      elif act_name == "elu":
          act = nn.ELU
      elif act_name == "tanh":
          act = nn.Tanh
      else:
          raise NotImplementedError()
  
      # construct sequential list of modules
      if hidden_depth == 0:
          mods = [nn.Linear(input_dim, output_dim)]
      else:
          mods = [nn.Linear(input_dim, hidden_dim), act()]
          for i in range(hidden_depth - 1):
              mods += [nn.Linear(hidden_dim, hidden_dim), act()]
          mods.append(nn.Linear(hidden_dim, output_dim))
      trunk = nn.Sequential(*mods)
      return trunk
  
  
  # Inputs are a batched matrix and batched rows/cols for indexing
  # Sizes are:
  #      matrix -> batch_size x num_rows x num_cols
  #      rows/cols -> batch_size x num_points
  #
  # If the matrix's batch size is 1, broadcasting is supported by expanding the matrix
  # so that it has the batch size of rows and cols.
  #
  # Returns:
  #      data -> batch_size x num_points
  # where
  #      data[i, j] = matrix[i, rows[i, j], cols[i, j]]
  def gather_from_rows_cols(matrix: torch.Tensor, rows: torch.Tensor, cols: torch.Tensor):
      if matrix.shape[0] == 1:
          matrix = matrix.expand(rows.shape[0], -1, -1)
      assert matrix.ndim == 3 and rows.ndim == 2 and rows.ndim == 2
      assert matrix.shape[0] == rows.shape[0] and matrix.shape[0] == cols.shape[0]
      assert rows.shape[1] == cols.shape[1]
      assert rows.min() >= 0 and rows.max() < matrix.shape[1]
      assert cols.min() >= 0 and cols.max() < matrix.shape[2]
  
      aux_idx = torch.arange(matrix.shape[0]).unsqueeze(-1).to(matrix.device)
      return matrix[aux_idx, rows, cols]
  
  
  # Computes jacobians of a function of LieGroup arguments using finite differences
  #   function: a callable representing the function whose jacobians will be computed
  #   group_args: the group values at which the jacobians should be computed
  #   function_dim: the dimension of the function. If None, it uses the dimension of the
  #      group args.
  #   delta_mag: magnitude of the differences used to compute the jacobian
  def numeric_jacobian(
      function: Callable[[List[th.LieGroup]], th.LieGroup],
      group_args: List[th.LieGroup],
      function_dim: Optional[int] = None,
      delta_mag: float = 1e-3,
  ):
      batch_size = max([group_arg.shape[0] for group_arg in group_args])
  
      def _compute(group_idx):
          dof = group_args[group_idx].dof()
          function_dim_ = function_dim or dof
          jac = torch.zeros(
              batch_size, function_dim_, dof, dtype=group_args[group_idx].dtype
          )
          for d in range(dof):
              delta = torch.zeros(1, dof).to(
                  device=group_args[0].device, dtype=group_args[group_idx].dtype
              )
              delta[:, d] = delta_mag
  
              group_plus = group_args[group_idx].retract(delta)
              group_minus = group_args[group_idx].retract(-delta)
              group_plus_args = [g for g in group_args]
              group_plus_args[group_idx] = group_plus
              group_minus_args = [g for g in group_args]
              group_minus_args[group_idx] = group_minus
  
              diff = function(group_minus_args).local(function(group_plus_args))
              jac[:, :, d] = diff / (2 * delta_mag)
          return jac
  
      jacs = []
      for group_idx in range(len(group_args)):
          jacs.append(_compute(group_idx))
      return jacs
  
  
  # Returns a function that approximates a function's gradient with a central difference.
  def numeric_grad(
      f: Callable[[np.ndarray], np.ndarray], h: float = 1e-4
  ) -> Callable[[np.ndarray], np.ndarray]:
      def df(x: np.ndarray):
          assert x.ndim == 1
          n = x.shape[0]
          g = np.zeros_like(x)
          for i in range(n):
              h_i = np.zeros_like(x)
              h_i[i] = h
              g[i] = (f(x + h_i) - f(x - h_i)) / (2.0 * h)
          return g
  
      return df
  
  
  # Updates the given variable with a random tensor of the same shape as the original.
  def _rand_fill_(v: th.Variable, batch_size: int):
      if isinstance(v, (th.SE2, th.SO3, th.SE3, th.SO3)):
          v.update(v.rand(batch_size, dtype=v.dtype, device=v.device).tensor)
      else:
          v.update(
              torch.rand((batch_size,) + v.shape[1:], dtype=v.dtype, device=v.device)
          )
  
  
  # Automatically checks the jacobians of the given cost function a number of times.
  #
  # Computes the manifold jacobians of the given cost function with respect to all
  # optimization variables, evaluated at randomly sampled values
  # of the optimization and auxiliary variable, and compares them with the corresponding
  # ones computed by torch autograd. By default, only checks once, but more checks can
  # be specified, with one set of sampled variables per each. The jacobians are
  # compared using the infinity norm of the jacobian matrix, at the specified tolerance.
  @torch.no_grad()
  def check_jacobians(cf: th.CostFunction, num_checks: int = 1, tol: float = 1.0e-3):
      from theseus.core.cost_function import _tmp_tensors
  
      optim_vars: List[th.Manifold] = list(cf.optim_vars)
      aux_vars = list(cf.aux_vars)
  
      def autograd_fn(*optim_var_tensors):
          for v, t in zip(optim_vars, optim_var_tensors):
              v.update(t)
          return cf.error()
  
      with _tmp_tensors(optim_vars), _tmp_tensors(aux_vars):
          for _ in range(num_checks):
              for v in optim_vars + aux_vars:
                  _rand_fill_(v, 1)
  
              autograd_jac = torch.autograd.functional.jacobian(
                  autograd_fn, tuple(v.tensor for v in optim_vars)
              )
              jac, _ = cf.jacobians()
              for idx, v in enumerate(optim_vars):
                  j1 = jac[idx][0]
                  j2 = autograd_jac[idx]
                  # In some "unfriendly" cost functions, the error's batch size could
                  # be different than the optim/aux vars batch size, if they save
                  # tensors that are not exposed as Theseus variables. To avoid issues,
                  # we just check the first element of the batch.
                  j2_sparse = j2[0, :, 0, :]
                  j2_sparse_manifold = v.project(j2_sparse, is_sparse=True)
                  if (j1 - j2_sparse_manifold).abs().max() > tol:
                      raise RuntimeError(
                          f"Jacobian for variable {v.name} appears incorrect to the "
                          "given tolerance."
                      )
  
  
  # A timer utility that adapts to the device. Useful for removing
  # boilerplate code when benchmarking tasks, and collect statistics.
  # For CPU it uses time.perf_counter_ns()
  # For GPU it uses torch.cuda.Event()
  #
  # Usage:
  #
  # from thesus.utils import Timer
  #
  # with Timer("cuda:0") as timer:
  #    do_some_stuff()
  # print(timer.elapsed_time)
  ##
  # timer = Timer("cuda:0")
  # timer.start()
  # do_some_stuff()
  # timer.end()
  # print(timer.elapsed_time)
  #
  # The timer can also optionally collect history of times in a dictionary
  # by adding caller ids to each context or start() calls.
  # For example,
  #
  # timer.start("f1")
  # f1()
  # timer.end()
  # with timer("f2"):
  #   f2()
  # print(timer.stats())
  #
  # The timer can also be set inactive either via constructor
  # (i.e., Timer(device, active=False))
  # or via start method (i.e., timer.start(call_id, deactivate=True)).
  class Timer:
      def __init__(self, device: th.DeviceType, active: bool = True) -> None:
          self.active = active
          self.device = torch.device(device)
          self.elapsed_time = 0.0
          self._stats: Dict[str, List[float]] = defaultdict(list)
          self._caller: Optional[str] = None
          self._tmp_deactivated = False
  
      def start(self, caller: Optional[str] = None, deactivate: bool = False) -> "Timer":
          if not self.active or deactivate:
              self._tmp_deactivated = deactivate
              return self
          if self.device.type == "cuda":
              self._start_event = torch.cuda.Event(enable_timing=True)
              self._end_event = torch.cuda.Event(enable_timing=True)
              self._start_event.record()
          else:
              self._start_time = time.perf_counter_ns()
          self._caller = caller
          return self
  
      def end(self) -> None:
          if not self.active or self._tmp_deactivated:
              self._tmp_deactivated = False
              return
          if self.device.type == "cuda":
              self._end_event.record()
              torch.cuda.synchronize()
              self.elapsed_time = self._start_event.elapsed_time(self._end_event) / 1e3
          else:
              self.elapsed_time = (time.perf_counter_ns() - self._start_time) / 1e9
          if self._caller is not None:
              self._stats[self._caller].append(self.elapsed_time)
  
      def __call__(
          self, caller: Optional[str] = None, deactivate: bool = False
      ) -> "Timer":
          self._caller = caller
          self._tmp_deactivated = deactivate
          return self
  
      def __enter__(self) -> "Timer":
          self.start(caller=self._caller, deactivate=self._tmp_deactivated)
          return self
  
      def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
          self.end()
  
      def stats(self) -> Dict[str, List[float]]:
          return self._stats
  
  
  # Wrapper for cProfile.Profile for easily make optional, turn on/off and printing
  class Profiler:
      def __init__(self, active: bool):
          self.c_profiler = cProfile.Profile()
          self.active = active
  
      def enable(self):
          if self.active:
              self.c_profiler.enable()
  
      def disable(self):
          if self.active:
              self.c_profiler.disable()
  
      def print(self):
          if self.active:
              s = io.StringIO()
              sortby = pstats.SortKey.CUMULATIVE
              ps = pstats.Stats(self.c_profiler, stream=s).sort_stats(sortby)
              ps.print_stats()
              print(s.getvalue())
  
      def create_stats(self):
          if self.active:
              self.c_profiler.create_stats()
  
      def dump_stats(self, filename: str):
          if self.active:
              self.c_profiler.dump_stats(filename)
  