#include <algorithm>
#include <atomic>
#include <chrono>
#include <cmath>
#include <csignal>
#include <fstream>
#include <iostream>
#include <memory>
#include <regex>
#include <sstream>
#include <string>
#include <thread>
#include <vector>

#include <Eigen/Core>

#include <unitree/robot/channel/channel_factory.hpp>
#include <unitree/dds_wrapper/robots/g1/g1.h>
#include <unitree/dds_wrapper/robots/r1/r1.h>

namespace {

std::atomic<bool> running{true};

void signalHandler(int)
{
    running = false;
}

struct PersonDetection {
    bool valid = false;
    float confidence = 0.0f;
    float x = 0.5f;
    float y = 0.5f;
};

std::string readFile(const std::string& path)
{
    std::ifstream file(path);

    if (!file)
        return {};

    std::ostringstream out;
    out << file.rdbuf();

    return out.str();
}

PersonDetection readBestPerson(
    const std::string& path)
{
    const std::string json =
        readFile(path);

    PersonDetection best;

    if (json.empty())
        return best;

    /*
     * Matches our yolo_detector.py output:
     *
     * "class": "person",
     * "confidence": 0.95,
     * ...
     * "normalized_center": {
     *     "x": 0.51,
     *     "y": 0.43
     * }
     */
    static const std::regex personRegex(
        R"REGEX(
"class"\s*:\s*"person"
[\s\S]*?
"confidence"\s*:\s*([0-9.]+)
[\s\S]*?
"normalized_center"\s*:\s*\{
\s*"x"\s*:\s*([0-9.]+)
\s*,\s*
"y"\s*:\s*([0-9.]+)
)REGEX",
        std::regex::ECMAScript
    );

    auto begin =
        std::sregex_iterator(
            json.begin(),
            json.end(),
            personRegex
        );

    auto end =
        std::sregex_iterator();

    for (auto it = begin;
         it != end;
         ++it) {

        try {
            const float confidence =
                std::stof((*it)[1].str());

            const float x =
                std::stof((*it)[2].str());

            const float y =
                std::stof((*it)[3].str());

            if (!best.valid ||
                confidence > best.confidence) {

                best.valid = true;
                best.confidence = confidence;
                best.x = x;
                best.y = y;
            }
        }
        catch (...) {
        }
    }

    return best;
}

class HeadTracker {
public:
    explicit HeadTracker(
        const std::string& interface)
    {
        unitree::robot::ChannelFactory::
            Instance()->Init(
                0,
                interface
            );

        lowstate_ =
            std::make_shared<
                unitree::robot::g1::subscription::LowState
            >();

        std::cout
            << "[MOTION] Waiting for lowstate...\n";

        lowstate_->wait_for_connection();

        armsdk_ =
            std::make_unique<
                unitree::robot::r1::publisher::ArmSdk
            >();

        std::cout
            << "[MOTION] Connected to R1\n";

        initialize();
    }

    ~HeadTracker()
    {
        release();
    }

    void run()
    {
        constexpr char DETECTIONS[] =
            "/dev/shm/r1_detections.json";

        while (running) {

            const PersonDetection person =
                readBestPerson(
                    DETECTIONS
                );

            if (person.valid &&
                person.confidence >= 0.50f) {

                track(person);
            }

            publish();

            std::this_thread::sleep_for(
                std::chrono::milliseconds(20)
            );
        }
    }

private:
    std::shared_ptr<
        unitree::robot::g1::subscription::LowState
    > lowstate_;

    std::unique_ptr<
        unitree::robot::r1::publisher::ArmSdk
    > armsdk_;

    Eigen::VectorXf q_des_;

    /*
     * ArmSdk::JOINTS ordering on R1:
     *
     *  0-4   left arm
     *  5-9   right arm
     *  10    waist yaw
     *  11    head pitch
     *  12    head yaw
     */
    static constexpr int HEAD_PITCH = 11;
    static constexpr int HEAD_YAW   = 12;

    /*
     * Start conservatively.
     *
     * We can widen these after verifying direction
     * and mechanical range on your R1.
     */
    static constexpr float MAX_YAW =
        0.70f;

    static constexpr float MAX_PITCH =
        0.25f;

    /*
     * Camera control.
     */
    static constexpr float DEAD_ZONE_X =
        0.07f;

    static constexpr float DEAD_ZONE_Y =
        0.08f;

    /*
     * Maximum movement per 20 ms update.
     *
     * 0.012 rad * 50 Hz ~= 0.6 rad/s.
     */
    static constexpr float MAX_YAW_STEP =
        0.012f;

    static constexpr float MAX_PITCH_STEP =
        0.008f;

    /*
     * If the head moves AWAY from the person,
     * flip one of these signs.
     */
    static constexpr float YAW_SIGN =
        -1.0f;

    static constexpr float PITCH_SIGN =
        1.0f;

    const std::vector<float> kp_ = {
        50, 50, 40, 40, 30,
        50, 50, 40, 40, 30,
        50,
        15,
        15
    };

    const std::vector<float> kd_ = {
        2, 2, 2, 2, 2,
        2, 2, 2, 2, 2,
        3,
        1,
        1
    };

    void initialize()
    {
        armsdk_->weight(1.0f);

        q_des_.resize(
            armsdk_->JOINTS.size()
        );

        for (size_t i = 0;
             i < armsdk_->JOINTS.size();
             ++i) {

            const int joint =
                static_cast<int>(
                    armsdk_->JOINTS[i]
                );

            const float current =
                lowstate_
                    ->msg_
                    .motor_state()
                    .at(joint)
                    .q();

            q_des_[i] = current;

            auto& command =
                armsdk_
                    ->msg_
                    .motor_cmd()
                    .at(joint);

            command.q(current);
            command.dq(0.0f);
            command.tau(0.0f);

            command.kp(
                kp_.at(i)
            );

            command.kd(
                kd_.at(i)
            );
        }

        armsdk_->unlockAndPublish();

        std::cout
            << "[MOTION] ArmSdk enabled\n"
            << "[MOTION] Initial head pitch: "
            << q_des_[HEAD_PITCH]
            << " rad\n"
            << "[MOTION] Initial head yaw: "
            << q_des_[HEAD_YAW]
            << " rad\n";
    }

    void track(
        const PersonDetection& person)
    {
        float error_x =
            person.x - 0.5f;

        float error_y =
            person.y - 0.5f;

        if (std::abs(error_x) <
            DEAD_ZONE_X) {

            error_x = 0.0f;
        }

        if (std::abs(error_y) <
            DEAD_ZONE_Y) {

            error_y = 0.0f;
        }

        const float yaw_step =
            YAW_SIGN *
            std::clamp(
                error_x * 0.05f,
                -MAX_YAW_STEP,
                MAX_YAW_STEP
            );

        const float pitch_step =
            PITCH_SIGN *
            std::clamp(
                error_y * 0.04f,
                -MAX_PITCH_STEP,
                MAX_PITCH_STEP
            );

        q_des_[HEAD_YAW] =
            std::clamp(
                q_des_[HEAD_YAW] +
                    yaw_step,

                -MAX_YAW,
                MAX_YAW
            );

        q_des_[HEAD_PITCH] =
            std::clamp(
                q_des_[HEAD_PITCH] +
                    pitch_step,

                -MAX_PITCH,
                MAX_PITCH
            );
    }

    void publish()
    {
        for (size_t i = 0;
             i < armsdk_->JOINTS.size();
             ++i) {

            const int joint =
                static_cast<int>(
                    armsdk_->JOINTS[i]
                );

            armsdk_
                ->msg_
                .motor_cmd()
                .at(joint)
                .q(q_des_[i]);
        }

        armsdk_->unlockAndPublish();
    }

    void release()
    {
        if (!armsdk_)
            return;

        std::cout
            << "[MOTION] Releasing ArmSdk...\n";

        for (int step = 20;
             step >= 0;
             --step) {

            const float weight =
                static_cast<float>(step) /
                20.0f;

            armsdk_->weight(weight);
            armsdk_->unlockAndPublish();

            std::this_thread::sleep_for(
                std::chrono::milliseconds(25)
            );
        }

        std::cout
            << "[MOTION] Released\n";
    }
};

} // namespace


int main(
    int argc,
    char** argv)
{
    if (argc < 2) {

        std::cerr
            << "Usage: "
            << argv[0]
            << " <network_interface>\n";

        return 1;
    }

    std::signal(
        SIGINT,
        signalHandler
    );

    std::signal(
        SIGTERM,
        signalHandler
    );

    try {

        HeadTracker tracker(
            argv[1]
        );

        std::cout
            << "[MOTION] Person tracking active\n"
            << "[MOTION] Ctrl+C to stop\n";

        tracker.run();
    }
    catch (const std::exception& e) {

        std::cerr
            << "[MOTION] Error: "
            << e.what()
            << '\n';

        return 1;
    }

    return 0;
}
