"""Base logger."""

import abc
from typing import Any, Mapping

import numpy as np
import tree


LoggingData = Mapping[str, Any]


class Logger(abc.ABC):
  """A logger has a `write` method."""

  @abc.abstractmethod
  def write(self, data: LoggingData) -> None:
    """Writes `data` to destination (file, terminal, database, etc)."""


class NoOpLogger(Logger):
  """Simple Logger which does nothing and outputs no logs.
  This should be used sparingly, but it can prove useful if we want to quiet an
  individual component and have it produce no logging whatsoever.
  """

  def write(self, data: LoggingData):
    pass


def tensor_to_numpy(value: Any):
  if hasattr(value, 'numpy'):
    return value.numpy()  # tf.Tensor (TF2).
  if hasattr(value, 'device_buffer'):
    return np.asarray(value)  # jnp.DeviceArray.
  return value


def to_numpy(values: Any):
  """Converts tensors in a nested structure to numpy.
  Converts tensors from TensorFlow to Numpy if needed without importing TF
  dependency.
  Args:
    values: nested structure with numpy and / or TF tensors.
  Returns:
    Same nested structure as values, but with numpy tensors.
  """
  return tree.map_structure(tensor_to_numpy, values)