#include "KnownGraspProviderSegment.h"

#include <RobotAPI/libraries/armem/server/MemoryToIceAdapter.h>
#include <ArmarXCore/core/rapidxml/wrapper/RapidXmlReader.h>

namespace armarx::armem::grasping::segment
{
    KnownGraspProviderSegment::KnownGraspProviderSegment(armem::server::MemoryToIceAdapter& iceMemory) :
        Base(iceMemory, PROVIDER_SEGMENT_NAME, CORE_SEGMENT_NAME)
    {

    }

    void KnownGraspProviderSegment::init()
    {
        Base::init();

        loadMemory();
    }

    std::optional<arondto::KnownGraspInfo> KnownGraspProviderSegment::knownGraspInfoFromObjectInfo(const ObjectInfo& info)
    {
        std::string objectClassName = info.className();
        auto fileLocInfo = info.file(".xml", "_Grasps");
        std::filesystem::path graspFilePath = fileLocInfo.absolutePath;

        if (std::filesystem::is_regular_file(graspFilePath))
        {
            auto reader = RapidXmlReader::FromFile(graspFilePath);
            RapidXmlReaderNode root = reader->getRoot();

            std::string objectClassNameXML = root.attribute_value("name");
            ARMARX_CHECK_EQUAL(objectClassName, objectClassNameXML);

            arondto::KnownGraspInfo ret;
            ret.correspondingObject.memoryName = "Object";
            ret.correspondingObject.coreSegmentName = "Class";
            ret.correspondingObject.providerSegmentName = "PriorKnowledgeData";
            ret.correspondingObject.entityName = info.idStr();
            ret.xml.package = fileLocInfo.package;
            ret.xml.path = fileLocInfo.relativePath;

            for (const auto& graspSetNode : root.nodes())
            {
                if (graspSetNode.name() != "GraspSet")
                {
                    continue;
                }

                arondto::KnownGraspSet retGraspSet;

                retGraspSet.name = graspSetNode.attribute_value("name");
                retGraspSet.robot = graspSetNode.attribute_value("RobotType");
                retGraspSet.endeffector = graspSetNode.attribute_value("EndEffector");

                for (const auto& graspNode : graspSetNode.nodes())
                {
                    if (graspNode.name() != "Grasp")
                    {
                        continue;
                    }

                    arondto::KnownGrasp retGrasp;

                    retGrasp.name = graspNode.attribute_value("name");
                    retGrasp.quality = std::stof(graspNode.attribute_value("quality"));
                    retGrasp.creator = graspNode.attribute_value("Creation");

                    ARMARX_CHECK(graspNode.nodes().size() == 2 or graspNode.nodes().size() == 1);
                    RapidXmlReaderNode transformNode = graspNode.nodes()[0];

                    ARMARX_CHECK_EQUAL(transformNode.name(), "Transform");
                    ARMARX_CHECK_EQUAL(transformNode.nodes().size(), 1);
                    RapidXmlReaderNode matrixNode = transformNode.nodes()[0];

                    ARMARX_CHECK_EQUAL(matrixNode.nodes().size(), 4);
                    RapidXmlReaderNode row0 = matrixNode.nodes()[0];
                    RapidXmlReaderNode row1 = matrixNode.nodes()[1];
                    RapidXmlReaderNode row2 = matrixNode.nodes()[2];
                    RapidXmlReaderNode row3 = matrixNode.nodes()[3];

                    retGrasp.pose(0, 0) = std::stof(row0.attribute_value("c1"));
                    retGrasp.pose(0, 1) = std::stof(row0.attribute_value("c2"));
                    retGrasp.pose(0, 2) = std::stof(row0.attribute_value("c3"));
                    retGrasp.pose(0, 3) = std::stof(row0.attribute_value("c4"));
                    retGrasp.pose(1, 0) = std::stof(row1.attribute_value("c1"));
                    retGrasp.pose(1, 1) = std::stof(row1.attribute_value("c2"));
                    retGrasp.pose(1, 2) = std::stof(row1.attribute_value("c3"));
                    retGrasp.pose(1, 3) = std::stof(row1.attribute_value("c4"));
                    retGrasp.pose(2, 0) = std::stof(row2.attribute_value("c1"));
                    retGrasp.pose(2, 1) = std::stof(row2.attribute_value("c2"));
                    retGrasp.pose(2, 2) = std::stof(row2.attribute_value("c3"));
                    retGrasp.pose(2, 3) = std::stof(row2.attribute_value("c4"));
                    retGrasp.pose(3, 0) = std::stof(row3.attribute_value("c1"));
                    retGrasp.pose(3, 1) = std::stof(row3.attribute_value("c2"));
                    retGrasp.pose(3, 2) = std::stof(row3.attribute_value("c3"));
                    retGrasp.pose(3, 3) = std::stof(row3.attribute_value("c4"));

                    ARMARX_VERBOSE << "Found grasp '" << retGrasp.name << "' in set '" << retGraspSet.name << "' for obj '" << objectClassName << "' with pose \n" << retGrasp.pose;

                    retGraspSet.grasps.push_back(retGrasp);
                }
                ret.graspSets[retGraspSet.name] = retGraspSet;
            }
            return ret;
        }
        return std::nullopt;
    }

    void KnownGraspProviderSegment::loadMemory()
    {
        // load data from prior knowledge
        ObjectFinder objectFinder;
        const auto now = armem::Time::Now();

        const bool checkPaths = false;
        std::vector<ObjectInfo> infos = objectFinder.findAllObjects(checkPaths);

        const MemoryID providerID = segmentPtr->id().withProviderSegmentName(objectFinder.getPackageName());
        ARMARX_INFO << "Checking up to " << infos.size() << " object classes from '" << objectFinder.getPackageName() << "' ...";

        Commit commit;
        for (ObjectInfo& info : infos)
        {
            info.setLogError(false);
            if (auto knownGraspCandidate = knownGraspInfoFromObjectInfo(info); knownGraspCandidate)
            {
                EntityUpdate& update = commit.add();
                update.entityID = providerID.withEntityName(info.id().str());
                update.entityID.timestamp = update.timeArrived = update.timeCreated = update.timeSent = now;

                update.instancesData =
                {
                    knownGraspCandidate->toAron()
                };
            }
        }
        ARMARX_INFO << "Loaded " << commit.updates.size() << " grasp candidates from object classes from '" << objectFinder.getPackageName() << "'.";
        auto result = iceMemory.commit(commit);
        if (!result.allSuccess())
        {
            ARMARX_WARNING << "Got errors for commit: " << result.allErrorMessages();
        }
    }
}