waypoint_editor.cpp

/*
MIT License

Copyright 2025 ymd-stella

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

#include <iostream>
#include <vector>
#include <string>
#include <fstream>
#include <set>
#include <memory>
#include <algorithm>
#include <iomanip>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <guik/viewer/light_viewer.hpp>
#include <guik/model_control.hpp>
#include <glk/primitives/primitives.hpp>
#include <imgui.h>
#include <ImGuizmo.h>
#include <portable-file-dialogs.h>
#include <nlohmann/json.hpp>

struct Waypoint {
  size_t id;
  Eigen::Vector3f position;
  Eigen::Quaternionf orientation;

  Eigen::Matrix4f matrix() const {
    Eigen::Matrix4f matrix = Eigen::Matrix4f::Identity();
    matrix.block<3, 3>(0, 0) = orientation.toRotationMatrix();
    matrix.block<3, 1>(0, 3) = position;
    return matrix;
  }
};

namespace Eigen {
// Eigen::Vector3f <-> json
void to_json(nlohmann::json& j, const Eigen::Vector3f& v) {
  j = {v.x(), v.y(), v.z()};
}
void from_json(const nlohmann::json& j, Eigen::Vector3f& v) {
  j.at(0).get_to(v.x());
  j.at(1).get_to(v.y());
  j.at(2).get_to(v.z());
}

// Eigen::Quaternionf <-> json
void to_json(nlohmann::json& j, const Eigen::Quaternionf& q) {
  j = {q.w(), q.x(), q.y(), q.z()};
}
void from_json(const nlohmann::json& j, Eigen::Quaternionf& q) {
  j.at(0).get_to(q.w());
  j.at(1).get_to(q.x());
  j.at(2).get_to(q.y());
  j.at(3).get_to(q.z());
}
}  // namespace Eigen

// Waypoint <-> json
void to_json(nlohmann::json& j, const Waypoint& wp) {
  j = nlohmann::json{{"id", wp.id}, {"position", wp.position}, {"orientation", wp.orientation}};
}
void from_json(const nlohmann::json& j, Waypoint& wp) {
  j.at("id").get_to(wp.id);
  j.at("position").get_to(wp.position);
  j.at("orientation").get_to(wp.orientation);
}

class WaypointEditor {
public:
  WaypointEditor() {
    viewer = guik::LightViewer::instance();
    gizmo_control = std::make_unique<guik::ModelControl>("gizmo");
    next_waypoint_id = 0;
  }

  // Main application loop
  void run() {
    setup_ui();
    viewer->spin();
  }

private:
  // Setup UI callbacks
  void setup_ui() {
    viewer->enable_info_buffer();
    viewer->register_ui_callback("waypoint_editor", [&]() {
      auto& io = ImGui::GetIO();
      if (!io.WantCaptureMouse && io.MouseClicked[ImGuiMouseButton_Left]) {
        float depth = viewer->pick_depth({io.MousePos.x, io.MousePos.y});
        if (depth < 1.0f) {
          Eigen::Vector4i info = viewer->pick_info({io.MousePos.x, io.MousePos.y});
          auto waypoint_id = info[0];
          bool is_selected = selected_indices.count(waypoint_id);
          if (!ImGui::GetIO().KeyCtrl) {
            selected_indices.clear();
          }
          if (is_selected) {
            selected_indices.erase(waypoint_id);
          } else {
            selected_indices.insert(waypoint_id);
            active_waypoint_index = waypoint_id;
          }
        }
      }
      update_ui();
      update_scene();
      handle_gizmo();
    });
  }

  // Draw all UI panels
  void update_ui() {
    update_main_control_panel();
    update_waypoint_inspector();
    update_properties_editor();
  }

  // Control Panel for file I/O and adding/deleting waypoints
  void update_main_control_panel() {
    ImGui::Begin("Control Panel");

    if (ImGui::Button("Add Waypoint")) {
      add_waypoint();
    }
    ImGui::SameLine();
    if (ImGui::Button("Delete Selected") && !selected_indices.empty()) {
      delete_selected();
    }

    ImGui::Separator();

    if (ImGui::Button("Load from File...")) {
      load_from_file();
    }
    ImGui::SameLine();
    if (ImGui::Button("Save to File...")) {
      save_to_file();
    }

    ImGui::Separator();
    ImGui::Text("Gizmo Controls");
    gizmo_control->draw_gizmo_ui();

    ImGui::End();
  }

  // Inspector panel to list and select waypoints
  void update_waypoint_inspector() {
    ImGui::Begin("Waypoint Inspector");
    ImGui::Text("%lu waypoints", waypoints.size());
    ImGui::BeginChild("WaypointList", ImVec2(0, 0), true);

    for (size_t i = 0; i < waypoints.size(); ++i) {
      const auto& wp = waypoints[i];
      std::string label = "Waypoint " + std::to_string(wp.id);
      bool is_selected = selected_indices.count(i);

      if (ImGui::Selectable(label.c_str(), is_selected)) {
        if (!ImGui::GetIO().KeyCtrl) {
          selected_indices.clear();
        }
        if (is_selected) {
          selected_indices.erase(i);
        } else {
          selected_indices.insert(i);
          active_waypoint_index = i;
        }
      }
    }
    ImGui::EndChild();
    ImGui::End();
  }

  // Properties editor for the active waypoint
  void update_properties_editor() {
    ImGui::Begin("Properties");
    if (selected_indices.empty() || active_waypoint_index >= waypoints.size()) {
      ImGui::Text("No waypoint selected.");
    } else {
      auto& wp = waypoints[active_waypoint_index];
      ImGui::Text("Editing Waypoint %zu", wp.id);

      ImGui::DragFloat3("Position", wp.position.data(), 0.1f);

      // Edit orientation as Euler angles for user-friendliness
      Eigen::Vector3f euler = wp.orientation.toRotationMatrix().eulerAngles(0, 1, 2);  // Roll, Pitch, Yaw
      euler *= 180.0f / M_PI;                                                          // Convert to degrees
      if (ImGui::DragFloat3("Rotation (RPY)", euler.data(), 1.0f)) {
        euler *= M_PI / 180.0f;  // Convert back to radians
        wp.orientation =
          Eigen::AngleAxisf(euler[0], Eigen::Vector3f::UnitX()) * Eigen::AngleAxisf(euler[1], Eigen::Vector3f::UnitY()) * Eigen::AngleAxisf(euler[2], Eigen::Vector3f::UnitZ());
      }
    }
    ImGui::End();
  }

  // Synchronize waypoint data with 3D scene drawables
  void update_scene() {
    // First, remove drawables for waypoints that no longer exist
    std::set<std::string> existing_drawable_names;
    for (size_t i = 0; i < waypoints.size(); ++i) {
      existing_drawable_names.insert("wp_" + std::to_string(waypoints[i].id));
    }

    while (!waypoints_deleted.empty()) {
      auto waypoint = waypoints_deleted.back();
      viewer->remove_drawable("wp_" + std::to_string(waypoint.id) + "_axis_x");
      waypoints_deleted.pop_back();
    }

    // Update/create drawables for all current waypoints
    for (size_t i = 0; i < waypoints.size(); ++i) {
      const auto& wp = waypoints[i];
      bool is_selected = selected_indices.count(i);
      update_waypoint(wp.id, wp.matrix(), is_selected);
    }
  }

  // Draw a single waypoint as a coordinate frame primitive
  void update_waypoint(const size_t waypoint_id, const Eigen::Matrix4f& pose, bool is_selected) {
    Eigen::Matrix4f transform = Eigen::Matrix4f::Identity();
    viewer->update_drawable(
      "wp_" + std::to_string(waypoint_id) + "_axis_x",
      glk::Primitives::coordinate_system(),
      guik::VertexColor(pose * transform).add("info_values", Eigen::Vector4i(waypoint_id, 0, 0, 0)));
  }

  // Handle ImGuizmo interaction and update waypoint data
  void handle_gizmo() {
    if (selected_indices.empty() || active_waypoint_index >= waypoints.size()) {
      return;
    }

    // Set the gizmo's transform to the active waypoint's transform
    gizmo_control->set_model_matrix(waypoints[active_waypoint_index].matrix());

    // Draw the gizmo and check for interaction
    gizmo_control->draw_gizmo();

    // Gizmo was manipulated
    Eigen::Matrix4f new_active_matrix = gizmo_control->model_matrix();
    Eigen::Matrix4f old_active_matrix = waypoints[active_waypoint_index].matrix();

    // Calculate delta transform
    Eigen::Matrix4f delta_transform = new_active_matrix * old_active_matrix.inverse();

    // Apply delta transform to all selected waypoints
    for (const auto& index : selected_indices) {
      if (index < waypoints.size()) {
        Eigen::Matrix4f new_matrix;
        if (index == active_waypoint_index) {
          new_matrix = new_active_matrix;
        } else {
          new_matrix = delta_transform * waypoints[index].matrix();
        }

        // Decompose matrix and update waypoint data
        waypoints[index].position = new_matrix.block<3, 1>(0, 3);
        waypoints[index].orientation = Eigen::Quaternionf(new_matrix.block<3, 3>(0, 0));
      }
    }
  }

  void add_waypoint() {
    Waypoint wp;
    wp.id = next_waypoint_id++;
    wp.position = Eigen::Vector3f::Random() * 2.0f;  // Place randomly for demonstration
    wp.orientation.setIdentity();
    waypoints.push_back(wp);
  }

  void delete_selected() {
    // Sort indices in descending order to avoid invalidating indices during removal
    std::vector<size_t> sorted_indices(selected_indices.begin(), selected_indices.end());
    std::sort(sorted_indices.rbegin(), sorted_indices.rend());

    for (const auto& index : sorted_indices) {
      if (index < waypoints.size()) {
        waypoints_deleted.push_back(waypoints[index]);
        waypoints.erase(waypoints.begin() + index);
      }
    }
    selected_indices.clear();
    active_waypoint_index = -1;  // Invalidate active index
  }

  void load_from_file() {
    auto selections = pfd::open_file("Open Waypoints", ".", {"JSON Files", "*.json"}).result();
    for (auto&& selection : selections) {
      std::ifstream file(selection);
      if (file.is_open()) {
        try {
          nlohmann::json j;
          file >> j;
          waypoints = j.get<std::vector<Waypoint>>();

          // Update next ID to avoid collisions
          size_t max_id = 0;
          for (const auto& wp : waypoints) {
            if (wp.id > max_id) max_id = wp.id;
          }
          next_waypoint_id = max_id + 1;

          selected_indices.clear();
        } catch (const std::exception& e) {
          pfd::message("Error", "Failed to parse JSON file: " + std::string(e.what()), pfd::choice::ok, pfd::icon::error);
        }
      }
    }
  }

  void save_to_file() {
    auto destination = pfd::save_file("Save Waypoints", "waypoints.json", {"JSON Files", "*.json"}).result();
    if (!destination.empty()) {
      nlohmann::json j = waypoints;
      std::ofstream file(destination);
      file << std::setw(4) << j << std::endl;
    }
  }

private:
  guik::LightViewer* viewer;
  std::vector<Waypoint> waypoints;
  std::vector<Waypoint> waypoints_deleted;
  std::set<size_t> selected_indices;
  size_t active_waypoint_index = -1;
  size_t next_waypoint_id;

  std::unique_ptr<guik::ModelControl> gizmo_control;
};

int main(int argc, char** argv) {
  WaypointEditor editor;
  editor.run();
  return 0;
}