text: # # Unity ML-Agents Toolkit
  from time import perf_counter
  
  from contextlib import contextmanager
  from typing import Any, Callable, Dict, Generator, TypeVar
  
  """
  Lightweight, hierarchical timers for profiling sections of code.
  
  Example:
  
  @timed
  def foo(t):
      time.sleep(t)
  
  def main():
      for i in range(3):
          foo(i + 1)
      with hierarchical_timer("context"):
          foo(1)
  
      print(get_timer_tree())
  
  This would produce a timer tree like
      (root)
          "foo"
          "context"
              "foo"
  
  The total time and counts are tracked for each block of code; in this example "foo" and "context.foo" are considered
  distinct blocks, and are tracked separately.
  
  The decorator and contextmanager are equivalent; the context manager may be more useful if you want more control
  over the timer name, or are splitting up multiple sections of a large function.
  """
  
  
  class TimerNode:
      """
      Represents the time spent in a block of code.
      """
  
      __slots__ = ["children", "total", "count", "is_parallel"]
  
      def __init__(self):
          # Note that since dictionary keys are the node names, we don't explicitly store the name on the TimerNode.
          self.children: Dict[str, TimerNode] = {}
          self.total: float = 0.0
          self.count: int = 0
          self.is_parallel = False
  
      def get_child(self, name: str) -> "TimerNode":
          """
          Get the child node corresponding to the name (and create if it doesn't already exist).
          """
          child = self.children.get(name)
          if child is None:
              child = TimerNode()
              self.children[name] = child
          return child
  
      def add_time(self, elapsed: float) -> None:
          """
          Accumulate the time spent in the node (and increment the count).
          """
          self.total += elapsed
          self.count += 1
  
      def merge(self, other: "TimerNode", root_name: str = None, is_parallel=True):
          """
          Add the other node to this node, then do the same recursively on its children.
          :param other: The other node to merge
          :param root_name: Optional name of the root node being merged.
          :param is_parallel: Whether or not the code block was executed in parallel.
          :return:
          """
          if root_name:
              node = self.get_child(root_name)
          else:
              node = self
  
          node.total += other.total
          node.count += other.count
          node.is_parallel |= is_parallel
          for other_child_name, other_child_node in other.children.items():
              child = node.get_child(other_child_name)
              child.merge(other_child_node, is_parallel=is_parallel)
  
  
  class TimerStack:
      """
      Tracks all the time spent. Users shouldn't use this directly, they should use the contextmanager below to make
      sure that pushes and pops are already matched.
      """
  
      __slots__ = ["root", "stack", "start_time"]
  
      def __init__(self):
          self.root = TimerNode()
          self.stack = [self.root]
          self.start_time = perf_counter()
  
      def reset(self):
          self.root = TimerNode()
          self.stack = [self.root]
          self.start_time = perf_counter()
  
      def push(self, name: str) -> TimerNode:
          """
          Called when entering a new block of code that is timed (e.g. with a contextmanager).
          """
          current_node: TimerNode = self.stack[-1]
          next_node = current_node.get_child(name)
          self.stack.append(next_node)
          return next_node
  
      def pop(self) -> None:
          """
          Called when exiting a new block of code that is timed (e.g. with a contextmanager).
          """
          self.stack.pop()
  
      def get_root(self) -> TimerNode:
          """
          Update the total time and count of the root name, and return it.
          """
          root = self.root
          root.total = perf_counter() - self.start_time
          root.count = 1
          return root
  
      def get_timing_tree(self, node: TimerNode = None) -> Dict[str, Any]:
          """
          Recursively build a tree of timings, suitable for output/archiving.
          """
          res: Dict[str, Any] = {}
          if node is None:
              # Special case the root - total is time since it was created, and count is 1
              node = self.get_root()
              res["name"] = "root"
  
          res["total"] = node.total
          res["count"] = node.count
  
          if node.is_parallel:
              # Note when the block ran in parallel, so that it's less confusing that a timer is less that its children.
              res["is_parallel"] = True
  
          child_total = 0.0
          child_list = []
          for child_name, child_node in node.children.items():
              child_res: Dict[str, Any] = {
                  "name": child_name,
                  **self.get_timing_tree(child_node),
              }
              child_list.append(child_res)
              child_total += child_res["total"]
  
          # "self" time is total time minus all time spent on children
          res["self"] = max(0.0, node.total - child_total)
          if child_list:
              res["children"] = child_list
  
          return res
  
  
  # Global instance of a TimerStack. This is generally all that we need for profiling, but you can potentially
  # create multiple instances and pass them to the contextmanager
  _global_timer_stack = TimerStack()
  
  
  @contextmanager
  def hierarchical_timer(name: str, timer_stack: TimerStack = None) -> Generator:
      """
      Creates a scoped timer around a block of code. This time spent will automatically be incremented when
      the context manager exits.
      """
      timer_stack = timer_stack or _global_timer_stack
      timer_node = timer_stack.push(name)
      start_time = perf_counter()
  
      try:
          # The wrapped code block will run here.
          yield timer_node
      finally:
          # This will trigger either when the context manager exits, or an exception is raised.
          # We'll accumulate the time, and the exception (if any) gets raised automatically.
          elapsed = perf_counter() - start_time
          timer_node.add_time(elapsed)
          timer_stack.pop()
  
  
  # This is used to ensure the signature of the decorated function is preserved
  # See also https://github.com/python/mypy/issues/3157
  FuncT = TypeVar("FuncT", bound=Callable[..., Any])
  
  
  def timed(func: FuncT) -> FuncT:
      """
      Decorator for timing a function or method. The name of the timer will be the qualified name of the function.
      Usage:
          @timed
          def my_func(x, y):
              return x + y
      Note that because this doesn't take arguments, the global timer stack is always used.
      """
  
      def wrapped(*args, **kwargs):
          with hierarchical_timer(func.__qualname__):
              return func(*args, **kwargs)
  
      return wrapped  # type: ignore
  
  
  def get_timer_tree(timer_stack: TimerStack = None) -> Dict[str, Any]:
      """
      Return the tree of timings from the TimerStack as a dictionary (or the global stack if none is provided)
      """
      timer_stack = timer_stack or _global_timer_stack
      return timer_stack.get_timing_tree()
  
  
  def get_timer_root(timer_stack: TimerStack = None) -> TimerNode:
      """
      Get the root TimerNode of the timer_stack (or the global TimerStack if not specified)
      """
      timer_stack = timer_stack or _global_timer_stack
      return timer_stack.get_root()
  
  
  def reset_timers(timer_stack: TimerStack = None) -> None:
      """
      Reset the timer_stack (or the global TimerStack if not specified)
      """
      timer_stack = timer_stack or _global_timer_stack
      timer_stack.reset()
  