/* +------------------------------------------------------------------------+
   |                     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/gui/CDisplayWindow.h>
#include <mrpt/img/CImage.h>
#include <mrpt/math/CMatrixF.h>
#include <mrpt/math/fourier.h>
#include <mrpt/system/CTicTac.h>

#include <iostream>

using namespace mrpt;
using namespace mrpt::img;
using namespace mrpt::gui;
using namespace mrpt::math;
using namespace mrpt::system;
using namespace std;

#include <mrpt/examples_config.h>
string myDataDir(MRPT_EXAMPLES_BASE_DIRECTORY + string("img_convolution_fft/"));

// ------------------------------------------------------
//				TestImageConvolutionFFT
// ------------------------------------------------------
void TestImageConvolutionFFT()
{
  CTicTac tictac;
  CImage img;
  CMatrixF imgCorr;

  // ====================  1  ===================
  if (!img.loadFromFile(myDataDir + string("test_image.jpg")))
    throw std::runtime_error("Cannot load test image!");

  printf(
      "Computing %ux%u image convolution ...", (unsigned)img.getWidth(), (unsigned)img.getHeight());

  CMatrixF res_R, res_I;

  double meanTime = 0;

  int N = 3;

  // Convolution using FFT 2D:
  for (int nTimes = 0; nTimes < N; nTimes++)
  {
    tictac.Tic();

    size_t x, y;
    size_t actual_lx = img.getWidth();
    size_t actual_ly = img.getHeight();
    size_t lx = mrpt::round2up(actual_lx);
    size_t ly = mrpt::round2up(actual_ly);

    CMatrixF i1(ly, lx), i2;
    // Get as matrixes, padded with zeros up to power-of-two sizes:
    img.getAsMatrix(i1, false);

    // imgWindow.getAsMatrix(i2,false);
    i2.loadFromTextFile(myDataDir + string("test_convolution_window.txt"));
    i2.setSize(ly, lx);

    if (nTimes == 0) printf("\nMax real:%f Min real:%f\n", i1.maxCoeff(), i1.minCoeff());

    // FFT:
    CMatrixF I1_R, I1_I, I2_R, I2_I;
    CMatrixF ZEROS(ly, lx);
    math::dft2_complex(i1, ZEROS, I1_R, I1_I);
    math::dft2_complex(i2, ZEROS, I2_R, I2_I);

    // Compute the COMPLEX MULTIPLICATION of I2 by I1:
    for (y = 0; y < ly; y++)
      for (x = 0; x < lx; x++)
      {
        float r1 = I1_R(y, x);
        float r2 = I2_R(y, x);
        float im1 = I1_I(y, x);
        float im2 = I2_I(y, x);

        I2_R(y, x) = r1 * r2 - im1 * im2;
        I2_I(y, x) = r2 * im1 + r1 * im2;
      }

    // IFFT:
    math::idft2_complex(I2_R, I2_I, res_R, res_I);
    res_R *= 1.0f;  // SCALE!

    meanTime += tictac.Tac();
    printf(" Done,%.06fms\n", tictac.Tac() * 1000.0f);
    printf("Max real:%f Min real:%f\n", res_R.maxCoeff(), res_R.minCoeff());
  }

  printf("Mean time: %.06fms\n", 1000.0f * meanTime / N);

  CDisplayWindow winR("real");

  CImage imgR;
  imgR.setFromMatrix(res_R, false /*it is not normalized */);
  winR.showImage(imgR);
  winR.waitForKey();

  //	DEBUG_SAVE_MATRIX(res_R);
  //	DEBUG_SAVE_MATRIX(res_I);
}

// ------------------------------------------------------
//						MAIN
// ------------------------------------------------------
int main()
{
  try
  {
    TestImageConvolutionFFT();

    return 0;
  }
  catch (const std::exception& e)
  {
    std::cerr << "MRPT error: " << mrpt::exception_to_str(e) << std::endl;
    return -1;
  }
}
