/* +------------------------------------------------------------------------+
   |                     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          |
   +------------------------------------------------------------------------+ */

#include <mrpt/config/CConfigFile.h>
#include <mrpt/gui/CDisplayWindow3D.h>
#include <mrpt/hwdrivers/CCameraSensor.h>
#include <mrpt/opengl/Scene.h>
#include <mrpt/system/CTimeLogger.h>
#include <mrpt/system/filesystem.h>  // for ASSERT_FILE_EXISTS_
#include <mrpt/vision/CStereoRectifyMap.h>

#include <chrono>
#include <iostream>
#include <thread>

using namespace mrpt;
using namespace mrpt::opengl;
using namespace mrpt::gui;
using namespace mrpt::vision;
using namespace mrpt::system;
using namespace mrpt::img;
using namespace mrpt::obs;
using namespace mrpt::system;
using namespace mrpt::config;
using namespace mrpt::img;
using namespace std;

// ------------------------------------------------------
//				TestStereoRectify
// ------------------------------------------------------
void TestStereoRectify(int argc, char** argv)
{
  CTimeLogger timlog;
  mrpt::vision::CStereoRectifyMap rectifyMap;

  // Parse optional arguments:
  if (argc != 1 && argc != 2)
  {
    cout << "Usage:\n"
         << argv[0] << " ==> Run with default camera parameters (from rawlog file)\n"
         << argv[0] << "[params.cfg] ==> Load stereo camera parameters from cfg file\n";
  }
  if (argc == 2)
  {
    const string sCfgFile = argv[1];
    ASSERT_FILE_EXISTS_(sCfgFile);

    // Load params from file:
    mrpt::img::TStereoCamera params;
    params.loadFromConfigFile("CAMERA_PARAMS", mrpt::config::CConfigFile(sCfgFile));

    // Prepare rectify map:
    timlog.enter("rectifyMap.setFromCamParams");
    rectifyMap.setFromCamParams(params);
    timlog.leave("rectifyMap.setFromCamParams");
  }

  // Show to the user a list of possible camera drivers and creates and open
  // the selected camera.
  cout << "Please, select the input stereo camera or rawlog file (with "
          "stereo images)...\n";

  mrpt::hwdrivers::CCameraSensor::Ptr cam = mrpt::hwdrivers::prepareVideoSourceFromUserSelection();
  if (!cam) return;

  cout << "Video stream open OK\n";

  // Create 3D window:
  CDisplayWindow3D win("Demo of stereo rectification", 1280, 600);

  // Create 2 viewports, one for each image:
  std::vector<Viewport::Ptr> gl_views(2);
  {
    Scene::Ptr& theScene = win.get3DSceneAndLock();
    gl_views[0] = theScene->getViewport("main");
    ASSERT_(gl_views[0]);
    gl_views[1] = theScene->createViewport("right_image");

    // Assign sizes:
    gl_views[0]->setViewportPosition(0, 0, .5, 1.);
    gl_views[1]->setViewportPosition(.5, 0, .5, 1.);

    // IMPORTANT!!! IF NOT UNLOCKED, THE WINDOW WILL NOT BE UPDATED!
    win.unlockAccess3DScene();
  }

  win.setPos(10, 10);
  //	win.addTextMessage(...

  bool enable_rectify = true;
  bool enable_draw_epipolar_lines = true;
  CImage img_left_rectified,
      img_right_rectified;  // Declared here to serve as a memory buffer
  // (avoid deallocating/allocating)

  cout << "Close the window to end.\n";
  while (win.isOpen())
  {
    win.addTextMessage(5, 5, format("%.02fFPS", win.getRenderingFPS()), 0);
    win.addTextMessage(
        5, 25,
        format(
            "'r': Switch rectify (Now is: %s) | '+'/'-': Modify "
            "alpha (Now is: %.02f)",
            enable_rectify ? "ON" : "OFF", rectifyMap.getAlpha()),
        10);
    win.addTextMessage(
        5, 50,
        format(
            "'s': Switch resize output to 320x240 (Now is: %s) | 'c': "
            "Switch no-disparity (Now is: %s) | 'e': Switch epipolar lines",
            rectifyMap.isEnabledResizeOutput() ? "ON" : "OFF",
            rectifyMap.isEnabledBothCentersCoincide() ? "ON" : "OFF"),
        11);

    std::this_thread::sleep_for(1ms);

    // Grab new video frame:
    CObservation::Ptr obs = cam->getNextFrame();
    if (obs)
    {
      if (IS_CLASS(*obs, CObservationStereoImages))
      {
        // Get the observation object:
        CObservationStereoImages::Ptr o = std::dynamic_pointer_cast<CObservationStereoImages>(obs);

        // If the rectification maps are still not ready, prepare them
        // now:
        if (!rectifyMap.isSet())
        {
          timlog.enter("rectifyMap.setFromCamParams");
          rectifyMap.setFromCamParams(*o);
          timlog.leave("rectifyMap.setFromCamParams");

          /*mrpt::img::TStereoCamera params;
          o->getStereoCameraParams(params);
          cout << params.dumpAsText() << endl;*/
        }

        win.get3DSceneAndLock();

        if (enable_rectify)
        {
          // Rectify:
          timlog.enter("rectifyMap.rectify()");

          rectifyMap.rectify(o->imageLeft, o->imageRight, img_left_rectified, img_right_rectified);

          timlog.leave("rectifyMap.rectify()");
        }
        else
        {
          // Don't rectify:
          img_left_rectified = o->imageLeft;
          img_right_rectified = o->imageRight;
        }

        // Draw lines:
        if (enable_draw_epipolar_lines)
        {
          const unsigned int LINES_SEP = 40;
          const unsigned int w = img_left_rectified.getWidth();
          const unsigned int h = img_left_rectified.getHeight();
          for (unsigned int y = 0; y < h; y += LINES_SEP)
          {
            img_left_rectified.line(0, y, w - 1, y, mrpt::img::TColor::red(), 2);
            img_right_rectified.line(0, y, w - 1, y, mrpt::img::TColor::red(), 2);
          }
        }

        gl_views[0]->setImageView(img_left_rectified);
        gl_views[1]->setImageView(img_right_rectified);

        win.addTextMessage(150, 5, mrpt::system::timeToString(o->timestamp), 2);

        win.unlockAccess3DScene();
        win.repaint();
      }

      if (win.keyHit())
      {
        mrptKeyModifier kmods;
        int key = win.getPushedKey(&kmods);

        if (key == MRPTK_ESCAPE) break;
        if (key == 'r' || key == 'R') enable_rectify = !enable_rectify;
        if (key == 'e' || key == 'E') enable_draw_epipolar_lines = !enable_draw_epipolar_lines;
        if (key == '+' || key == '-')
        {
          double alpha = rectifyMap.getAlpha() + (key == '-' ? -0.1 : 0.1);
          alpha = std::min(1., std::max(0., alpha));
          rectifyMap.setAlpha(alpha);
        }
        if (key == 's' || key == 'S')
        {
          rectifyMap.enableResizeOutput(!rectifyMap.isEnabledResizeOutput(), 320, 240);
        }
        if (key == 'c' || key == 'C')
        {
          rectifyMap.enableBothCentersCoincide(!rectifyMap.isEnabledBothCentersCoincide());
        }
      }
    }
  }
}

// ------------------------------------------------------
//						MAIN
// ------------------------------------------------------
int main(int argc, char** argv)
{
  try
  {
    TestStereoRectify(argc, argv);
    return 0;
  }
  catch (const std::exception& e)
  {
    std::cerr << "MRPT error: " << mrpt::exception_to_str(e) << std::endl;
    return -1;
  }
}
