#include <array>
#include <iostream>
#include <string>
#include <vector>
#include <unistd.h>

#include <opencv2/opencv.hpp>
#include <pinocchio/parsers/urdf.hpp>

#include <unitree/robot/go2/video/video_client.hpp>

#include "robot_common.hh"
#include "r1_vision/pinocchio_kinematics.hh"

using namespace r1_vision;

struct ClickState {
    bool clicked = false;
    cv::Point point{};
};

static void MouseCallback(int event, int x, int y, int, void* user) {
    auto* s = static_cast<ClickState*>(user);
    if (event == cv::EVENT_LBUTTONDOWN) {
        s->clicked = true;
        s->point = cv::Point(x, y);
        std::cout << "[CLICK] " << x << ", " << y << std::endl;
    }
}

static cv::Mat GetFrame(unitree::robot::go2::VideoClient& camera) {
    std::vector<uint8_t> jpeg;
    if (camera.GetImageSample(jpeg) != 0 || jpeg.empty()) return {};
    cv::Mat encoded(1, static_cast<int>(jpeg.size()), CV_8UC1, jpeg.data());
    return cv::imdecode(encoded, cv::IMREAD_COLOR);
}

int main(int argc, char** argv) {
    if (argc < 3) {
        std::cerr << "Usage: r1_hand_calibrate <network_interface> <camera_intrinsics.yaml>\n";
        return 1;
    }

    cv::FileStorage fs(argv[2], cv::FileStorage::READ);
    if (!fs.isOpened()) {
        std::cerr << "Cannot open intrinsics: " << argv[2] << std::endl;
        return 1;
    }
    cv::Mat K, D;
    fs["camera_matrix"] >> K;
    fs["distortion"] >> D;
    fs.release();

    if (K.empty() || D.empty()) {
        std::cerr << "Invalid camera calibration file." << std::endl;
        return 1;
    }

    pinocchio::Model model;
    pinocchio::urdf::buildModel(
        "/home/jvizner/unitree_ros/robots/r1_description/R1.urdf",
        model);
    pinocchio::Data data(model);

    RobotController robot(argv[1]);

    unitree::robot::go2::VideoClient camera;
    camera.SetTimeout(1.0f);
    camera.Init();

    const std::string window = "R1 Wrist Extrinsic Calibration";
    cv::namedWindow(window, cv::WINDOW_NORMAL);
    ClickState click;
    cv::setMouseCallback(window, MouseCallback, &click);

    std::vector<cv::Point3f> robot_points;
    std::vector<cv::Point2f> image_points;

    std::cout << "\nControls:\n"
              << "  left click = mark LEFT wrist\n"
              << "  c = capture point\n"
              << "  p/P = left shoulder pitch +/- 0.08 rad\n"
              << "  r/R = left shoulder roll +/- 0.08 rad\n"
              << "  y/Y = left shoulder yaw +/- 0.08 rad\n"
              << "  e/E = left elbow +/- 0.08 rad\n"
              << "  q = solve and quit\n";

    while (true) {
        cv::Mat frame = GetFrame(camera);
        if (frame.empty()) {
            usleep(10000);
            continue;
        }

        const auto q = robot.CurrentQ();
        const Eigen::Vector3d wrist = FramePosition(
            model, data, q, "left_wrist_roll_link");

        cv::Mat display = frame.clone();
        cv::putText(display,
                    "samples=" + std::to_string(image_points.size()),
                    {20, 35}, cv::FONT_HERSHEY_SIMPLEX, 0.8,
                    {0, 255, 0}, 2);

        cv::putText(display,
                    "wrist R1: " + std::to_string(wrist.x()) + ", " +
                    std::to_string(wrist.y()) + ", " +
                    std::to_string(wrist.z()),
                    {20, 70}, cv::FONT_HERSHEY_SIMPLEX, 0.55,
                    {255, 255, 0}, 2);

        if (click.clicked)
            cv::circle(display, click.point, 7, {0, 0, 255}, -1);

        cv::imshow(window, display);
        const int key = cv::waitKey(10) & 0xff;

        if (key == 'q') break;

        if (key == 'c') {
            if (!click.clicked) {
                std::cout << "[WARN] Click the wrist first." << std::endl;
                continue;
            }
            image_points.emplace_back(
                static_cast<float>(click.point.x),
                static_cast<float>(click.point.y));
            robot_points.emplace_back(
                static_cast<float>(wrist.x()),
                static_cast<float>(wrist.y()),
                static_cast<float>(wrist.z()));
            std::cout << "[CAPTURE] #" << image_points.size()
                      << " pixel=" << click.point.x << "," << click.point.y
                      << " robot=" << wrist.transpose() << std::endl;
            click.clicked = false;
        }

        constexpr float step = 0.08f;
        switch (key) {
            case 'p': robot.MoveJointRelative(0, +step); break;
            case 'P': robot.MoveJointRelative(0, -step); break;
            case 'r': robot.MoveJointRelative(1, +step); break;
            case 'R': robot.MoveJointRelative(1, -step); break;
            case 'y': robot.MoveJointRelative(2, +step); break;
            case 'Y': robot.MoveJointRelative(2, -step); break;
            case 'e': robot.MoveJointRelative(3, +step); break;
            case 'E': robot.MoveJointRelative(3, -step); break;
            default: break;
        }
    }

    cv::destroyAllWindows();

    if (image_points.size() < 8) {
        std::cerr << "Need at least 8 captured points; got "
                  << image_points.size() << std::endl;
        return 1;
    }

    cv::Mat rvec, tvec;
    std::vector<int> inliers;
    if (!cv::solvePnPRansac(
            robot_points, image_points, K, D,
            rvec, tvec, false, 200, 5.0, 0.99,
            inliers, cv::SOLVEPNP_ITERATIVE)) {
        std::cerr << "solvePnPRansac failed." << std::endl;
        return 1;
    }

    cv::Mat R;
    cv::Rodrigues(rvec, R);

    std::vector<cv::Point2f> projected;
    cv::projectPoints(robot_points, rvec, tvec, K, D, projected);

    double sum_sq = 0.0;
    for (size_t i = 0; i < image_points.size(); ++i) {
        const double dx = image_points[i].x - projected[i].x;
        const double dy = image_points[i].y - projected[i].y;
        sum_sq += dx * dx + dy * dy;
    }
    const double rms = std::sqrt(sum_sq / image_points.size());

    std::cout << "\n========================================\n"
              << "EXTRINSIC CALIBRATION\n"
              << "========================================\n"
              << "R robot->camera:\n" << R << "\n"
              << "t robot->camera:\n" << tvec << "\n"
              << "RMS reprojection error = " << rms << " px\n"
              << "Inliers = " << inliers.size() << " / " << image_points.size()
              << std::endl;

    const std::string output =
        "/home/jvizner/r1_camera_calibration/camera_extrinsics.yaml";

    cv::FileStorage out(output, cv::FileStorage::WRITE);
    out << "camera_matrix" << K;
    out << "distortion" << D;
    out << "rotation_robot_to_camera" << R;
    out << "translation_robot_to_camera" << tvec;
    out << "rms_reprojection_error" << rms;
    out.release();

    std::cout << "[SYS] Saved: " << output << std::endl;
    return 0;
}
