// Copyright 2022 DeepMind Technologies Limited
//
// 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.

#include <algorithm>
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <utility>
#include <variant>
#include <vector>

#include "indexers.h"
#include "mjdata_meta.h"
#include "raw.h"
#include "util/crossplatform.h"

namespace mujoco::python {

namespace py = ::pybind11;

namespace {
// Parses raw mjModel to create a mapping from names to array indices.
//
// Args:
//   count: Number of names in the map.
//   name_offsets: Array consisting of indices that correspond to the start of
//     each name in the `names` array.
//   names: Character arrays consisting of the concatenation of all names.
template <typename IntPtr, typename CharPtr>
NameToID MakeNameToID(int count, IntPtr name_offsets, CharPtr names)
{
  NameToID name_to_id;
  for (int index = 0; index < count; ++index) {
    const char* name = &names[name_offsets[index]];
    if (name[0] != '\0') {
      name_to_id.insert({name, index});
    }
  }
  return name_to_id;
}

// Parses raw mjModel to create a mapping from array indices to names.
//
// Args:
//   count: Number of names in the map.
//   name_offsets: Array consisting of indices that correspond to the start of
//     each name in the `names` array.
//   names: Character arrays consisting of the concatenation of all names.
template <typename IntPtr, typename CharPtr>
IDToName MakeIDToName(int count, IntPtr name_offsets, CharPtr names) {
  IDToName id_to_name;
  for (int index = 0; index < count; ++index) {
    const char* name = &names[name_offsets[index]];
    id_to_name.emplace_back(name);
  }
  return id_to_name;
};

// Makes an array view into an mjModel/mjData struct field at a given index.
//
// Template args:
//   MjSize: An (int M::*) specifying the number of entities of the same type
//     as the one whose array view is to be created. This is a way to indicate
//     the MuJoCo category of the entity itself, e.g. nbody indicates that the
//     field belongs to a body.
//   T: Scalar data type of the field.
//   M: Either raw::MjModel or MjDataMetadata.
//
// Args:
//   base_ptr: Pointer to the first entry in the entire field.
//   index: Entity ID at which a view is to be created.
//   shape: The shape of the array as defined by MuJoCo. If MjSize corresponds
//     to a ragged or indirect field, the shape may be prepended by additional
//     dimensions as appropriate, e.g. if MjSize is nv then the view gains an
//     additional dimension of the size len(qvel) of the particular joint.
//   m: Used for dereferencing MjSize.
//   owner: The base object whose lifetime is tied to the returned array.
template <auto MjSize, typename T, typename M>
py::array_t<T> MakeArray(T* base_ptr, int index, std::vector<int>&& shape,
                         const M& m, py::handle owner) {
  int offset;
  if (MjSize == &M::nq) {
    offset = m.jnt_qposadr[index];
    shape.insert(
        shape.begin(),
        ((index < m.njnt-1) ? m.jnt_qposadr[index+1] : m.nq) - offset);
  } else if (MjSize == &M::nv) {
    offset = m.jnt_dofadr[index];
    shape.insert(
        shape.begin(),
        ((index < m.njnt-1) ? m.jnt_dofadr[index+1] : m.nv) - offset);
  } else if (MjSize == &M::nhfielddata) {
    offset = m.hfield_adr[index];
    shape.insert(shape.begin(), m.hfield_ncol[index]);
    shape.insert(shape.begin(), m.hfield_nrow[index]);
  } else if (MjSize == &M::ntexdata) {
    offset = m.tex_adr[index];
    shape.insert(shape.begin(), m.tex_width[index]);
    shape.insert(shape.begin(), m.tex_height[index]);
  } else if (MjSize == &M::nsensordata) {
    offset = m.sensor_adr[index];
    shape.insert(shape.begin(), m.sensor_dim[index]);
  } else if (MjSize == &M::nnumericdata) {
    offset = m.numeric_adr[index];
    shape.insert(shape.begin(), m.numeric_size[index]);
  } else if (MjSize == &M::ntupledata) {
    offset = m.tuple_adr[index];
    shape.insert(shape.begin(), m.tuple_size[index]);
  } else {
    // Do not return a NumPy array with shape () since these aren't very nice
    // to work with. Instead, always return singleton arrays with shape (1,).
    if (shape.empty()) {
      shape.push_back(1);
    }
    int size = 1;
    for (int s : shape) {
      size *= s;
    }
    offset = index * size;
  }
  return py::array_t<T>(std::move(shape), base_ptr + offset,
                        py::reinterpret_borrow<py::object>(owner));
}
}  // namespace

// M is either a raw::MjModel or MjDataMetadata.
template <typename M>
NameToIDMappings::NameToIDMappings(const M& m)
    : body(MakeNameToID(m.nbody, m.name_bodyadr, m.names)),
      jnt(MakeNameToID(m.njnt, m.name_jntadr, m.names)),
      geom(MakeNameToID(m.ngeom, m.name_geomadr, m.names)),
      site(MakeNameToID(m.nsite, m.name_siteadr, m.names)),
      cam(MakeNameToID(m.ncam, m.name_camadr, m.names)),
      light(MakeNameToID(m.nlight, m.name_lightadr, m.names)),
      mesh(MakeNameToID(m.nmesh, m.name_meshadr, m.names)),
      skin(MakeNameToID(m.nskin, m.name_skinadr, m.names)),
      hfield(MakeNameToID(m.nhfield, m.name_hfieldadr, m.names)),
      tex(MakeNameToID(m.ntex, m.name_texadr, m.names)),
      mat(MakeNameToID(m.nmat, m.name_matadr, m.names)),
      pair(MakeNameToID(m.npair, m.name_pairadr, m.names)),
      exclude(MakeNameToID(m.nexclude, m.name_excludeadr, m.names)),
      eq(MakeNameToID(m.neq, m.name_eqadr, m.names)),
      tendon(MakeNameToID(m.ntendon, m.name_tendonadr, m.names)),
      actuator(MakeNameToID(m.nu, m.name_actuatoradr, m.names)),
      sensor(MakeNameToID(m.nsensor, m.name_sensoradr, m.names)),
      numeric(MakeNameToID(m.nnumeric, m.name_numericadr, m.names)),
      text(MakeNameToID(m.ntext, m.name_textadr, m.names)),
      tuple(MakeNameToID(m.ntuple, m.name_tupleadr, m.names)),
      key(MakeNameToID(m.nkey, m.name_keyadr, m.names)) {}

// M is either a raw::MjModel or MjDataMetadata.
template <typename M>
IDToNameMappings::IDToNameMappings(const M& m)
    : body(MakeIDToName(m.nbody, m.name_bodyadr, m.names)),
      jnt(MakeIDToName(m.njnt, m.name_jntadr, m.names)),
      geom(MakeIDToName(m.ngeom, m.name_geomadr, m.names)),
      site(MakeIDToName(m.nsite, m.name_siteadr, m.names)),
      cam(MakeIDToName(m.ncam, m.name_camadr, m.names)),
      light(MakeIDToName(m.nlight, m.name_lightadr, m.names)),
      mesh(MakeIDToName(m.nmesh, m.name_meshadr, m.names)),
      skin(MakeIDToName(m.nskin, m.name_skinadr, m.names)),
      hfield(MakeIDToName(m.nhfield, m.name_hfieldadr, m.names)),
      tex(MakeIDToName(m.ntex, m.name_texadr, m.names)),
      mat(MakeIDToName(m.nmat, m.name_matadr, m.names)),
      pair(MakeIDToName(m.npair, m.name_pairadr, m.names)),
      exclude(MakeIDToName(m.nexclude, m.name_excludeadr, m.names)),
      eq(MakeIDToName(m.neq, m.name_eqadr, m.names)),
      tendon(MakeIDToName(m.ntendon, m.name_tendonadr, m.names)),
      actuator(MakeIDToName(m.nu, m.name_actuatoradr, m.names)),
      sensor(MakeIDToName(m.nsensor, m.name_sensoradr, m.names)),
      numeric(MakeIDToName(m.nnumeric, m.name_numericadr, m.names)),
      text(MakeIDToName(m.ntext, m.name_textadr, m.names)),
      tuple(MakeIDToName(m.ntuple, m.name_tupleadr, m.names)),
      key(MakeIDToName(m.nkey, m.name_keyadr, m.names)) {}

MjModelIndexer::MjModelIndexer(raw::MjModel* m, py::handle owner)
    : m_(m),
      owner_(owner),
      name_to_id_(*m),
      id_to_name_(*m)
#define XGROUP(MjModelFieldGroupedViews, field, nfield, FIELD_XMACROS) \
  , field##_(m->nfield, std::nullopt)
      MJMODEL_VIEW_GROUPS
#undef XGROUP
{}

#define XGROUP(MjModelFieldGroupedViews, field, nfield, FIELD_XMACROS) \
  MjModelFieldGroupedViews& MjModelIndexer::field(int i) {             \
    if (i >= field##_.size() || i < 0) {                               \
      throw py::index_error(IndexErrorMessage(i, field##_.size()));    \
    }                                                                  \
    auto& indexer = field##_[i];                                       \
    if (!indexer.has_value()) {                                        \
      const std::string& name = id_to_name_.field[i];                  \
      indexer.emplace(i, name, m_, owner_);                            \
    }                                                                  \
    return *indexer;                                                   \
  }
MJMODEL_VIEW_GROUPS
#undef XGROUP

#define XGROUP(MjModelFieldGroupedViews, field, nfield, FIELD_XMACROS) \
  MjModelFieldGroupedViews& MjModelIndexer::field##_by_name(           \
      std::string_view name) {                                         \
    auto item = name_to_id_.field.find(name);                          \
    if (item == name_to_id_.field.end()) {                             \
      throw py::key_error(KeyErrorMessage(name_to_id_.field, name));   \
    }                                                                  \
    return field(item->second);                                        \
  }
MJMODEL_VIEW_GROUPS
#undef XGROUP

MjDataIndexer::MjDataIndexer(raw::MjData* d, const MjDataMetadata* m,
                             py::handle owner)
    : d_(d),
      m_(m),
      owner_(owner),
      name_to_id_(*m),
      id_to_name_(*m)
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
  , field##_(m->nfield, std::nullopt)
      MJDATA_VIEW_GROUPS
#undef XGROUP
{}

#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS)    \
  MjDataGroupedViews& MjDataIndexer::field(int i) {                 \
    if (i >= field##_.size() || i < 0) {                            \
      throw py::index_error(IndexErrorMessage(i, field##_.size())); \
    }                                                               \
    auto& indexer = field##_[i];                                    \
    if (!indexer.has_value()) {                                     \
      const std::string& name = id_to_name_.field[i];               \
      indexer.emplace(i, name, d_, m_, owner_);                     \
    }                                                               \
    return *indexer;                                                \
  }
MJDATA_VIEW_GROUPS
#undef XGROUP

#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS)              \
  MjDataGroupedViews& MjDataIndexer::field##_by_name(std::string_view name) { \
    auto item = name_to_id_.field.find(name);                                 \
    if (item == name_to_id_.field.end()) {                                    \
      throw py::key_error(KeyErrorMessage(name_to_id_.field, name));          \
    }                                                                         \
    return field(item->second);                                               \
  }
MJDATA_VIEW_GROUPS
#undef XGROUP

#define MAKE_SHAPE(dim)                                              \
  MUJOCO_DIAG_IGNORE_UNUSED_LAMBDA_CAPTURE                           \
    [n = (dim)]() -> std::vector<int> {                              \
    if constexpr (std::string_view(#dim) == std::string_view("1")) { \
      return {};                                                     \
    } else {                                                         \
      return {n};                                                    \
    }                                                                \
  }()                                                                \
  MUJOCO_DIAG_UNIGNORE_UNUSED_LAMBDA_CAPTURE

#undef MJ_M
#define MJ_M(n) m_->n
#define X(type, prefix, var, dim0, dim1)                            \
  py::array_t<type> XGROUP::var() {                                 \
    if (!var##_.has_value()) {                                      \
      var##_.emplace(MakeArray<&raw::MjModel::dim0>(                \
          m_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
    }                                                               \
    return *var##_;                                                 \
  }

#define XGROUP MjModelActuatorViews
MJMODEL_ACTUATOR
#undef XGROUP

#define XGROUP MjModelBodyViews
MJMODEL_BODY
#undef XGROUP

#define XGROUP MjModelCameraViews
MJMODEL_CAMERA
#undef XGROUP

#define XGROUP MjModelEqualityViews
MJMODEL_EQUALITY
#undef XGROUP

#define XGROUP MjModelExcludeViews
MJMODEL_EXCLUDE
#undef XGROUP

#define XGROUP MjModelGeomViews
MJMODEL_GEOM
#undef XGROUP

#define XGROUP MjModelHfieldViews
MJMODEL_HFIELD
#undef XGROUP

#define XGROUP MjModelJointViews
MJMODEL_JOINT
#undef XGROUP

#define XGROUP MjModelLightViews
MJMODEL_LIGHT
#undef XGROUP

#define XGROUP MjModelMaterialViews
MJMODEL_MATERIAL
#undef XGROUP

#define XGROUP MjModelMeshViews
MJMODEL_MESH
#undef XGROUP

#define XGROUP MjModelNumericViews
MJMODEL_NUMERIC
#undef XGROUP

#define XGROUP MjModelPairViews
MJMODEL_PAIR
#undef XGROUP

#define XGROUP MjModelSensorViews
MJMODEL_SENSOR
#undef XGROUP

#define XGROUP MjModelSiteViews
MJMODEL_SITE
#undef XGROUP

#define XGROUP MjModelSkinViews
MJMODEL_SKIN
#undef XGROUP

#define XGROUP MjModelTendonViews
MJMODEL_TENDON
#undef XGROUP

#define XGROUP MjModelTextureViews
MJMODEL_TEXTURE
#undef XGROUP

#define XGROUP MjModelTupleViews
MJMODEL_TUPLE
#undef XGROUP

#define XGROUP MjModelKeyframeViews
MJMODEL_KEYFRAME
#undef XGROUP

#undef X
#define MJ_M(n) m_->n
#define X(type, prefix, var, dim0, dim1)                            \
  py::array_t<type> XGROUP::var() {                               \
    if (!var##_.has_value()) {                                      \
      var##_.emplace(MakeArray<&MjDataMetadata::dim0>(              \
          d_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
    }                                                               \
    return *var##_;                                                 \
  }

#define XGROUP MjDataActuatorViews
MJDATA_ACTUATOR
#undef XGROUP

#define XGROUP MjDataBodyViews
MJDATA_BODY
#undef XGROUP

#define XGROUP MjDataCameraViews
MJDATA_CAMERA
#undef XGROUP

#define XGROUP MjDataGeomViews
MJDATA_GEOM
#undef XGROUP

#define XGROUP MjDataJointViews
MJDATA_JOINT
#undef XGROUP

#define XGROUP MjDataLightViews
MJDATA_LIGHT
#undef XGROUP

#define XGROUP MjDataSensorViews
MJDATA_SENSOR
#undef XGROUP

#define XGROUP MjDataSiteViews
MJDATA_SITE
#undef XGROUP

#define XGROUP MjDataTendonViews
MJDATA_TENDON
#undef XGROUP

#undef X
#undef MJ_M
#define MJ_M(n) n

// Returns an error message when a non-existent name is requested from an
// indexer, which includes all valid names.
std::string KeyErrorMessage(const NameToID& map, std::string_view name) {
  // Make a sorted list of valid names
  std::vector<std::string_view> valid_names;
  valid_names.reserve(map.size());
  for (const auto& [key, value] : map) {
    valid_names.push_back(key);
  }
  std::sort(valid_names.begin(), valid_names.end());

  // Construct the error message
  std::ostringstream message;
  message << "Invalid name '" << name << "'. Valid names: [";
  int i = 0;
  for (const auto& key : valid_names) {
    message << "'" << key << "'";
    if (i < map.size() - 1) message << ", ";
    i++;
  }
  message << "]";
  return message.str();
}

std::string IndexErrorMessage(int index, int size) {
  std::ostringstream message;
  message << "Invalid index " << index << ". Valid indices from 0 to "
          << size - 1;
  return message.str();
}
}  // namespace mujoco::python