//
// Copyright (c) 2016-2024 CNRS INRIA
//

#include <cassert>
#include <algorithm>
#include <cmath>
#include <cstddef>
#include <memory>
#include <sstream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>

#include <Eigen/Core>
#include <Eigen/Geometry>

#include <boost/math/constants/constants.hpp>
#include <boost/optional/optional.hpp>
#include <boost/property_tree/ptree.hpp>

#include "pinocchio/macros.hpp"
#include "pinocchio/spatial.hpp"
#include "pinocchio/multibody.hpp"
#include "pinocchio/geometry.hpp"
#include "pinocchio/parsers/mjcf.hpp"

#ifdef PINOCCHIO_WITH_COLLISION

  #include <coal/BV/OBBRSS.h>
  #include <coal/BVH/BVH_model.h>
  #include <coal/mesh_loader/loader.h>
  #include <coal/shape/geometric_shapes.h>

#endif // PINOCCHIO_WITH_COLLISION

namespace pinocchio
{
  namespace mjcf
  {
    namespace details
    {
      static double computeVolume(const Eigen::VectorXd & size, const std::string & geomType)
      {
        double pi = boost::math::constants::pi<double>();
        if (geomType == "box")
        {
          return size.prod();
        }
        else if (geomType == "cylinder")
        {
          return 2 * pi * std::pow(size(0), 2) * size(1);
        }
        else if (geomType == "sphere")
        {
          return 4.0 / 3 * pi * std::pow(size(0), 3);
        }
        else if (geomType == "capsule")
        {
          return 4.0 / 3 * pi * std::pow(size(0), 3) + 2 * pi * std::pow(size(0), 2) * size(1);
        }
        else if (geomType == "ellipsoid")
        {
          return 4.0 / 3 * pi * size.prod();
        }
        else if (geomType == "plane")
        {
          return 0.0;
        }
        else
        {
          throw std::invalid_argument("geometry type does not exist");
        }
      }

      bool isType(const MjcfGeom & geom, const GeometryType & type)
      {
        switch (type)
        {
        case ::pinocchio::COLLISION:
          return geom.geomKind != MjcfGeom::VISUAL;
        default:
        case ::pinocchio::VISUAL:
          return geom.geomKind != MjcfGeom::COLLISION;
        }
      }

      /**
       * @brief      Get a coal::CollisionGeometry from an mjcf geometry, searching
       *             for it in specified package directories
       *
       * @param[in]  geom  A mjcf geometry object
       * @return     A shared pointer on the geometry converted as a coal::CollisionGeometry
       */
#ifdef PINOCCHIO_WITH_COLLISION
      static std::shared_ptr<coal::CollisionGeometry> retrieveCollisionGeometry(
        ::coal::MeshLoaderPtr & meshLoader,
        const MjcfGeom & geom,
        const MjcfGraph & currentGraph,
        std::string & meshPath,
        Eigen::Vector3d & meshScale)
      {
        if (geom.geomType == "mesh")
        {
          if (currentGraph.mapOfMeshes.find(geom.meshName) == currentGraph.mapOfMeshes.end())
          {
            std::stringstream ss;
            ss << "Cannot find mesh " << geom.meshName << " for geometry " << geom.geomName;
            std::string error_msg = ss.str();
            PINOCCHIO_THROW_PRETTY(std::invalid_argument, error_msg);
          }
          MjcfMesh currentMesh = currentGraph.mapOfMeshes.at(geom.meshName);
          if (currentMesh.vertices.size() > 0)
          {
            auto vertices = currentMesh.vertices;
            // Scale vertices
            for (Eigen::Index i = 0; i < vertices.rows(); ++i)
              vertices.row(i) = vertices.row(i).cwiseProduct(currentMesh.scale.transpose());
            auto model = std::make_shared<coal::BVHModel<coal::OBBRSS>>();
            model->beginModel();
            model->addVertices(vertices);
            model->endModel();
            model->buildConvexHull(true, "Qt");
            return model->convex;
          }
          meshPath = currentMesh.filePath;
          meshScale = currentMesh.scale;
          coal::BVHModelPtr_t bvh = meshLoader->load(meshPath, meshScale);
          return bvh;
        }
        else if (geom.geomType == "cylinder")
        {
          meshPath = "CYLINDER";
          meshScale << 1, 1, 1;
          double radius = geom.size(0);
          double height = geom.size(1) * 2;
          return std::make_shared<coal::Cylinder>(radius, height);
        }
        else if (geom.geomType == "box")
        {
          meshPath = "BOX";
          meshScale << 1, 1, 1;
          double x = geom.size(0);
          double y = geom.size(1);
          double z = geom.size(2);
          return std::make_shared<coal::Box>(x, y, z);
        }
        else if (geom.geomType == "sphere")
        {
          meshPath = "SPHERE";
          meshScale << 1, 1, 1;
          double radius = geom.size(0);
          return std::make_shared<coal::Sphere>(radius);
        }
        else if (geom.geomType == "ellipsoid")
        {
          meshPath = "ELLIPSOID";
          meshScale << 1, 1, 1;
          double rx = geom.size(0);
          double ry = geom.size(1);
          double rz = geom.size(2);
          return std::make_shared<coal::Ellipsoid>(rx, ry, rz);
        }
        else if (geom.geomType == "capsule")
        {
          meshPath = "CAPSULE";
          meshScale << 1, 1, 1;
          double radius = geom.size(0);
          double height = geom.size(1) * 2;
          return std::make_shared<coal::Capsule>(radius, height);
        }
        else if (geom.geomType == "plane")
        {
          meshPath = "PLANE";
          meshScale << 1, 1, 1;
          return std::make_shared<coal::Halfspace>(0, 0, 1, 0);
        }
        else
        {
          // Missing plane, hfield and sdf
          std::stringstream ss;
          ss << "Unknown geometry type: " << geom.geomType << "\n";
          PINOCCHIO_THROW(std::invalid_argument, ss.str());
        }
      }
#else
      static std::shared_ptr<coal::CollisionGeometry> retrieveCollisionGeometry(
        ::coal::MeshLoaderPtr &,
        const MjcfGeom &,
        const MjcfGraph &,
        std::string &,
        Eigen::Vector3d &)
      {
        return std::make_shared<coal::CollisionGeometry>();
      }
#endif

      /// @brief Add all geometries associated with a body
      /// @param currentBody Body from which geometries will be collected
      /// @param type Type of geometry to parse (COLLISION or VISUAL)
      /// @param geomModel geometry model to fill
      /// @param meshLoader mesh loader from coal
      static void addLinksToGeomModel(
        const MjcfBody & currentBody,
        const MjcfGraph & currentGraph,
        GeometryModel & geomModel,
        ::coal::MeshLoaderPtr & meshLoader,
        const GeometryType & type,
        bool isWorldBody = false)
      {
        const std::string & bodyName = currentBody.bodyName;
        FrameIndex frame_id;
        Frame frame;
        if (isWorldBody)
        {
          frame_id = 0;
          // "universe" frame should be of index 0
          assert(frame_id == currentGraph.mjcfVisitor.model.getFrameId(bodyName));
          frame = currentGraph.mjcfVisitor.model.frames[frame_id];
        }
        else
        {
          frame_id = currentGraph.mjcfVisitor.getBodyId(bodyName);
          frame = currentGraph.mjcfVisitor.getBodyFrame(bodyName);
        }

        const SE3 & bodyPlacement = frame.placement;

        std::size_t objectId = 0;
        for (const auto & geom : currentBody.geomChildren)
        {
          std::string meshPath = "";
          Eigen::Vector3d meshScale(Eigen::Vector3d::Ones());
          if (isType(geom, type))
          {
            const GeometryObject::CollisionGeometryPtr geometry =
              retrieveCollisionGeometry(meshLoader, geom, currentGraph, meshPath, meshScale);

            bool overrideMaterial = false;
            Eigen::Vector4d meshColor = geom.rgba;
            std::string texturePath;
            if (!geom.materialName.empty())
            {
              if (
                currentGraph.mapOfMaterials.find(geom.materialName)
                == currentGraph.mapOfMaterials.end())
              {
                std::stringstream ss;
                ss << "Cannot find material " << geom.materialName << " for geometry "
                   << geom.geomName;
                std::string error_msg = ss.str();
                PINOCCHIO_THROW_PRETTY(std::invalid_argument, error_msg);
              }
              MjcfMaterial mat = currentGraph.mapOfMaterials.at(geom.materialName);
              meshColor = mat.rgba;
              overrideMaterial = true;
              if (!mat.texture.empty())
              {
                if (
                  currentGraph.mapOfTextures.find(mat.texture) == currentGraph.mapOfTextures.end())
                {
                  std::stringstream ss;
                  ss << "Cannot find texture " << mat.texture << " for material "
                     << geom.materialName << " for geometry " << geom.geomName;
                  std::string error_msg = ss.str();
                  PINOCCHIO_THROW_PRETTY(std::invalid_argument, error_msg);
                }
                MjcfTexture text = currentGraph.mapOfTextures.at(mat.texture);
                texturePath = text.filePath;
              }
            }
            const SE3 geomPlacement = bodyPlacement * geom.geomPlacement;
            std::ostringstream geometry_object_suffix;
            geometry_object_suffix << "_" << objectId;
            std::string geometry_object_name = geom.geomName;
            if (geom.geomName.empty())
            {
              geometry_object_name = std::string(bodyName + geometry_object_suffix.str());
            }

            GeometryObject geometry_object(
              geometry_object_name, frame.parentJoint, frame_id, geomPlacement, geometry, meshPath,
              meshScale, overrideMaterial, meshColor, texturePath);
            geomModel.addGeometryObject(geometry_object);
            ++objectId;
          }
        }
      }

      void MjcfMaterial::goThroughElement(const ptree & el)
      {
        // texture name
        auto text_s = el.get_optional<std::string>("<xmlattr>.texture");
        if (text_s)
          texture = *text_s;

        // Emission
        auto single_v = el.get_optional<float>("<xmlattr>.emission");
        if (single_v)
          emission = *single_v;

        // Specular
        if ((single_v = el.get_optional<float>("<xmlattr>.specular")))
          specular = *single_v;

        // Shininess
        if ((single_v = el.get_optional<float>("<xmlattr>.shininess")))
          shininess = *single_v;

        if ((single_v = el.get_optional<float>("<xmlattr>.reflectance")))
          reflectance = *single_v;

        auto rgba_ = el.get_optional<std::string>("<xmlattr>.rgba");
        if (rgba_)
          rgba = internal::getVectorFromStream<4>(*rgba_);
      }

      void MjcfGeom::goThroughElement(const ptree & el, const MjcfGraph & currentGraph)
      {
        if (el.get_child_optional("<xmlattr>.pos") && el.get_child_optional("<xmlattr>.fromto"))
          throw std::invalid_argument("Both pos and fromto are defined in geom object");

        // Placement
        geomPlacement = currentGraph.convertPosition(el);

        // density
        auto value_v = el.get_optional<double>("<xmlattr>.density");
        if (value_v)
          density = *value_v;

        // contype
        auto val_v = el.get_optional<int>("<xmlattr>.contype");
        if (val_v)
          contype = *val_v;
        // condim
        if ((val_v = el.get_optional<int>("<xmlattr>.conaffinity")))
          conaffinity = *val_v;
        // group
        if ((val_v = el.get_optional<int>("<xmlattr>.group")))
          group = *val_v;

        // mass
        massGeom = el.get_optional<double>("<xmlattr>.mass");

        // shellinertia
        auto value_s = el.get_optional<std::string>("<xmlattr>.shellinertia");
        if (value_s)
          throw std::invalid_argument("Shell inertia computation is not supported yet.");

        // type
        if ((value_s = el.get_optional<std::string>("<xmlattr>.type")))
          geomType = *value_s;

        // material name
        if ((value_s = el.get_optional<std::string>("<xmlattr>.material")))
          materialName = *value_s;

        if (geomType == "mesh")
        {
          value_s = el.get_optional<std::string>("<xmlattr>.mesh");
          if (value_s)
            meshName = *value_s;
        }

        // rgba
        if ((value_s = el.get_optional<std::string>("<xmlattr>.rgba")))
          rgba = internal::getVectorFromStream<4>(*value_s);

        // size
        if ((value_s = el.get_optional<std::string>("<xmlattr>.size")))
          sizeS = *value_s;
        // fromto
        if (!fromtoS)
          fromtoS = el.get_optional<std::string>("<xmlattr>.fromto");
        if (fromtoS)
        {
          Eigen::VectorXd poses = internal::getVectorFromStream<6>(*fromtoS);
          geomPlacement.translation() = (poses.head(3) + poses.tail(3)) / 2;

          Eigen::Vector3d zaxis = poses.tail(3) - poses.head(3);
          // Compute the rotation matrix that maps z_axis to unit z
          geomPlacement.rotation() =
            Eigen::Quaterniond::FromTwoVectors(Eigen::Vector3d::UnitZ(), zaxis).toRotationMatrix();
        }
      }

      void MjcfGeom::findKind()
      {
        // Common mechanism to set visual only geometry
        if (contype == 0 && conaffinity == 0)
          geomKind = VISUAL;
        else if (group > 2) // Mechanism for Collision only geometries
          geomKind = COLLISION;
        else
          geomKind = BOTH;
      }

      void MjcfGeom::computeSize()
      {
        std::vector<std::string> forbiddenListFromTo = {"plane", "hfield", "mesh", "sphere"};
        if (
          fromtoS
          && std::find(forbiddenListFromTo.begin(), forbiddenListFromTo.end(), geomType)
               != forbiddenListFromTo.end())
          throw std::invalid_argument(
            "fromto tag can only be used with capsules, boxes, cylinders and ellipsoids");

        // First deal with cases that do not work with fromto
        if (geomType == "sphere")
          size = internal::getVectorFromStream<1>(sizeS);
        else if (geomType == "plane")
          size = internal::getVectorFromStream<3>(sizeS);
        else if (geomType == "box")
        {
          if (!fromtoS)
            size = internal::getVectorFromStream<3>(sizeS) * 2; // half x, half y, half z
          else
          {
            size = Eigen::Vector3d::Zero();
            size(0) = internal::getVectorFromStream<1>(sizeS)(0) * 2;
            size(1) = size(0);
            Eigen::VectorXd poses = internal::getVectorFromStream<6>(*fromtoS);
            Eigen::Vector3d zaxis = poses.tail(3) - poses.head(3);
            size(2) = zaxis.norm();
          }
        }
        else if (geomType == "ellipsoid")
        {
          if (!fromtoS)
            size = internal::getVectorFromStream<3>(sizeS); // radius and half length
          else
          {
            size = Eigen::Vector3d::Zero();
            size(0) = internal::getVectorFromStream<1>(sizeS)(0);
            size(1) = size(0);
            Eigen::VectorXd poses = internal::getVectorFromStream<6>(*fromtoS);
            Eigen::Vector3d zaxis = poses.tail(3) - poses.head(3);
            size(2) = zaxis.norm() / 2; // to get radius;
          }
        }
        else if (geomType == "cylinder" || geomType == "capsule")
        {
          if (!fromtoS)
            size = internal::getVectorFromStream<2>(sizeS);
          else
          {
            size = Eigen::Vector2d::Zero();
            size(0) = internal::getVectorFromStream<1>(sizeS)(0);
            Eigen::VectorXd poses = internal::getVectorFromStream<6>(*fromtoS);
            Eigen::Vector3d zaxis = poses.tail(3) - poses.head(3);
            size(1) = zaxis.norm() / 2; // to get half height
          }
        }
        else if (geomType == "mesh" || geomType == "plane")
          return;
        else
          throw std::invalid_argument("geomType does not exist");
      }

      void MjcfGeom::computeInertia()
      {
        if (geomType == "mesh" || geomType == "plane" || geomType == "hfield")
          return;

        double mass;
        if (massGeom)
          mass = *massGeom;
        else
          mass = computeVolume(size, geomType) * density;

        if (geomType == "box")
        {
          geomInertia = Inertia::FromBox(mass, size(0), size(1), size(2));
        }
        else if (geomType == "cylinder")
        {
          geomInertia = Inertia::FromCylinder(mass, size(0), size(1) * 2);
        }
        else if (geomType == "ellipsoid")
        {
          geomInertia = Inertia::FromEllipsoid(mass, size(0), size(1), size(2));
        }
        else if (geomType == "sphere")
        {
          geomInertia = Inertia::FromSphere(mass, size(0));
        }
        else if (geomType == "capsule")
        {
          geomInertia = Inertia::FromCapsule(mass, size(0), size(1) * 2);
        }
        else
        {
          throw std::invalid_argument("Unsupported geometry type");
        }
      }

      void
      MjcfGeom::fill(const ptree & el, const MjcfBody & currentBody, const MjcfGraph & currentGraph)
      {
        // Name
        auto name_s = el.get_optional<std::string>("<xmlattr>.name");
        if (name_s)
          geomName = *name_s;
        else
          geomName =
            currentBody.bodyName + "Geom_" + std::to_string(currentBody.geomChildren.size());

        // default < ChildClass < Class < Real Joint
        if (currentGraph.mapOfClasses.find("mujoco_default") != currentGraph.mapOfClasses.end())
        {
          const MjcfClass & classD = currentGraph.mapOfClasses.at("mujoco_default");
          if (auto geom_p = classD.classElement.get_child_optional("geom"))
            goThroughElement(*geom_p, currentGraph);
        }
        //  childClass
        if (currentBody.childClass != "")
        {
          const MjcfClass & classE = currentGraph.mapOfClasses.at(currentBody.childClass);
          if (auto geom_p = classE.classElement.get_child_optional("geom"))
            goThroughElement(*geom_p, currentGraph);
        }

        // Class
        auto cl_s = el.get_optional<std::string>("<xmlattr>.class");
        if (cl_s)
        {
          std::string className = *cl_s;
          const MjcfClass & classE = currentGraph.mapOfClasses.at(className);
          if (auto geom_p = classE.classElement.get_child_optional("geom"))
            goThroughElement(*geom_p, currentGraph);
        }

        // Geom
        goThroughElement(el, currentGraph);
        if (geomType == "mesh" && meshName.empty())
          throw std::invalid_argument("Type is mesh but no mesh file were provided");
        // Compute Kind of geometry
        findKind();

        computeSize();

        // Compute Mass and inertia of geom object
        computeInertia();
      }

      void MjcfSite::goThroughElement(const ptree & el, const MjcfGraph & currentGraph)
      {
        if (el.get_child_optional("<xmlattr>.pos") && el.get_child_optional("<xmlattr>.fromto"))
          throw std::invalid_argument("Both pos and fromto are defined in site object");

        // Placement
        sitePlacement = currentGraph.convertPosition(el);

        auto fromtoS = el.get_optional<std::string>("<xmlattr>.fromto");
        if (fromtoS)
        {
          Eigen::VectorXd poses = internal::getVectorFromStream<6>(*fromtoS);
          sitePlacement.translation() = (poses.head(3) + poses.tail(3)) / 2;

          Eigen::Vector3d zaxis = poses.tail(3) - poses.head(3);
          // Compute the rotation matrix that maps z_axis to unit z
          sitePlacement.rotation() =
            Eigen::Quaterniond::FromTwoVectors(Eigen::Vector3d::UnitZ(), zaxis).toRotationMatrix();
        }
      }

      void
      MjcfSite::fill(const ptree & el, const MjcfBody & currentBody, const MjcfGraph & currentGraph)
      {
        // Name
        auto name_s = el.get_optional<std::string>("<xmlattr>.name");
        if (name_s)
          siteName = *name_s;
        else
          siteName =
            currentBody.bodyName + "Site_" + std::to_string(currentBody.siteChildren.size());

        // default < ChildClass < Class < Real Joint
        if (currentGraph.mapOfClasses.find("mujoco_default") != currentGraph.mapOfClasses.end())
        {
          const MjcfClass & classD = currentGraph.mapOfClasses.at("mujoco_default");
          if (auto site_p = classD.classElement.get_child_optional("site"))
            goThroughElement(*site_p, currentGraph);
        }
        //  childClass
        if (currentBody.childClass != "")
        {
          const MjcfClass & classE = currentGraph.mapOfClasses.at(currentBody.childClass);
          if (auto site_p = classE.classElement.get_child_optional("site"))
            goThroughElement(*site_p, currentGraph);
        }

        // Class
        auto cl_s = el.get_optional<std::string>("<xmlattr>.class");
        if (cl_s)
        {
          std::string className = *cl_s;
          const MjcfClass & classE = currentGraph.mapOfClasses.at(className);
          if (auto site_p = classE.classElement.get_child_optional("site"))
            goThroughElement(*site_p, currentGraph);
        }

        // Site
        goThroughElement(el, currentGraph);
      }

      void MjcfGraph::parseGeomTree(
        const GeometryType & type, GeometryModel & geomModel, ::coal::MeshLoaderPtr & meshLoader)
      {
#ifdef PINOCCHIO_WITH_COLLISION
        if (!meshLoader)
          meshLoader = std::make_shared<coal::MeshLoader>(coal::MeshLoader());
#endif // ifdef PINOCCHIO_WITH_COLLISION

        addLinksToGeomModel(worldBody, *this, geomModel, meshLoader, type, true);
        for (const auto & entry : bodiesList)
        {
          const MjcfBody & currentBody = mapOfBodies.at(entry);

          addLinksToGeomModel(currentBody, *this, geomModel, meshLoader, type);
        }
      }
    } // namespace details
  } // namespace mjcf
} // namespace pinocchio
