ObjectFinder.h

#pragma once

#include <filesystem>
#include <unordered_map>

#include <VirtualRobot/VirtualRobot.h>

#include <ArmarXCore/core/logging/Logging.h>

#include "ObjectInfo.h"
#include "ObjectPose.h"

#include <RobotAPI/libraries/armem_objects/types.h>

namespace armarx
{
    /**
     * @brief Used to find objects in the ArmarX objects repository [1].
     *
     * @see [1] https://gitlab.com/ArmarX/ArmarXObjects
     */
    class ObjectFinder : Logging
    {
    public:
        using path = std::filesystem::path;
        inline static const std::string DefaultObjectsPackageName = "ArmarXObjects";

    public:
        ObjectFinder(const std::string& objectsPackageName = DefaultObjectsPackageName);

        ObjectFinder(ObjectFinder&&)                 = default;
        ObjectFinder(const ObjectFinder&)            = default;
        ObjectFinder& operator=(ObjectFinder&&)      = default;
        ObjectFinder& operator=(const ObjectFinder&) = default;

        void setPath(const std::string& path);

        std::string getPackageName() const;

        std::optional<ObjectInfo> findObject(const std::string& dataset, const std::string& name) const;
        std::optional<ObjectInfo> findObject(const std::string& nameOrID) const;
        std::optional<ObjectInfo> findObject(const ObjectID& id) const;
        std::optional<ObjectInfo> findObject(const objpose::ObjectPose& obj) const;

        std::vector<std::string> getDatasets() const;
        std::vector<path> getDatasetDirectories() const;

        std::vector<ObjectInfo> findAllObjects(bool checkPaths = true) const;
        std::map<std::string, std::vector<ObjectInfo>> findAllObjectsByDataset(bool checkPaths = true) const;
        std::vector<ObjectInfo> findAllObjectsOfDataset(const std::string& dataset, bool checkPaths = true) const;

        std::vector<armem::articulated_object::ArticulatedObjectDescription> findAllArticulatedObjects(bool checkPaths) const;
        std::vector<armem::articulated_object::ArticulatedObjectDescription> findAllArticulatedObjectsOfDataset(const std::string& dataset, bool checkPaths) const;
        std::unordered_map<std::string, std::vector<armem::articulated_object::ArticulatedObjectDescription>> findAllArticulatedObjectsByDataset(bool checkPaths = true) const;

        VirtualRobot::ManipulationObjectPtr
        static loadManipulationObject(const std::optional<ObjectInfo>& ts);
        VirtualRobot::ManipulationObjectPtr
        loadManipulationObject(const objpose::ObjectPose& obj) const;

        VirtualRobot::ObstaclePtr
        static loadObstacle(const std::optional<ObjectInfo>& ts);
        VirtualRobot::ObstaclePtr
        loadObstacle(const objpose::ObjectPose& obj) const;


        /**
         * @brief Load names to use when matched when recognizing an object by name.
         *
         * If the object's names JSON file does not exist, no names will be added from a file.
         * If you would like to detect this case, first `findObject()`, then use
         * `ObjectInfo::loadRecognizedNames()`, which returns a `std::optional`.
         *
         * @param includeClassName If true, include the raw class name in the result.
         * @see `ObjectInfo::loadRecognizedNames()`
         */
        std::vector<std::string> loadRecognizedNames(const ObjectID& objectID, bool includeClassName = false) const;
        /**
         * @brief Load names to use when verbalizing an object name.
         *
         * If the object's names JSON file does not exist, no names will be added from a file.
         * If you would like to detect this case, first `findObject()`, then use
         * `ObjectInfo::loadSpokenNames()`, which returns a `std::optional`.
         *
         * @param includeClassName If true, include the raw class name in the result.
         * @see `ObjectInfo::loadSpokenNames()`
         */
        std::vector<std::string> loadSpokenNames(const ObjectID& objectID, bool includeClassName = false) const;


    private:

        void init() const;
        bool isDatasetDirValid(const std::filesystem::path& path) const;

        path _rootDirAbs() const;
        path _rootDirRel() const;

        bool _ready() const;


    private:

        /// Name of package containing the object models (ArmarXObjects by default).
        mutable std::string packageName;

        /**
         * @brief Absolute path to data directory (e.g. "/.../repos/ArmarXObjects/data").
         * Empty if package could not be found.
         */
        mutable path packageDataDir;

    };
}