/*
 * This file is part of ArmarX.
 *
 * ArmarX is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 *
 * ArmarX is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 *
 * @package    RobotAPI::ArmarXObjects::RobotSensorMemory
 * @author     Rainer Kartmann ( rainer dot kartmann at kit dot edu )
 * @date       2020
 * @copyright  http://www.gnu.org/licenses/gpl-2.0.txt
 *             GNU General Public License
 */

#include "RobotStateWriter.h"

// STL
#include <chrono>

// Simox
#include <SimoxUtility/math/convert/rpy_to_mat3f.h>

// ArmarX
#include <ArmarXCore/core/exceptions/local/ExpressionException.h>
#include <ArmarXCore/core/time/CycleUtil.h>
#include <ArmarXCore/libraries/ArmarXCoreComponentPlugins/DebugObserverComponentPlugin.h>

#include <RobotAPI/libraries/core/FramedPose.h>
#include <RobotAPI/libraries/aron/core/navigator/data/AllNavigators.h>
#include <RobotAPI/libraries/armem/core/Time.h>
#include <RobotAPI/libraries/armem/core/ice_conversions.h>
#include <RobotAPI/libraries/armem/server/MemoryToIceAdapter.h>
#include <RobotAPI/libraries/armem_robot_state/aron/Transform.aron.generated.h>
#include <RobotAPI/libraries/armem_robot_state/server/localization/Segment.h>


namespace armarx::armem::server::robot_state
{

    void RobotStateWriter::connect(armarx::plugins::DebugObserverComponentPlugin& debugObserver)
    {
        {
            // Thread-local copy of debug observer helper.
            this->debugObserver = DebugObserverHelper(
                                      Logging::tag.tagName, debugObserver.getDebugObserver(), true
                                  );
        }
    }


    static float toDurationMs(
        std::chrono::high_resolution_clock::time_point start,
        std::chrono::high_resolution_clock::time_point end)
    {
        auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
        return duration.count() / 1000.f;
    }


    void RobotStateWriter::run(float pollFrequency,
                               std::queue<RobotUnitData>& dataQueue, std::mutex& dataMutex,
                               MemoryToIceAdapter& memory, localization::Segment& localizationSegment)
    {
        CycleUtil cycle(static_cast<int>(1000.f / pollFrequency));
        while (task and not task->isStopped())
        {
            size_t queueSize = 0;
            std::queue<RobotUnitData> batch;
            {
                std::lock_guard lock{dataMutex};
                queueSize = dataQueue.size();
                if (dataQueue.size() >= properties.memoryBatchSize)
                {
                    std::swap(batch, dataQueue);
                }
            }
            if (debugObserver)
            {
                debugObserver->setDebugObserverDatafield("RobotStateWriter | Queue Size", queueSize);
            }
            if (batch.size() > 0)
            {
                Update update = buildUpdate(batch);

                auto start = std::chrono::high_resolution_clock::now();
                auto endProprioception = start, endLocalization = start;

                armem::CommitResult result;
                {
                    // Commits lock the core segments.

                    result = memory.commitLocking(update.proprioception);
                    endProprioception = std::chrono::high_resolution_clock::now();

                    localizationSegment.commitTransformLocking(update.localization);
                    endLocalization = std::chrono::high_resolution_clock::now();
                }
                if (not result.allSuccess())
                {
                    ARMARX_WARNING << "Could not add data to memory. Error message: " << result.allErrorMessages();
                }
                if (debugObserver)
                {
                    auto end = std::chrono::high_resolution_clock::now();

                    debugObserver->setDebugObserverDatafield("RobotStateWriter | t Commit [ms]", toDurationMs(start, end));
                    debugObserver->setDebugObserverDatafield("RobotStateWriter | t Commit Proprioception [ms]", toDurationMs(start, endProprioception));
                    debugObserver->setDebugObserverDatafield("RobotStateWriter | t Commit Localization [ms]", toDurationMs(endProprioception, endLocalization));
                }
            }

            if (debugObserver)
            {
                debugObserver->sendDebugObserverBatch();
            }
            cycle.waitForCycleDuration();
        }
    }


    RobotStateWriter::Update RobotStateWriter::buildUpdate(std::queue<RobotUnitData> dataQueue)
    {
        ARMARX_CHECK_GREATER(dataQueue.size(), 0);

        ARMARX_DEBUG << "RobotStateWriter: Commit batch of " << dataQueue.size() << " timesteps to memory...";
        auto start = std::chrono::high_resolution_clock::now();

        // Send batch to memory
        Update update;

        while (dataQueue.size() > 0)
        {
            const RobotUnitData& convertedAndGroupedData = dataQueue.front();

            for (const auto& [encName, encTimestep] : convertedAndGroupedData.groups)
            {
                ARMARX_CHECK_EQUAL(encName, encTimestep.name);

                armem::EntityUpdate& up = update.proprioception.add();
                up.entityID = properties.robotUnitProviderID.withEntityName(encName);
                up.timeCreated = encTimestep.timestamp;
                up.instancesData = { encTimestep.data };
            }

            // odometry pose -> localization segment
            auto it = convertedAndGroupedData.groups.find(properties.platformGroupName);
            if (it != convertedAndGroupedData.groups.end())
            {
                ARMARX_DEBUG << "Found odometry data.";

                const RobotUnitData::RobotUnitDataGroup& timestep = it->second;
                const float relPosX = aron::datanavigator::FloatNavigator::DynamicCast(timestep.data->getElement("relativePositionX"))->getValue();
                const float relPosY = aron::datanavigator::FloatNavigator::DynamicCast(timestep.data->getElement("relativePositionY"))->getValue();
                const float relOrientation = aron::datanavigator::FloatNavigator::DynamicCast(timestep.data->getElement("relativePositionRotation"))->getValue();

                Eigen::Affine3f odometryPose = Eigen::Affine3f::Identity();
                odometryPose.translation() << relPosX, relPosY, 0; // TODO set height
                odometryPose.linear() = simox::math::rpy_to_mat3f(0.F, 0.F, relOrientation);

                // const auto timestamp = armem::Time::microSeconds(it->second.timestamp);

                armem::robot_state::Transform& transform = update.localization;
                transform.header.parentFrame = armarx::OdometryFrame;
                transform.header.frame = "root"; // TODO: robot root node
                transform.header.agent = properties.robotUnitProviderID.providerSegmentName;
                transform.header.timestamp = it->second.timestamp;
                transform.transform = odometryPose;
            }
            else if (!noOdometryDataLogged)
            {
                ARMARX_INFO << "No odometry data! (No group with name '" << properties.platformGroupName << "'.). \n"
                            << "If you are using a mobile platform this should not have happened.";
                noOdometryDataLogged = true;
            }

            dataQueue.pop();
        }

        auto stop = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::microseconds>(stop - start);
        ARMARX_DEBUG << "RobotUnitData: The total runtime of sending a batch to the memory is: " << duration;
        if (debugObserver)
        {
            debugObserver->setDebugObserverDatafield("RobotStateWriter | t Build Update [ms]", duration.count() / 1000.f);
        }

        return update;
    }

}