# Copyright 2017 The dm_control Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

"""Main user-facing classes and utility functions for loading MuJoCo models."""

import contextlib
import copy
import ctypes
from typing import Union
import weakref

from absl import logging

from dm_control.mujoco.wrapper import util
# Some clients explicitly import core.mjlib.
from dm_control.mujoco.wrapper.mjbindings import mjlib  # pylint: disable=unused-import
import mujoco
import numpy as np

# Unused internal import: resources.

_FAKE_BINARY_FILENAME = "model.mjb"

_CONTACT_ID_OUT_OF_RANGE = (
    "`contact_id` must be between 0 and {max_valid} (inclusive), got: {actual}."
)


class Error(Exception):
  """Base class for MuJoCo exceptions."""
  pass


if mujoco.mjVERSION_HEADER != mujoco.mj_version():
  raise Error("MuJoCo library version ({0}) does not match header version "
              "({1})".format(mujoco.mjVERSION_HEADER, mujoco.mj_version()))

# This is used to keep track of the `MJMODEL` pointer that was most recently
# loaded by `_get_model_ptr_from_xml`. Only this model can be saved to XML.
_LAST_PARSED_MODEL_PTR = None

_NOT_LAST_PARSED_ERROR = (
    "Only the model that was most recently loaded from an XML file or string "
    "can be saved to an XML file.")

import time

# NB: Python functions that are called from C are defined at module-level to
# ensure they won't be garbage-collected before they are called.
@ctypes.CFUNCTYPE(None, ctypes.c_char_p)
def _warning_callback(message):
  logging.warning(util.to_native_string(message))


@ctypes.CFUNCTYPE(None, ctypes.c_char_p)
def _error_callback(message):
  logging.fatal(util.to_native_string(message))


# Override MuJoCo's callbacks for handling warnings
mujoco.set_mju_user_warning(_warning_callback)


def enable_timer(enabled=True):
  if enabled:
    set_callback("mjcb_time", time.time)
  else:
    set_callback("mjcb_time", None)


def _str2type(type_str):
  type_id = mujoco.mju_str2Type(util.to_binary_string(type_str))
  if not type_id:
    raise Error("{!r} is not a valid object type name.".format(type_str))
  return type_id


def _type2str(type_id):
  type_str_ptr = mujoco.mju_type2Str(type_id)
  if not type_str_ptr:
    raise Error("{!r} is not a valid object type ID.".format(type_id))
  return type_str_ptr


def set_callback(name, new_callback=None):
  """Sets a user-defined callback function to modify MuJoCo's behavior.

  Callback functions should have the following signature:
    func(const_mjmodel_ptr, mjdata_ptr) -> None

  Args:
    name: Name of the callback to set. Must be a field in
      `functions.function_pointers`.
    new_callback: The new callback. This can be one of the following:
      * A Python callable
      * A C function exposed by a `ctypes.CDLL` object
      * An integer specifying the address of a callback function
      * None, in which case any existing callback of that name is removed
  """
  getattr(mujoco, "set_" + name)(new_callback)


@contextlib.contextmanager
def callback_context(name, new_callback=None):
  """Context manager that temporarily overrides a MuJoCo callback function.

  On exit, the callback will be restored to its original value (None if the
  callback was not already overridden when the context was entered).

  Args:
    name: Name of the callback to set. Must be a field in
      `mjbindings.function_pointers`.
    new_callback: The new callback. This can be one of the following:
      * A Python callable
      * A C function exposed by a `ctypes.CDLL` object
      * An integer specifying the address of a callback function
      * None, in which case any existing callback of that name is removed

  Yields:
    None
  """
  old_callback = getattr(mujoco, "get_" + name)()
  set_callback(name, new_callback)
  try:
    yield
  finally:
    # Ensure that the callback is reset on exit, even if an exception is raised.
    set_callback(name, old_callback)


def get_schema():
  """Returns a string containing the schema used by the MuJoCo XML parser."""
  buf = ctypes.create_string_buffer(100000)
  mujoco.mj_printSchema(None, buf, len(buf), 0, 0)
  return buf.value


def _get_model_ptr_from_xml(xml_path=None, xml_string=None, assets=None):
  """Parses a model XML file, compiles it, and returns a pointer to an mjModel.

  Args:
    xml_path: None or a path to a model XML file in MJCF or URDF format.
    xml_string: None or an XML string containing an MJCF or URDF model
      description.
    assets: None or a dict containing external assets referenced by the model
      (such as additional XML files, textures, meshes etc.), in the form of
      `{filename: contents_string}` pairs. The keys should correspond to the
      filenames specified in the model XML. Ignored if `xml_string` is None.

    One of `xml_path` or `xml_string` must be specified.

  Returns:
    A `ctypes.POINTER` to a new `mjbindings.types.MJMODEL` instance.

  Raises:
    TypeError: If both or neither of `xml_path` and `xml_string` are specified.
    Error: If the model is not created successfully.
  """
  if xml_path is None and xml_string is None:
    raise TypeError(
        "At least one of `xml_path` or `xml_string` must be specified.")
  elif xml_path is not None and xml_string is not None:
    raise TypeError(
        "Only one of `xml_path` or `xml_string` may be specified.")

  if xml_string is not None:
    ptr = mujoco.MjModel.from_xml_string(xml_string, assets or {})
  else:
    ptr = mujoco.MjModel.from_xml_path(xml_path, assets or {})

  global _LAST_PARSED_MODEL_PTR
  _LAST_PARSED_MODEL_PTR = ptr

  return ptr


def save_last_parsed_model_to_xml(xml_path, check_model=None):
  """Writes a description of the most recently loaded model to an MJCF XML file.

  Args:
    xml_path: Path to the output XML file.
    check_model: Optional `MjModel` instance. If specified, this model will be
      checked to see if it is the most recently parsed one, and a ValueError
      will be raised otherwise.
  Raises:
    Error: If MuJoCo encounters an error while writing the XML file.
    ValueError: If `check_model` was passed, and this model is not the most
      recently parsed one.
  """
  if check_model and check_model.ptr is not _LAST_PARSED_MODEL_PTR:
    raise ValueError(_NOT_LAST_PARSED_ERROR)
  mujoco.mj_saveLastXML(xml_path, _LAST_PARSED_MODEL_PTR)


def _get_model_ptr_from_binary(binary_path=None, byte_string=None):
  """Returns a pointer to an mjModel from the contents of a MuJoCo model binary.

  Args:
    binary_path: Path to an MJB file (as produced by MjModel.save_binary).
    byte_string: String of bytes (as returned by MjModel.to_bytes).

    One of `binary_path` or `byte_string` must be specified.

  Returns:
    A `ctypes.POINTER` to a new `mjbindings.types.MJMODEL` instance.

  Raises:
    TypeError: If both or neither of `byte_string` and `binary_path`
      are specified.
  """
  if binary_path is None and byte_string is None:
    raise TypeError(
        "At least one of `byte_string` or `binary_path` must be specified.")
  elif binary_path is not None and byte_string is not None:
    raise TypeError(
        "Only one of `byte_string` or `binary_path` may be specified.")

  if byte_string is not None:
    assets = {_FAKE_BINARY_FILENAME: byte_string}
    return mujoco.MjModel.from_binary_path(_FAKE_BINARY_FILENAME, assets)
  return mujoco.MjModel.from_binary_path(binary_path, {})


class _MjModelMeta(type):
  """Metaclass which allows MjModel below to delegate to mujoco.MjModel."""

  def __new__(mcs, name, bases, dct):
    for attr in dir(mujoco.MjModel):
      if not attr.startswith("_"):
        if attr not in dct:
          # pylint: disable=protected-access
          fget = lambda self, attr=attr: getattr(self._model, attr)
          fset = (
              lambda self, value, attr=attr: setattr(self._model, attr, value))
          # pylint: enable=protected-access
          dct[attr] = property(fget, fset)
    return super().__new__(mcs, name, bases, dct)


class MjModel(metaclass=_MjModelMeta):
  """Wrapper class for a MuJoCo 'mjModel' instance.

  MjModel encapsulates features of the model that are expected to remain
  constant. It also contains simulation and visualization options which may be
  changed occasionally, although this is done explicitly by the user.
  """
  _HAS_DYNAMIC_ATTRIBUTES = True

  def __init__(self, model_ptr):
    """Creates a new MjModel instance from a mujoco.MjModel."""
    self._model = model_ptr

  @property
  def ptr(self):
    """The lower level MjModel instance."""
    return self._model

  def __getstate__(self):
    return self._model

  def __setstate__(self, state):
    self._model = state

  def __copy__(self):
    new_model_ptr = copy.copy(self._model)
    return self.__class__(new_model_ptr)

  @classmethod
  def from_xml_string(cls, xml_string, assets=None):
    """Creates an `MjModel` instance from a model description XML string.

    Args:
      xml_string: String containing an MJCF or URDF model description.
      assets: Optional dict containing external assets referenced by the model
        (such as additional XML files, textures, meshes etc.), in the form of
        `{filename: contents_string}` pairs. The keys should correspond to the
        filenames specified in the model XML.

    Returns:
      An `MjModel` instance.
    """
    model_ptr = _get_model_ptr_from_xml(xml_string=xml_string, assets=assets)
    return cls(model_ptr)

  @classmethod
  def from_byte_string(cls, byte_string):
    """Creates an MjModel instance from a model binary as a string of bytes."""
    model_ptr = _get_model_ptr_from_binary(byte_string=byte_string)
    return cls(model_ptr)

  @classmethod
  def from_xml_path(cls, xml_path):
    """Creates an MjModel instance from a path to a model XML file."""
    model_ptr = _get_model_ptr_from_xml(xml_path=xml_path)
    return cls(model_ptr)

  @classmethod
  def from_binary_path(cls, binary_path):
    """Creates an MjModel instance from a path to a compiled model binary."""
    model_ptr = _get_model_ptr_from_binary(binary_path=binary_path)
    return cls(model_ptr)

  def save_binary(self, binary_path):
    """Saves the MjModel instance to a binary file."""
    mujoco.mj_saveModel(self.ptr, binary_path, None)

  def to_bytes(self):
    """Serialize the model to a string of bytes."""
    bufsize = mujoco.mj_sizeModel(self.ptr)
    buf = np.zeros(shape=(bufsize,), dtype=np.uint8)
    mujoco.mj_saveModel(self.ptr, None, buf)
    return buf.tobytes()

  def copy(self):
    """Returns a copy of this MjModel instance."""
    return self.__copy__()

  def free(self):
    """Frees the native resources held by this MjModel.

    This is an advanced feature for use when manual memory management is
    necessary. This MjModel object MUST NOT be used after this function has
    been called.
    """
    del self._ptr

  def name2id(self, name, object_type):
    """Returns the integer ID of a specified MuJoCo object.

    Args:
      name: String specifying the name of the object to query.
      object_type: The type of the object. Can be either a lowercase string
        (e.g. 'body', 'geom') or an `mjtObj` enum value.

    Returns:
      An integer object ID.

    Raises:
      Error: If `object_type` is not a valid MuJoCo object type, or if no object
        with the corresponding name and type was found.
    """
    if isinstance(object_type, str):
      object_type = _str2type(object_type)
    obj_id = mujoco.mj_name2id(self.ptr, object_type, name)
    if obj_id == -1:
      raise Error("Object of type {!r} with name {!r} does not exist.".format(
          _type2str(object_type), name))
    return obj_id

  def id2name(self, object_id, object_type):
    """Returns the name associated with a MuJoCo object ID, if there is one.

    Args:
      object_id: Integer ID.
      object_type: The type of the object. Can be either a lowercase string
        (e.g. 'body', 'geom') or an `mjtObj` enum value.

    Returns:
      A string containing the object name, or an empty string if the object ID
      either doesn't exist or has no name.

    Raises:
      Error: If `object_type` is not a valid MuJoCo object type.
    """
    if isinstance(object_type, str):
      object_type = _str2type(object_type)
    return mujoco.mj_id2name(self.ptr, object_type, object_id) or ""

  @contextlib.contextmanager
  def disable(self, *flags):
    """Context manager for temporarily disabling MuJoCo flags.

    Args:
      *flags: Positional arguments specifying flags to disable. Can be either
        lowercase strings (e.g. 'gravity', 'contact') or `mjtDisableBit` enum
        values.

    Yields:
      None

    Raises:
      ValueError: If any item in `flags` is neither a valid name nor a value
        from `mujoco.mjtDisableBit`.
    """
    old_bitmask = self.opt.disableflags
    new_bitmask = old_bitmask
    for flag in flags:
      if isinstance(flag, str):
        try:
          field_name = "mjDSBL_" + flag.upper()
          flag = getattr(mujoco.mjtDisableBit, field_name)
        except AttributeError:
          valid_names = [
              field_name.split("_")[1].lower()
              for field_name in list(mujoco.mjtDisableBit.__members__)[:-1]
          ]
          raise ValueError("'{}' is not a valid flag name. Valid names: {}"
                           .format(flag, ", ".join(valid_names))) from None
      elif isinstance(flag, int):
        flag = mujoco.mjtDisableBit(flag)
      new_bitmask |= flag.value
    self.opt.disableflags = new_bitmask
    try:
      yield
    finally:
      self.opt.disableflags = old_bitmask

  @property
  def name(self):
    """Returns the name of the model."""
    # The model's name is the first null terminated string in _model.names
    return str(self._model.names[:self._model.names.find(b"\0")], "ascii")


class _MjDataMeta(type):
  """Metaclass which allows MjData below to delegate to mujoco.MjData."""

  def __new__(mcs, name, bases, dct):
    for attr in dir(mujoco.MjData):
      if not attr.startswith("_"):
        if attr not in dct:
          # pylint: disable=protected-access
          fget = lambda self, attr=attr: getattr(self._data, attr)
          fset = lambda self, value, attr=attr: setattr(self._data, attr, value)
          # pylint: enable=protected-access
          dct[attr] = property(fget, fset)
    return super().__new__(mcs, name, bases, dct)


class MjData(metaclass=_MjDataMeta):
  """Wrapper class for a MuJoCo 'mjData' instance.

  MjData contains all of the dynamic variables and intermediate results produced
  by the simulation. These are expected to change on each simulation timestep.
  """

  _HAS_DYNAMIC_ATTRIBUTES = True

  def __init__(self, model_or_data: Union[MjModel, mujoco.MjData]):
    """Constructs a new MjData instance.

    Args:
      model_or_data: dm_control.mujoco.wrapper.MjModel instance, or
          mujoco.MjData.
    """
    if isinstance(model_or_data, MjModel):
      self._model = model_or_data
      self._data = mujoco.MjData(model_or_data._model)
    elif isinstance(model_or_data, mujoco.MjData):
      self._data = model_or_data
      self._model = MjModel(self._data.model)

  def __getstate__(self):
    return self._data

  def __setstate__(self, state):
    self._data = state
    self._model = MjModel(self._data.model)

  def __copy__(self):
    # This makes a shallow copy that shares the same parent MjModel instance.
    return self._make_copy(share_model=True)

  def _make_copy(self, share_model):
    if share_model:
      return self.__class__(copy.copy(self._data))
    return self.__class__(copy.deepcopy(self._data))

  def copy(self):
    """Returns a copy of this MjData instance with the same parent MjModel."""
    return self.__copy__()

  def object_velocity(self, object_id, object_type, local_frame=False):
    """Returns the 6D velocity (linear, angular) of a MuJoCo object.

    Args:
      object_id: Object identifier. Can be either integer ID or String name.
      object_type: The type of the object. Can be either a lowercase string
        (e.g. 'body', 'geom') or an `mjtObj` enum value.
      local_frame: Boolean specifiying whether the velocity is given in the
        global (worldbody), or local (object) frame.

    Returns:
      2x3 array with stacked (linear_velocity, angular_velocity)

    Raises:
      Error: If `object_type` is not a valid MuJoCo object type, or if no object
        with the corresponding name and type was found.
    """
    if not isinstance(object_type, int):
      object_type = _str2type(object_type)
    velocity = np.empty(6, dtype=np.float64)
    if not isinstance(object_id, int):
      object_id = self.model.name2id(object_id, object_type)
    mujoco.mj_objectVelocity(self._model.ptr, self._data, object_type,
                             object_id, velocity, local_frame)
    #  MuJoCo returns velocities in (angular, linear) order, which we flip here.
    return velocity.reshape(2, 3)[::-1]

  def contact_force(self, contact_id):
    """Returns the wrench of a contact as a 2 x 3 array of (forces, torques).

    Args:
      contact_id: Integer, the index of the contact within the contact buffer
        (`self.contact`).

    Returns:
      2x3 array with stacked (force, torque). Note that the order of dimensions
        is (normal, tangent, tangent), in the contact's frame.

    Raises:
      ValueError: If `contact_id` is negative or bigger than ncon-1.
    """
    if not 0 <= contact_id < self.ncon:
      raise ValueError(_CONTACT_ID_OUT_OF_RANGE
                       .format(max_valid=self.ncon-1, actual=contact_id))

    # Run the portion of `mj_step2` that are needed for correct contact forces.
    mujoco.mj_fwdActuation(self._model.ptr, self._data)
    mujoco.mj_fwdAcceleration(self._model.ptr, self._data)
    mujoco.mj_fwdConstraint(self._model.ptr, self._data)

    wrench = np.empty(6, dtype=np.float64)
    mujoco.mj_contactForce(self._model.ptr, self._data, contact_id, wrench)
    return wrench.reshape(2, 3)

  @property
  def ptr(self):
    """The lower level MjData instance."""
    return self._data

  @property
  def model(self):
    """The parent MjModel for this MjData instance."""
    return self._model

  @property
  def contact(self):
    """Variable-length recarray containing all current contacts."""
    return self._data.contact[:self.ncon]


# Docstrings for these subclasses are inherited from their parent class.


class MjvCamera(mujoco.MjvCamera):  # pylint: disable=missing-docstring

  # Provide this alias for the "type" property for backwards compatibility.
  @property
  def type_(self):
    return self.type

  @type_.setter
  def type_(self, t):
    self.type = t

  @property
  def ptr(self):
    return self


class MjvOption(mujoco.MjvOption):  # pylint: disable=missing-docstring

  def __init__(self):
    super().__init__()
    self.flags[mujoco.mjtVisFlag.mjVIS_RANGEFINDER] = False

  # Provided for backwards compatibility
  @property
  def ptr(self):
    return self


class MjrContext:
  """Wrapper for mujoco.MjrContext."""

  def __init__(self,
               model,
               gl_context,
               font_scale=mujoco.mjtFontScale.mjFONTSCALE_150):
    """Initializes this MjrContext instance.

    Args:
      model: An `MjModel` instance.
      gl_context: A `render.ContextBase` instance.
      font_scale: Integer controlling the font size for text. Must be a value
        in `mujoco.mjtFontScale`.

    Raises:
      ValueError: If `font_scale` is invalid.
    """
    if not isinstance(font_scale, mujoco.mjtFontScale):
      font_scale = mujoco.mjtFontScale(font_scale)
    self._gl_context = gl_context
    with gl_context.make_current() as ctx:
      ptr = ctx.call(mujoco.MjrContext, model.ptr, font_scale)
      ctx.call(mujoco.mjr_setBuffer, mujoco.mjtFramebuffer.mjFB_OFFSCREEN, ptr)
    gl_context.keep_alive(ptr)
    gl_context.increment_refcount()
    self._ptr = weakref.ref(ptr)

  @property
  def ptr(self):
    return self._ptr()

  def free(self):
    """Frees the native resources held by this MjrContext.

    This is an advanced feature for use when manual memory management is
    necessary. This MjrContext object MUST NOT be used after this function has
    been called.
    """
    if self._gl_context and not self._gl_context.terminated:
      ptr = self.ptr
      if ptr:
        self._gl_context.dont_keep_alive(ptr)
        with self._gl_context.make_current() as ctx:
          ctx.call(ptr.free)

    if self._gl_context:
      self._gl_context.decrement_refcount()
      self._gl_context.free()
      self._gl_context = None

  def __del__(self):
    self.free()

# A mapping from human-readable short names to mjtRndFlag enum values, i.e.
# {'shadow': mjtRndFlag.mjRND_SHADOW, 'fog': mjtRndFlag.mjRND_FOG, ...}
_NAME_TO_RENDER_FLAG_ENUM_VALUE = {
    name[len("mjRND_"):].lower(): getattr(mujoco.mjtRndFlag, name).value
    for name in mujoco.mjtRndFlag.__members__
    if name != "mjRND_NUMRNDFLAG"
}


def _estimate_max_renderable_geoms(model):
  """Estimates the maximum number of renderable geoms for a given model."""
  # Only one type of object frame can be rendered at once.
  max_nframes = max(
      [model.nbody, model.ngeom, model.nsite, model.ncam, model.nlight])
  # This is probably an underestimate, but it is unlikely that all possible
  # rendering options will be enabled simultaneously, or that all renderable
  # geoms will be present within the viewing frustum at the same time.
  return (
      3 * max_nframes +  # 1 geom per axis for each frame.
      4 * model.ngeom +  # geom itself + contacts + 2 * split contact forces.
      3 * model.nbody +  # COM + inertia box + perturbation force.
      model.nsite +
      model.ntendon +
      model.njnt +
      model.nu +
      model.nskin +
      model.ncam +
      model.nlight)


class MjvScene(mujoco.MjvScene):  # pylint: disable=missing-docstring

  def __init__(self, model=None, max_geom=None):
    """Initializes a new `MjvScene` instance.

    Args:
      model: (optional) An `MjModel` instance.
      max_geom: (optional) An integer specifying the maximum number of geoms
        that can be represented in the scene. If None, this will be chosen
        automatically based on `model`.
    """
    if model is None:
      super().__init__()
    else:
      if max_geom is None:
        if model is None:
          max_renderable_geoms = 0
        else:
          max_renderable_geoms = _estimate_max_renderable_geoms(model)
        max_geom = max(1000, max_renderable_geoms)

      super().__init__(model.ptr, max_geom)

  @property
  def ptr(self):
    return self

  @contextlib.contextmanager
  def override_flags(self, overrides):
    """Context manager for temporarily overriding rendering flags.

    Args:
      overrides: A mapping specifying rendering flags to override. The keys can
        be either lowercase strings or `mjtRndFlag` enum values, and the values
        are the overridden flag values, e.g. `{'wireframe': True}` or
        `{mujoco.mjtRndFlag.mjRND_WIREFRAME: True}`. See `mujoco.mjtRndFlag` for
        the set of valid flags.

    Yields:
      None
    """
    if not overrides:
      yield
    else:
      original_flags = self.flags.copy()
      for key, value in overrides.items():
        index = _NAME_TO_RENDER_FLAG_ENUM_VALUE.get(key, key)
        self.flags[index] = value
      try:
        yield
      finally:
        np.copyto(self.flags, original_flags)

  def free(self):
    """Frees the native resources held by this MjvScene.

    This is an advanced feature for use when manual memory management is
    necessary. This MjvScene object MUST NOT be used after this function has
    been called.
    """
    pass

  @property
  def geoms(self):
    """Variable-length recarray containing all geoms currently in the buffer."""
    return super().geoms[:super().ngeom]


class MjvPerturb(mujoco.MjvPerturb):  # pylint: disable=missing-docstring

  @property
  def ptr(self):
    return self


class MjvFigure(mujoco.MjvFigure):  # pylint: disable=missing-docstring

  @property
  def ptr(self):
    return self

  @property
  def range_(self):
    return self.range

  @range_.setter
  def range_(self, value):
    self.range = value