/* +------------------------------------------------------------------------+
   |                     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 "poses-precomp.h"  // Precompiled headers
//
#include <mrpt/math/matrix_serialization.h>
#include <mrpt/math/ops_matrices.h>
#include <mrpt/math/wrap2pi.h>
#include <mrpt/poses/CPose3D.h>
#include <mrpt/poses/CPose3DPDFGaussian.h>
#include <mrpt/poses/CPose3DPDFGaussianInf.h>
#include <mrpt/poses/CPose3DQuatPDFGaussian.h>
#include <mrpt/poses/CPosePDFGaussianInf.h>
#include <mrpt/random.h>
#include <mrpt/serialization/CArchive.h>
#include <mrpt/system/os.h>

#include <Eigen/Dense>

using namespace mrpt;
using namespace mrpt::poses;
using namespace mrpt::math;
using namespace mrpt::random;
using namespace mrpt::system;
using namespace std;

IMPLEMENTS_SERIALIZABLE(CPose3DPDFGaussianInf, CPose3DPDF, mrpt::poses)

/*---------------------------------------------------------------
  Constructor
  ---------------------------------------------------------------*/
CPose3DPDFGaussianInf::CPose3DPDFGaussianInf() : mean(0, 0, 0), cov_inv() {}
/*---------------------------------------------------------------
  Constructor
  ---------------------------------------------------------------*/
CPose3DPDFGaussianInf::CPose3DPDFGaussianInf(
    [[maybe_unused]] TConstructorFlags_Poses constructor_dummy_param) :
    mean(UNINITIALIZED_POSE), cov_inv(UNINITIALIZED_MATRIX)
{
}

/*---------------------------------------------------------------
  Constructor
  ---------------------------------------------------------------*/
CPose3DPDFGaussianInf::CPose3DPDFGaussianInf(
    const CPose3D& init_Mean, const CMatrixDouble66& init_Cov) :
    mean(init_Mean), cov_inv(init_Cov)
{
}

/*---------------------------------------------------------------
  Constructor
  ---------------------------------------------------------------*/
CPose3DPDFGaussianInf::CPose3DPDFGaussianInf(const CPose3D& init_Mean) : mean(init_Mean), cov_inv()
{
}

/*---------------------------------------------------------------
          CPose3DPDFGaussianInf
 ---------------------------------------------------------------*/
CPose3DPDFGaussianInf::CPose3DPDFGaussianInf(const CPose3DQuatPDFGaussian& o) :
    mean(UNINITIALIZED_POSE), cov_inv(UNINITIALIZED_MATRIX)
{
  this->copyFrom(o);
}

/*---------------------------------------------------------------
            copyFrom
 ---------------------------------------------------------------*/
void CPose3DPDFGaussianInf::copyFrom(const CPose3DQuatPDFGaussian& o)
{
  const CPose3DPDFGaussian p(o);
  this->copyFrom(p);
}

uint8_t CPose3DPDFGaussianInf::serializeGetVersion() const { return 0; }
void CPose3DPDFGaussianInf::serializeTo(mrpt::serialization::CArchive& out) const
{
  out << mean;
  mrpt::math::serializeSymmetricMatrixTo(cov_inv, out);
}
void CPose3DPDFGaussianInf::serializeFrom(mrpt::serialization::CArchive& in, uint8_t version)
{
  switch (version)
  {
    case 0:
    {
      in >> mean;
      mrpt::math::deserializeSymmetricMatrixFrom(cov_inv, in);
    }
    break;
    default:
      MRPT_THROW_UNKNOWN_SERIALIZATION_VERSION(version);
  };
}

void CPose3DPDFGaussianInf::copyFrom(const CPose3DPDF& o)
{
  if (this == &o) return;  // It may be used sometimes

  if (IS_CLASS(o, CPose3DPDFGaussianInf))
  {  // It's my same class:
    const auto* ptr = dynamic_cast<const CPose3DPDFGaussianInf*>(&o);
    ASSERT_(ptr != nullptr);
    mean = ptr->mean;
    cov_inv = ptr->cov_inv;
  }
  else
  {
    // Convert to gaussian pdf:
    CMatrixDouble66 cov(UNINITIALIZED_MATRIX);
    o.getCovarianceAndMean(cov, mean);
    this->cov_inv = cov.inverse_LLt();
  }
}

void CPose3DPDFGaussianInf::copyFrom(const CPosePDF& o)
{
  if (IS_CLASS(o, CPosePDFGaussianInf))
  {  // cov is already inverted, but it's
     // a 2D pose:
    const auto* ptr = dynamic_cast<const CPosePDFGaussianInf*>(&o);
    ASSERT_(ptr != nullptr);

    mean = CPose3D(ptr->mean);

    // 3x3 inv_cov -> 6x6 inv_cov
    cov_inv.setZero();
    cov_inv(0, 0) = ptr->cov_inv(0, 0);
    cov_inv(1, 1) = ptr->cov_inv(1, 1);
    cov_inv(3, 3) = ptr->cov_inv(2, 2);

    cov_inv(0, 1) = cov_inv(1, 0) = ptr->cov_inv(0, 1);
    cov_inv(0, 3) = cov_inv(3, 0) = ptr->cov_inv(0, 2);
    cov_inv(1, 3) = cov_inv(3, 1) = ptr->cov_inv(1, 2);
  }
  else
  {
    CPose3DPDFGaussian p(UNINITIALIZED_POSE);
    p.copyFrom(o);
    this->copyFrom(p);
  }
}

/*---------------------------------------------------------------

  ---------------------------------------------------------------*/
bool CPose3DPDFGaussianInf::saveToTextFile(const string& file) const
{
  FILE* f = os::fopen(file.c_str(), "wt");
  if (!f) return false;

  os::fprintf(
      f, "%e %e %e %e %e %e\n", mean.x(), mean.y(), mean.z(), mean.yaw(), mean.pitch(),
      mean.roll());

  for (unsigned int i = 0; i < 6; i++)
    os::fprintf(
        f, "%e %e %e %e %e %e\n", cov_inv(i, 0), cov_inv(i, 1), cov_inv(i, 2), cov_inv(i, 3),
        cov_inv(i, 4), cov_inv(i, 5));

  os::fclose(f);
  return true;
}

/*---------------------------------------------------------------
            changeCoordinatesReference
 ---------------------------------------------------------------*/
void CPose3DPDFGaussianInf::changeCoordinatesReference(const CPose3D& newReferenceBase)
{
  MRPT_START

  CPose3DPDFGaussian a;
  a.copyFrom(*this);
  a.changeCoordinatesReference(newReferenceBase);
  this->copyFrom(a);

  MRPT_END
}

/*---------------------------------------------------------------
          drawSingleSample
 ---------------------------------------------------------------*/
void CPose3DPDFGaussianInf::drawSingleSample([[maybe_unused]] CPose3D& outPart) const
{
  MRPT_START

  const CMatrixDouble66 cov = cov_inv.inverse_LLt();

  CVectorDouble v;
  getRandomGenerator().drawGaussianMultivariate(v, cov);

  outPart.setFromValues(
      mean.x() + v[0], mean.y() + v[1], mean.z() + v[2], mean.yaw() + v[3], mean.pitch() + v[4],
      mean.roll() + v[5]);

  MRPT_END_WITH_CLEAN_UP(cov_inv.saveToTextFile("__DEBUG_EXC_DUMP_drawSingleSample_COV_INV.txt"););
}

/*---------------------------------------------------------------
          drawManySamples
 ---------------------------------------------------------------*/
void CPose3DPDFGaussianInf::drawManySamples(size_t N, vector<CVectorDouble>& outSamples) const
{
  MRPT_START

  CMatrixDouble66 cov = this->cov_inv.inverse_LLt();

  getRandomGenerator().drawGaussianMultivariateMany(outSamples, N, cov);

  for (auto& outSample : outSamples)
  {
    outSample[0] += mean.x();
    outSample[1] += mean.y();
    outSample[2] += mean.z();
    outSample[3] = math::wrapToPi(outSample[3] + mean.yaw());
    outSample[4] = math::wrapToPi(outSample[4] + mean.pitch());
    outSample[5] = math::wrapToPi(outSample[5] + mean.roll());
  }

  MRPT_END
}

/*---------------------------------------------------------------
          bayesianFusion
 ---------------------------------------------------------------*/
void CPose3DPDFGaussianInf::bayesianFusion(
    [[maybe_unused]] const CPose3DPDF& p1_, [[maybe_unused]] const CPose3DPDF& p2_)
{
  THROW_EXCEPTION("TO DO!!!");
}

/*---------------------------------------------------------------
          inverse
 ---------------------------------------------------------------*/
void CPose3DPDFGaussianInf::inverse(CPose3DPDF& o) const
{
  ASSERT_(o.GetRuntimeClass() == CLASS_ID(CPose3DPDFGaussianInf));
  auto& out = dynamic_cast<CPose3DPDFGaussianInf&>(o);

  // This is like: b=(0,0,0)
  //  OUT = b - THIS
  CPose3DPDFGaussianInf b;  // Init: all zeros.
  out = b - *this;
}

/*---------------------------------------------------------------
              +=
 ---------------------------------------------------------------*/
void CPose3DPDFGaussianInf::operator+=(const CPose3D& Ap)
{
  // COV:
  const CMatrixDouble66 OLD_COV_INV = this->cov_inv;
  CMatrixDouble66 df_dx(UNINITIALIZED_MATRIX), df_du(UNINITIALIZED_MATRIX);

  CPose3DPDF::jacobiansPoseComposition(
      this->mean,  // x
      Ap,          // u
      df_dx, df_du);

  // this->cov = H1*this->cov*H1' + H2*Ap.cov*H2';
  // cov_inv   = ... => The same than above!
  cov_inv = mrpt::math::multiply_HCHt(df_dx, OLD_COV_INV);
  // df_du: Nothing to do, since COV(Ap) = zeros

  // MEAN:
  this->mean = this->mean + Ap;
}

/*---------------------------------------------------------------
              +=
 ---------------------------------------------------------------*/
void CPose3DPDFGaussianInf::operator+=(const CPose3DPDFGaussianInf& Ap)
{
  CPose3DPDFGaussian a(UNINITIALIZED_POSE);
  CPose3DPDFGaussian b(UNINITIALIZED_POSE);
  a.copyFrom(*this);
  b.copyFrom(Ap);

  a += b;

  this->mean = a.mean;
  cov_inv = a.cov.inverse_LLt();
}

/*---------------------------------------------------------------
              -=
 ---------------------------------------------------------------*/
void CPose3DPDFGaussianInf::operator-=(const CPose3DPDFGaussianInf& Ap)
{
  CPose3DPDFGaussian a(UNINITIALIZED_POSE);
  CPose3DPDFGaussian b(UNINITIALIZED_POSE);
  a.copyFrom(*this);
  b.copyFrom(Ap);

  a -= b;

  this->mean = a.mean;
  cov_inv = a.cov.inverse_LLt();
}

/*---------------------------------------------------------------
            evaluatePDF
 ---------------------------------------------------------------*/
double CPose3DPDFGaussianInf::evaluatePDF([[maybe_unused]] const CPose3D& x) const
{
  THROW_EXCEPTION("TO DO!!!");
}

/*---------------------------------------------------------------
            evaluateNormalizedPDF
 ---------------------------------------------------------------*/
double CPose3DPDFGaussianInf::evaluateNormalizedPDF([[maybe_unused]] const CPose3D& x) const
{
  THROW_EXCEPTION("TO DO!!!");
}

/*---------------------------------------------------------------
            enforceCovSymmetry
 ---------------------------------------------------------------*/
void CPose3DPDFGaussianInf::enforceCovSymmetry()
{
  // Differences, when they exist, appear in the ~15'th significant
  //  digit, so... just take one of them arbitrarily!
  for (int i = 0; i < cov_inv.rows() - 1; i++)
    for (int j = i + 1; j < cov_inv.rows(); j++) cov_inv(i, j) = cov_inv(j, i);
}

/*---------------------------------------------------------------
            mahalanobisDistanceTo
 ---------------------------------------------------------------*/
double CPose3DPDFGaussianInf::mahalanobisDistanceTo(const CPose3DPDFGaussianInf& theOther)
{
  MRPT_START

  const CMatrixDouble66 cov = this->cov_inv.inverse_LLt();
  const CMatrixDouble66 cov2 = theOther.cov_inv.inverse_LLt();

  CMatrixDouble66 COV_ = cov + cov2;
  CMatrixDouble61 MU = CMatrixDouble61(theOther.mean) - CMatrixDouble61(mean);

  for (int i = 0; i < 6; i++)
  {
    if (COV_(i, i) == 0)
    {
      if (MU(i, 0) != 0)
        return std::numeric_limits<double>::infinity();
      else
        COV_(i, i) = 1;  // Any arbitrary value since
                         // MU(i)=0, and this value doesn't
                         // affect the result.
    }
  }

  return std::sqrt(mrpt::math::multiply_HtCH_scalar(MU, COV_.inverse_LLt()));

  MRPT_END
}

/*---------------------------------------------------------------
            operator <<
 ---------------------------------------------------------------*/
ostream& mrpt::poses::operator<<(ostream& out, const CPose3DPDFGaussianInf& obj)
{
  out << "Mean: " << obj.mean << "\n";
  out << "Inverse cov:\n" << obj.cov_inv << "\n";

  return out;
}

/*---------------------------------------------------------------
            getInvCovSubmatrix2D
 ---------------------------------------------------------------*/
void CPose3DPDFGaussianInf::getInvCovSubmatrix2D(CMatrixDouble& out_cov) const
{
  out_cov.setSize(3, 3);

  for (int i = 0; i < 3; i++)
  {
    int a = i == 2 ? 3 : i;
    for (int j = i; j < 3; j++)
    {
      int b = j == 2 ? 3 : j;
      double f = cov_inv(a, b);
      out_cov(i, j) = f;
      out_cov(j, i) = f;
    }
  }
}

bool mrpt::poses::operator==(const CPose3DPDFGaussianInf& p1, const CPose3DPDFGaussianInf& p2)
{
  return p1.mean == p2.mean && p1.cov_inv == p2.cov_inv;
}

CPose3DPDFGaussianInf mrpt::poses::operator+(
    const CPose3DPDFGaussianInf& x, const CPose3DPDFGaussianInf& u)
{
  CPose3DPDFGaussianInf res(x);
  res += u;
  return res;
}

CPose3DPDFGaussianInf mrpt::poses::operator-(
    const CPose3DPDFGaussianInf& x, const CPose3DPDFGaussianInf& u)
{
  CPose3DPDFGaussianInf res(x);
  res -= u;
  return res;
}
