/**
 * Copyright (C) 2010  Hauke Strasdat
 *                     Imperial College London
 * Copyright (c) 2005-2024, Individual contributors, see AUTHORS file
 *   See: https://www.mrpt.org/Authors - All rights reserved.
 *
 * bundle_adjuster.h is part of RobotVision.
 *
 * RobotVision is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or any later version.
 *
 * RobotVision is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * and the GNU Lesser General Public License along with this program.
 * If not, see <http://www.gnu.org/licenses/>.
 */
#pragma once

#include <mrpt/math/CMatrixFixed.h>
#include <mrpt/math/CVectorFixed.h>
#include <mrpt/poses/CPose3D.h>
#include <mrpt/vision/types.h>

#include <Eigen/Dense>
#include <array>
#include <vector>

// Declarations shared between ba_*.cpp files, but which are private to MRPT
//  not to be seen by an MRPT API user.

namespace mrpt::vision
{
using mrpt::math::CVectorFixedDouble;  // Allow these "using"s since these
                                       // headers are
// internal to mrpt
using mrpt::math::CMatrixFixed;
using std::vector;

#define VERBOSE_COUT \
  if (verbose) cout

// Auxiliary struct for keeping the list of all Jacobians in a sparse, efficient
// way.
template <int FrameDof, int PointDof, int ObsDim>
struct JacData
{
  inline JacData() :
      J_frame(mrpt::math::UNINITIALIZED_MATRIX),
      J_point(mrpt::math::UNINITIALIZED_MATRIX),
      J_frame_valid(false),
      J_point_valid(false)
  {
  }

  TCameraPoseID frame_id;
  TLandmarkID point_id;

  // Jacobians of the observation wrt the camera pose & the point:
  mrpt::math::CMatrixFixed<double, ObsDim, FrameDof> J_frame;
  mrpt::math::CMatrixFixed<double, ObsDim, PointDof> J_point;
  bool J_frame_valid, J_point_valid;

  // Needed by any struct having
  // Eigen::Matrix<> fields
};

/** The projective camera 2x6 Jacobian \f$ \frac{\partial h(f,p)}{\partial p}
 * \f$ (wrt the 6D camera pose)
 * \note Jacobians as described in
 * https://www.mrpt.org/6D_poses:equivalences_compositions_and_uncertainty
 * (Appendix A)
 */
template <bool POSES_ARE_INVERSE>
void frameJac(
    const mrpt::img::TCamera& camera_params,
    const mrpt::poses::CPose3D& cam_pose,
    const mrpt::math::TPoint3D& landmark_global,
    mrpt::math::CMatrixFixed<double, 2, 6>& out_J)
{
  using mrpt::square;

  double x, y, z;  // wrt cam (local coords)
  if (POSES_ARE_INVERSE)
    cam_pose.composePoint(landmark_global.x, landmark_global.y, landmark_global.z, x, y, z);
  else
    cam_pose.inverseComposePoint(landmark_global.x, landmark_global.y, landmark_global.z, x, y, z);

  ASSERT_(z != 0);
  const double _z = 1.0 / z;
  const double fx_z = camera_params.fx() * _z;
  const double fy_z = camera_params.fy() * _z;
  const double _z2 = square(_z);
  const double fx_z2 = camera_params.fx() * _z2;
  const double fy_z2 = camera_params.fy() * _z2;

  if (POSES_ARE_INVERSE)
  {
    const double xy = x * y;

    const double J_vals[] = {
        fx_z,       0,       -fx_z2 * x, -fx_z2 * xy, camera_params.fx() * (1 + square(x * _z)),
        -fx_z * y,  0,       fy_z,       -fy_z2 * y,  -camera_params.fy() * (1 + square(y * _z)),
        fy_z2 * xy, fy_z * x};
    out_J.loadFromArray(J_vals);
  }
  else
  {
    const mrpt::math::CMatrixDouble33& R = cam_pose.getRotationMatrix();

    const double jac_proj_vals[] = {fx_z, 0, -fx_z2 * x, 0, fy_z, -fy_z2 * y};
    const mrpt::math::CMatrixFixed<double, 2, 3> jac_proj(jac_proj_vals);

    const double p_x = cam_pose.x();
    const double p_y = cam_pose.y();
    const double p_z = cam_pose.z();
    const double tx = landmark_global.x - p_x;
    const double ty = landmark_global.y - p_y;
    const double tz = landmark_global.z - p_z;

    const double aux_vals[] = {
        -R(0, 0),
        -R(1, 0),
        -R(2, 0),
        tz * R(1, 0) - ty * R(2, 0) + R(1, 0) * p_z - R(2, 0) * p_y,
        tx * R(2, 0) - tz * R(0, 0) - R(0, 0) * p_z + R(2, 0) * p_x,
        ty * R(0, 0) - tx * R(1, 0) + R(0, 0) * p_y - R(1, 0) * p_x,
        -R(0, 1),
        -R(1, 1),
        -R(2, 1),
        tz * R(1, 1) - ty * R(2, 1) + R(1, 1) * p_z - R(2, 1) * p_y,
        tx * R(2, 1) - tz * R(0, 1) - R(0, 1) * p_z + R(2, 1) * p_x,
        ty * R(0, 1) - tx * R(1, 1) + R(0, 1) * p_y - R(1, 1) * p_x,
        -R(0, 2),
        -R(1, 2),
        -R(2, 2),
        tz * R(1, 2) - ty * R(2, 2) + R(1, 2) * p_z - R(2, 2) * p_y,
        tx * R(2, 2) - tz * R(0, 2) - R(0, 2) * p_z + R(2, 2) * p_x,
        ty * R(0, 2) - tx * R(1, 2) + R(0, 2) * p_y - R(1, 2) * p_x};
    const mrpt::math::CMatrixFixed<double, 3, 6> vals(aux_vals);
    out_J = jac_proj * vals;
  }
}

/** Jacobians wrt the point
 * \note Jacobians as described in
 * https://www.mrpt.org/6D_poses:equivalences_compositions_and_uncertainty
 * (Appendix A)
 */
template <bool POSES_ARE_INVERSE>
void pointJac(
    const mrpt::img::TCamera& camera_params,
    const mrpt::poses::CPose3D& cam_pose,
    const mrpt::math::TPoint3D& landmark_global,
    mrpt::math::CMatrixFixed<double, 2, 3>& out_J)
{
  using namespace mrpt::math;

  mrpt::math::TPoint3D l;  // Local point, wrt camera

  mrpt::math::CMatrixDouble33 dp_point(mrpt::math::UNINITIALIZED_MATRIX);

  if (POSES_ARE_INVERSE)
    cam_pose.composePoint(
        landmark_global.x, landmark_global.y, landmark_global.z, l.x, l.y, l.z, dp_point);
  else
    cam_pose.inverseComposePoint(
        landmark_global.x, landmark_global.y, landmark_global.z, l.x, l.y, l.z, dp_point);

  ASSERT_(l.z != 0);
  const double _z = 1.0 / l.z;
  const double _z2 = square(_z);

  const double tmp_vals[] = {camera_params.fx() * _z,         0,
                             -camera_params.fx() * l.x * _z2, 0,
                             camera_params.fy() * _z,         -camera_params.fy() * l.y * _z2};
  const mrpt::math::CMatrixFixed<double, 2, 3> tmp(tmp_vals);

  out_J = tmp * dp_point;
}

// === Compute sparse Jacobians ====
// Case: 6D poses + 3D points + 2D (x,y) observations
// For the case of *inverse* or *normal* frame poses being estimated.
// Made inline so immediate values in "poses_are_inverses" are propragated by
// the compiler
template <bool POSES_ARE_INVERSE>
void ba_compute_Jacobians(
    const TFramePosesVec& frame_poses,
    const TLandmarkLocationsVec& landmark_points,
    const mrpt::img::TCamera& camera_params,
    std::vector<JacData<6, 3, 2>>& jac_data_vec,
    size_t num_fix_frames,
    size_t num_fix_points)
{
  MRPT_START

  // num_fix_frames & num_fix_points: Are relative to the order in frame_poses
  // & landmark_points
  ASSERT_(!frame_poses.empty() && !landmark_points.empty());
  const size_t N = jac_data_vec.size();

  for (size_t i = 0; i < N; i++)
  {
    JacData<6, 3, 2>& D = jac_data_vec[i];

    const TCameraPoseID i_f = D.frame_id;
    const TLandmarkID i_p = D.point_id;

    ASSERTDEB_(i_f < frame_poses.size());
    ASSERTDEB_(i_p < landmark_points.size());

    if (i_f >= num_fix_frames)
    {
      frameJac<POSES_ARE_INVERSE>(camera_params, frame_poses[i_f], landmark_points[i_p], D.J_frame);
      D.J_frame_valid = true;
    }

    if (i_p >= num_fix_points)
    {
      pointJac<POSES_ARE_INVERSE>(camera_params, frame_poses[i_f], landmark_points[i_p], D.J_point);
      D.J_point_valid = true;
    }
  }
  MRPT_END
}

/** Construct the BA linear system.
 *  Set kernel_1st_deriv!=nullptr if using robust kernel.
 */
void ba_build_gradient_Hessians(
    const TSequenceFeatureObservations& observations,
    const std::vector<std::array<double, 2>>& residual_vec,
    const std::vector<JacData<6, 3, 2>>& jac_data_vec,
    std::vector<mrpt::math::CMatrixFixed<double, 6, 6>>& U,
    std::vector<CVectorFixedDouble<6>>& eps_frame,
    std::vector<mrpt::math::CMatrixFixed<double, 3, 3>>& V,
    std::vector<CVectorFixedDouble<3>>& eps_point,
    size_t num_fix_frames,
    size_t num_fix_points,
    const vector<double>* kernel_1st_deriv);
}  // namespace mrpt::vision
