/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2024, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

#ifdef __EMSCRIPTEN__  // are we actually in a Wasm JS instance?
#include <emscripten.h>
#include <emscripten/html5.h>

#define GL_GLEXT_PROTOTYPES
#define EGL_EGLEXT_PROTOTYPES
#include <GLFW/glfw3.h>

#include <functional>

std::function<void()> loop;
void main_loop() { loop(); }
#endif

#include <mrpt/core/exceptions.h>
#include <mrpt/core/lock_helper.h>
#include <mrpt/gui/CDisplayWindowGUI.h>
#include <mrpt/opengl/CAxis.h>
#include <mrpt/opengl/CGridPlaneXY.h>
#include <mrpt/opengl/stock_objects.h>
#include <mrpt/poses/Lie/SO.h>

#include <Eigen/Dense>
#include <Eigen/SVD>
#include <iostream>

// We need to define all variables here to avoid emscripten memory errors
// if they are defined as *locals* in the function.
struct AppData
{
  mrpt::gui::CDisplayWindowGUI::Ptr win;
  mrpt::opengl::CSetOfObjects::Ptr gl_corner_user;
  mrpt::opengl::CAxis::Ptr gl_corner_reference;

  // Input variables (they are bound to the GUI controls):
  mrpt::math::CQuaternionDouble in_quat;
  mrpt::math::CMatrixFixed<double, 3, 3> in_rot;
  std::array<double, 3> in_ypr = {0, 0, 0};
  mrpt::math::TVector3D in_axisangle_ax{0, 0, 0};
  double in_axisangle_ang = 0;
  mrpt::math::CVectorFixed<double, 3> in_lie_log = mrpt::math::CVectorFixed<double, 3>::Zero();
  bool units_radians = true;

  // In/out UI control declarations (declared here so we can use them in the
  // lambda below):
  nanogui::TabWidget* tabWidget = nullptr;
  nanogui::TextBox* ed_in_rot[3][3];
  nanogui::TextBox *edOutMatrix = nullptr, *edOutQuat = nullptr, *edOutAxisAngle_Ax = nullptr,
                   *edOutAxisAngle_An = nullptr, *edOutLogSO3 = nullptr;
  nanogui::TextBox* ed_out_rot[3][3];
  nanogui::Slider* sl_in_ypr[3] = {nullptr, nullptr, nullptr};
};

static AppData app;

// The main function: update all calculations:
static void recalcAll()
{
  mrpt::poses::CPose3D userPose;
  switch (app.tabWidget->activeTab())
  {
    // YPR
    case 0:
    {
      const double K = app.units_radians ? 1.0 : (M_PI / 180.0);
      userPose.setFromValues(0, 0, 0, app.in_ypr[0] * K, app.in_ypr[1] * K, app.in_ypr[2] * K);
    }
    break;
    // Rot matrix:
    case 1:
    {
      // Run a SVD to ensure we take the closest SO(3) matrix to the
      // user input:
      Eigen::Matrix3d M = app.in_rot.asEigen();
      Eigen::JacobiSVD<Eigen::Matrix3d> svd(M, Eigen::ComputeFullU | Eigen::ComputeFullV);

      auto R = mrpt::math::CMatrixDouble33((svd.matrixU() * svd.matrixV().transpose()).eval());
      userPose.setRotationMatrix(R);
      for (int r = 0; r < 3; r++)
        for (int c = 0; c < 3; c++) app.ed_in_rot[r][c]->setValue(mrpt::format("%.05f", R(r, c)));
    }
    break;
    // Quaternion
    case 2:
    {
      userPose = mrpt::poses::CPose3D(app.in_quat, 0, 0, 0);
    }
    break;
    // axis+angle
    case 3:
    {
      const double K = app.units_radians ? 1.0 : (M_PI / 180.0);
      mrpt::math::TVector3D v = app.in_axisangle_ax * (app.in_axisangle_ang * K);
      mrpt::math::CVectorFixed<double, 3> vn;
      for (int i = 0; i < 3; i++) vn[i] = v[i];

      userPose.setRotationMatrix(mrpt::poses::Lie::SO<3>::exp(vn));
    }
    break;
    // log(SO(3))
    case 4:
    {
      userPose.setRotationMatrix(mrpt::poses::Lie::SO<3>::exp(app.in_lie_log));
    }
    break;
  };

  // Set 3D view corner:
  app.gl_corner_user->setPose(userPose);

  // Update text output:
  const mrpt::math::CMatrixDouble33 Rout = userPose.getRotationMatrix();

  // matrix:
  app.edOutMatrix->setValue(Rout.inMatlabFormat());
  for (int r = 0; r < 3; r++)
    for (int c = 0; c < 3; c++) app.ed_out_rot[r][c]->setValue(mrpt::format("%.05f", Rout(r, c)));

  // quat:
  mrpt::math::CQuaternionDouble q;
  userPose.getAsQuaternion(q);
  app.edOutQuat->setValue(q.asString());

  // SO(3) log:
  const auto log_R = mrpt::poses::Lie::SO<3>::log(Rout);
  app.edOutLogSO3->setValue(mrpt::format("[%.05f %.05f %.05f]", log_R[0], log_R[1], log_R[2]));

  // axis-angle:
  mrpt::math::TVector3D axis(log_R[0], log_R[1], log_R[2]);
  if (axis.norm() > 1e-20) axis = axis.unitarize();

  app.edOutAxisAngle_Ax->setValue(mrpt::format("[%.05f %.05f %.05f]", axis[0], axis[1], axis[2]));
  app.edOutAxisAngle_An->setValue(mrpt::format(
      "%.05f %s", log_R.norm() * (app.units_radians ? 1.0 : 180.0 / M_PI),
      (app.units_radians ? "rad" : "deg")));
};

static void AppRotationConverter()
{
  nanogui::init();

  // Create main window:
  mrpt::gui::CDisplayWindowGUI_Params cp;
  // cp.fullscreen = true;

  // Init values:
  app.in_rot.setIdentity();

  app.gl_corner_user = mrpt::opengl::stock_objects::CornerXYZ(1.0f);
  app.gl_corner_reference =
      mrpt::opengl::CAxis::Create(-1.0f, -1.0f, -1.0f, 1.0f, 1.0f, 1.0f, 0.2f, 1.0f, true);
  app.gl_corner_reference->setTextScale(0.04);

  // Create GUI:
  app.win = mrpt::gui::CDisplayWindowGUI::Create("3D rotation converter", 900, 700, cp);

  // Add INPUT window:
  // -----------------------------
  nanogui::Window* winInput = new nanogui::Window(&(*app.win), "Rotation input");
  winInput->setPosition(nanogui::Vector2i(10, 50));
  winInput->setLayout(new nanogui::GroupLayout());
  winInput->setFixedWidth(350);

  app.tabWidget = winInput->add<nanogui::TabWidget>();

  {  // Yaw/pitch/roll  tab:
    nanogui::Widget* layer = app.tabWidget->createTab("Yaw-pitch-roll");
    auto layout = new nanogui::GridLayout(
        nanogui::Orientation::Horizontal, 3, nanogui::Alignment::Fill, 5, 0);
    layer->setLayout(layout);

    const char* lb[3] = {"yaw (Z)=", "pitch (Y)=", "roll (x)="};

    for (int i = 0; i < 3; i++)
    {
      layer->add<nanogui::Label>(lb[i]);
      nanogui::TextBox* ed = layer->add<nanogui::TextBox>(mrpt::format("%.2f", app.in_quat[i]));
      ed->setEditable(true);
      ed->setFormat("[-+]?[0-9.e+-]*");
      ed->setCallback(
          [i](const std::string& s)
          {
            app.in_ypr[i] = std::stod(s);
            app.sl_in_ypr[i]->setValue(app.in_ypr[i] / (app.units_radians ? M_PIf : 180.0f));
            recalcAll();
            return true;
          });
      nanogui::Slider* sl = layer->add<nanogui::Slider>();
      app.sl_in_ypr[i] = sl;
      sl->setRange({-1.0f, 1.0f});
      sl->setCallback(
          [&, i, ed](float val)
          {
            val *= app.units_radians ? M_PIf : 180.0f;
            if (i == 1) val *= 0.5f;  // Pitch
            ed->setValue(mrpt::format("%.03f", val));
            ed->callback()(ed->value());
          });
    }
  }

  {  // rotation matrix tab:
    nanogui::Widget* layer = app.tabWidget->createTab("SO(3) matrix");

    auto layout = new nanogui::GridLayout(
        nanogui::Orientation::Horizontal, 3, nanogui::Alignment::Fill, 5, 0);
    layer->setLayout(layout);

    for (int r = 0; r < 3; r++)
    {
      for (int c = 0; c < 3; c++)
      {
        nanogui::TextBox* ed =
            layer->add<nanogui::TextBox>(mrpt::format("%.04f", app.in_rot(r, c)));
        ed->setEditable(true);
        ed->setFormat("[-+]?[0-9.e+-]*");
        ed->setCallback(
            [r, c](const std::string& s)
            {
              app.in_rot(r, c) = std::stod(s);
              return true;
            });
        app.ed_in_rot[r][c] = ed;
      }
    }
  }

  {  // Quaternion tab:
    nanogui::Widget* layer = app.tabWidget->createTab("Quaternion");
    auto layout = new nanogui::GridLayout(
        nanogui::Orientation::Horizontal, 2, nanogui::Alignment::Fill, 5, 0);
    layer->setLayout(layout);

    const char* lb[4] = {"w (real)=", "x=", "y=", "z="};

    for (int i = 0; i < 4; i++)
    {
      layer->add<nanogui::Label>(lb[i]);

      nanogui::TextBox* ed = layer->add<nanogui::TextBox>(mrpt::format("%.2f", app.in_quat[i]));
      ed->setEditable(true);
      ed->setFormat("[-+]?[0-9.e+-]*");
      ed->setCallback(
          [i](const std::string& s)
          {
            app.in_quat[i] = std::stod(s);
            return true;
          });
    }
  }

  {  // Axis-angle tab:
    nanogui::Widget* layer = app.tabWidget->createTab("Axis-angle");
    layer->setLayout(new nanogui::GridLayout(
        nanogui::Orientation::Horizontal, 2, nanogui::Alignment::Fill, 5, 0));
    layer->add<nanogui::Label>("Axis:");

    {
      nanogui::Widget* panel = layer->add<nanogui::Widget>();
      panel->setLayout(new nanogui::GridLayout(
          nanogui::Orientation::Horizontal, 3, nanogui::Alignment::Fill, 5, 0));

      for (int r = 0; r < 3; r++)
      {
        nanogui::TextBox* ed =
            panel->add<nanogui::TextBox>(mrpt::format("%.04f", app.in_axisangle_ax[r]));
        ed->setEditable(true);
        ed->setFormat("[-+]?[0-9.e+-]*");
        ed->setCallback(
            [r](const std::string& s)
            {
              app.in_axisangle_ax[r] = std::stod(s);
              return true;
            });
      }
    }
    {
      layer->add<nanogui::Label>("Angle:");
      nanogui::TextBox* ed =
          layer->add<nanogui::TextBox>(mrpt::format("%.04f", app.in_axisangle_ang));
      ed->setEditable(true);
      ed->setFormat("[-+]?[0-9.e+-]*");
      ed->setCallback(
          [](const std::string& s)
          {
            app.in_axisangle_ang = std::stod(s);
            return true;
          });
    }
  }

  {  // axis with magnitude = log(R) in Lie group
    nanogui::Widget* layer = app.tabWidget->createTab("log(SO(3))");

    auto layout = new nanogui::GridLayout(
        nanogui::Orientation::Horizontal, 1, nanogui::Alignment::Fill, 5, 0);
    layer->setLayout(layout);

    layer->add<nanogui::Label>("Axis with magnitude", "sans-bold");
    layer->add<nanogui::Label>("(i.e. vee(log(R)) in SO(3))");
    {
      nanogui::Widget* panel = layer->add<nanogui::Widget>();
      panel->setLayout(new nanogui::GridLayout(
          nanogui::Orientation::Horizontal, 3, nanogui::Alignment::Fill, 5, 0));

      for (int r = 0; r < 3; r++)
      {
        nanogui::TextBox* ed =
            layer->add<nanogui::TextBox>(mrpt::format("%.04f", app.in_lie_log[r]));
        ed->setEditable(true);
        ed->setFormat("[-+]?[0-9.e+-]*");
        ed->setCallback(
            [r](const std::string& s)
            {
              app.in_lie_log[r] = std::stod(s);
              return true;
            });
      }
    }
  }

  app.tabWidget->setActiveTab(0);

  // Apply button:
  winInput->add<nanogui::Button>("Apply", ENTYPO_ICON_CALCULATOR)->setCallback(recalcAll);

  // Add top menu subwindow:
  // -----------------------------
  {
    nanogui::ref<nanogui::Window> winMenu = new nanogui::Window(&(*app.win), "");
    winMenu->setPosition(nanogui::Vector2i(0, 0));
    winMenu->setLayout(
        new nanogui::BoxLayout(nanogui::Orientation::Horizontal, nanogui::Alignment::Middle, 5));
    nanogui::Theme* modTheme = new nanogui::Theme(app.win->screen()->nvgContext());
    modTheme->mWindowHeaderHeight = 1;
    winMenu->setTheme(modTheme);

#ifndef __EMSCRIPTEN__  // are we actually in a Wasm JS instance?
    winMenu->add<nanogui::Button>("Quit", ENTYPO_ICON_ARROW_BOLD_LEFT)
        ->setCallback([]() { app.win->setVisible(false); });
#endif

    winMenu->add<nanogui::Label>("      ");  // separator

    winMenu
        ->add<nanogui::CheckBox>(
            "Show reference frame", [&](bool b) { app.gl_corner_reference->setVisibility(b); })
        ->setChecked(true);

    winMenu
        ->add<nanogui::CheckBox>(
            "Show rotated frame", [&](bool b) { app.gl_corner_user->setVisibility(b); })
        ->setChecked(true);
    winMenu
        ->add<nanogui::CheckBox>(
            "Ortho. view", [&](bool b) { app.win->camera().setCameraProjective(!b); })
        ->setChecked(false);

    winMenu->add<nanogui::Label>("Units:");
    winMenu->add<nanogui::ComboBox>(std::vector<std::string>({"radians", "degrees"}))
        ->setCallback(
            [&](int index)
            {
              // On units update:
              app.units_radians = (index == 0);
              recalcAll();
            });
  }

  // Add OUTPUT window:
  // -----------------------------
  {
    nanogui::Window* winOutput = new nanogui::Window(&(*app.win), "Rotation output");
    winOutput->setPosition(nanogui::Vector2i(10, 320));
    auto layout = new nanogui::GridLayout(
        nanogui::Orientation::Horizontal, 1, nanogui::Alignment::Fill, 5, 0);
    winOutput->setLayout(layout);
    winOutput->setFixedWidth(350);

    winOutput->add<nanogui::Label>("SO(3) rotation matrix", "sans-bold");
    {
      nanogui::Widget* panel = winOutput->add<nanogui::Widget>();
      panel->setLayout(new nanogui::GridLayout(
          nanogui::Orientation::Horizontal, 3, nanogui::Alignment::Fill, 5, 0));
      for (int r = 0; r < 3; r++)
        for (int c = 0; c < 3; c++) app.ed_out_rot[r][c] = panel->add<nanogui::TextBox>("");
    }

    winOutput->add<nanogui::Label>("In MATLAB-like notation:");
    app.edOutMatrix = winOutput->add<nanogui::TextBox>("");
    app.edOutMatrix->setEditable(true);

    winOutput->add<nanogui::Label>("Quaternion (r,x,y,z)", "sans-bold");
    app.edOutQuat = winOutput->add<nanogui::TextBox>("");
    app.edOutQuat->setEditable(true);

    winOutput->add<nanogui::Label>("Axis-angle (r,x,y,z)", "sans-bold");
    app.edOutAxisAngle_Ax = winOutput->add<nanogui::TextBox>("");
    app.edOutAxisAngle_Ax->setEditable(true);
    app.edOutAxisAngle_An = winOutput->add<nanogui::TextBox>("");
    app.edOutAxisAngle_An->setEditable(true);

    winOutput->add<nanogui::Label>("Axis with angle (log(R))", "sans-bold");
    app.edOutLogSO3 = winOutput->add<nanogui::TextBox>("");
    app.edOutLogSO3->setEditable(true);
  }

  // Add a background scene:
  // -----------------------------
  {
    auto scene = mrpt::opengl::Scene::Create();
    scene->insert(mrpt::opengl::CGridPlaneXY::Create());

    scene->insert(app.gl_corner_user);
    scene->insert(app.gl_corner_reference);

    auto lck = mrpt::lockHelper(app.win->background_scene_mtx);
    app.win->background_scene = std::move(scene);
  }

  app.win->performLayout();

  app.win->camera().setZoomDistance(5.0f);

  // Update view and process events:
  app.win->drawAll();
  app.win->setVisible(true);

#if !defined(__EMSCRIPTEN__)  // are we in a Wasm JS instance?
  // No: regular procedure:

  nanogui::mainloop();
  nanogui::shutdown();
#else

  // Yes, we are in a web browser running on JS:
  loop = []
  {
    // Check if any events have been activated (key pressed, mouse moved
    // etc.) and call corresponding response functions
    glfwPollEvents();

    glClearColor(0.2f, 0.25f, 0.3f, 1.0f);
    glClear(GL_COLOR_BUFFER_BIT);

    // Draw nanogui
    app.win->drawContents();
    app.win->drawWidgets();

    glfwSwapBuffers(app.win->glfwWindow());
  };
  emscripten_set_main_loop(main_loop, 0, true);

  glfwDestroyWindow(app.win->nanogui_screen()->glfwWindow());

  // Terminate GLFW, clearing any resources allocated by GLFW.
  glfwTerminate();

#endif
}

int main()
{
  try
  {
    AppRotationConverter();
    return 0;
  }
  catch (const std::exception& e)
  {
    std::cerr << mrpt::exception_to_str(e) << std::endl;
    return -1;
  }
}
