/*
 * Software License Agreement (BSD License)
 *
 *  Point Cloud Library (PCL) - www.pointclouds.org
 *  Copyright (c) 2010-2012, Willow Garage, Inc.
 *  Copyright (c) 2012-, Open Perception, Inc.
 *
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions
 *  are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials provided
 *     with the distribution.
 *   * Neither the name of the copyright holder(s) nor the names of its
 *     contributors may be used to endorse or promote products derived
 *     from this software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 *  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 *  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 *  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 *  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 *  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 *  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 *  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 *  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 *
 * $Id: lum.hpp 5663 2012-05-02 13:49:39Z florentinus $
 *
 */

#ifndef PCL_REGISTRATION_IMPL_LUM_HPP_
#define PCL_REGISTRATION_IMPL_LUM_HPP_

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> inline void
pcl::registration::LUM<PointT>::setLoopGraph (const SLAMGraphPtr &slam_graph)
{
  slam_graph_ = slam_graph;
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> inline typename pcl::registration::LUM<PointT>::SLAMGraphPtr
pcl::registration::LUM<PointT>::getLoopGraph () const
{
  return (slam_graph_);
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> typename pcl::registration::LUM<PointT>::SLAMGraph::vertices_size_type
pcl::registration::LUM<PointT>::getNumVertices () const
{
  return (num_vertices (*slam_graph_));
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> void
pcl::registration::LUM<PointT>::setMaxIterations (int max_iterations)
{
  max_iterations_ = max_iterations;
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> inline int
pcl::registration::LUM<PointT>::getMaxIterations () const
{
  return (max_iterations_);
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> void
pcl::registration::LUM<PointT>::setConvergenceThreshold (float convergence_threshold)
{
  convergence_threshold_ = convergence_threshold;
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> inline float
pcl::registration::LUM<PointT>::getConvergenceThreshold () const
{
  return (convergence_threshold_);
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> typename pcl::registration::LUM<PointT>::Vertex
pcl::registration::LUM<PointT>::addPointCloud (const PointCloudPtr &cloud, const Eigen::Vector6f &pose)
{
  Vertex v = add_vertex (*slam_graph_);
  (*slam_graph_)[v].cloud_ = cloud;
  if (v == 0 && pose != Eigen::Vector6f::Zero ())
  {
    PCL_WARN("[pcl::registration::LUM::addPointCloud] The pose estimate is ignored for the first cloud in the graph since that will become the reference pose.\n");
    (*slam_graph_)[v].pose_ = Eigen::Vector6f::Zero ();
    return (v);
  }
  (*slam_graph_)[v].pose_ = pose;
  return (v);
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> inline void
pcl::registration::LUM<PointT>::setPointCloud (const Vertex &vertex, const PointCloudPtr &cloud)
{
  if (vertex >= getNumVertices ())
  {
    PCL_ERROR("[pcl::registration::LUM::setPointCloud] You are attempting to set a point cloud to a non-existing graph vertex.\n");
    return;
  }
  (*slam_graph_)[vertex].cloud_ = cloud;
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> inline typename pcl::registration::LUM<PointT>::PointCloudPtr
pcl::registration::LUM<PointT>::getPointCloud (const Vertex &vertex) const
{
  if (vertex >= getNumVertices ())
  {
    PCL_ERROR("[pcl::registration::LUM::getPointCloud] You are attempting to get a point cloud from a non-existing graph vertex.\n");
    return (PointCloudPtr ());
  }
  return ((*slam_graph_)[vertex].cloud_);
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> inline void
pcl::registration::LUM<PointT>::setPose (const Vertex &vertex, const Eigen::Vector6f &pose)
{
  if (vertex >= getNumVertices ())
  {
    PCL_ERROR("[pcl::registration::LUM::setPose] You are attempting to set a pose estimate to a non-existing graph vertex.\n");
    return;
  }
  if (vertex == 0)
  {
    PCL_ERROR("[pcl::registration::LUM::setPose] The pose estimate is ignored for the first cloud in the graph since that will become the reference pose.\n");
    return;
  }
  (*slam_graph_)[vertex].pose_ = pose;
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> inline Eigen::Vector6f
pcl::registration::LUM<PointT>::getPose (const Vertex &vertex) const
{
  if (vertex >= getNumVertices ())
  {
    PCL_ERROR("[pcl::registration::LUM::getPose] You are attempting to get a pose estimate from a non-existing graph vertex.\n");
    return (Eigen::Vector6f::Zero ());
  }
  return ((*slam_graph_)[vertex].pose_);
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> inline Eigen::Affine3f
pcl::registration::LUM<PointT>::getTransformation (const Vertex &vertex) const
{
  Eigen::Vector6f pose = getPose (vertex);
  return (pcl::getTransformation (pose (0), pose (1), pose (2), pose (3), pose (4), pose (5)));
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> void
pcl::registration::LUM<PointT>::setCorrespondences (const Vertex &source_vertex, const Vertex &target_vertex, const pcl::CorrespondencesPtr &corrs)
{
  if (source_vertex >= getNumVertices () || target_vertex >= getNumVertices () || source_vertex == target_vertex)
  {
    PCL_ERROR("[pcl::registration::LUM::setCorrespondences] You are attempting to set a set of correspondences between non-existing or identical graph vertices.\n");
    return;
  }
  Edge e;
  bool present;
  boost::tuples::tie (e, present) = edge (source_vertex, target_vertex, *slam_graph_);
  if (!present)
    boost::tuples::tie (e, present) = add_edge (source_vertex, target_vertex, *slam_graph_);
  (*slam_graph_)[e].corrs_ = corrs;
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> inline pcl::CorrespondencesPtr
pcl::registration::LUM<PointT>::getCorrespondences (const Vertex &source_vertex, const Vertex &target_vertex) const
{
  if (source_vertex >= getNumVertices () || target_vertex >= getNumVertices ())
  {
    PCL_ERROR("[pcl::registration::LUM::getCorrespondences] You are attempting to get a set of correspondences between non-existing graph vertices.\n");
    return (pcl::CorrespondencesPtr ());
  }
  Edge e;
  bool present;
  boost::tuples::tie (e, present) = edge (source_vertex, target_vertex, *slam_graph_);
  if (!present)
  {
    PCL_ERROR("[pcl::registration::LUM::getCorrespondences] You are attempting to get a set of correspondences from a non-existing graph edge.\n");
    return (pcl::CorrespondencesPtr ());
  }
  return ((*slam_graph_)[e].corrs_);
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> void
pcl::registration::LUM<PointT>::compute ()
{
  int n = static_cast<int> (getNumVertices ());
  if (n < 2)
  {
    PCL_ERROR("[pcl::registration::LUM::compute] The slam graph needs at least 2 vertices.\n");
    return;
  }
  for (int i = 0; i < max_iterations_; ++i)
  {
    // Linearized computation of C^-1 and C^-1*D and convergence checking for all edges in the graph (results stored in slam_graph_)
    typename SLAMGraph::edge_iterator e, e_end;
    for (boost::tuples::tie (e, e_end) = edges (*slam_graph_); e != e_end; ++e)
      computeEdge (*e);

    // Declare matrices G and B
    Eigen::MatrixXf G = Eigen::MatrixXf::Zero (6 * (n - 1), 6 * (n - 1));
    Eigen::VectorXf B = Eigen::VectorXf::Zero (6 * (n - 1));

    // Start at 1 because 0 is the reference pose
    for (int vi = 1; vi != n; ++vi)
    {
      for (int vj = 0; vj != n; ++vj)
      {
        // Attempt to use the forward edge, otherwise use backward edge, otherwise there was no edge
        Edge e;
        bool present1, present2;
        boost::tuples::tie (e, present1) = edge (vi, vj, *slam_graph_);
        if (!present1)
        {
          boost::tuples::tie (e, present2) = edge (vj, vi, *slam_graph_);
          if (!present2)
            continue;
        }

        // Fill in elements of G and B
        if (vj > 0)
          G.block (6 * (vi - 1), 6 * (vj - 1), 6, 6) = -(*slam_graph_)[e].cinv_;
        G.block (6 * (vi - 1), 6 * (vi - 1), 6, 6) += (*slam_graph_)[e].cinv_;
        B.segment (6 * (vi - 1), 6) += (present1 ? 1 : -1) * (*slam_graph_)[e].cinvd_;
      }
    }

    // Computation of the linear equation system: GX = B
    // TODO investigate accuracy vs. speed tradeoff and find the best solving method for our type of linear equation (sparse)
    Eigen::VectorXf X = G.colPivHouseholderQr ().solve (B);

    // Update the poses
    float sum = 0.0;
    for (int vi = 1; vi != n; ++vi)
    {
      Eigen::Vector6f difference_pose = static_cast<Eigen::Vector6f> (-incidenceCorrection (getPose (vi)).inverse () * X.segment (6 * (vi - 1), 6));
      sum += difference_pose.norm ();
      setPose (vi, getPose (vi) + difference_pose);
    }

    // Convergence check
    if (sum <= convergence_threshold_ * static_cast<float> (n - 1))
      return;
  }
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> typename pcl::registration::LUM<PointT>::PointCloudPtr
pcl::registration::LUM<PointT>::getTransformedCloud (const Vertex &vertex) const
{
  PointCloudPtr out (new PointCloud);
  pcl::transformPointCloud (*getPointCloud (vertex), *out, getTransformation (vertex));
  return (out);
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> typename pcl::registration::LUM<PointT>::PointCloudPtr
pcl::registration::LUM<PointT>::getConcatenatedCloud () const
{
  PointCloudPtr out (new PointCloud);
  typename SLAMGraph::vertex_iterator v, v_end;
  for (boost::tuples::tie (v, v_end) = vertices (*slam_graph_); v != v_end; ++v)
  {
    PointCloud temp;
    pcl::transformPointCloud (*getPointCloud (*v), temp, getTransformation (*v));
    *out += temp;
  }
  return (out);
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> void
pcl::registration::LUM<PointT>::computeEdge (const Edge &e)
{
  // Get necessary local data from graph
  PointCloudPtr source_cloud = (*slam_graph_)[source (e, *slam_graph_)].cloud_;
  PointCloudPtr target_cloud = (*slam_graph_)[target (e, *slam_graph_)].cloud_;
  Eigen::Vector6f source_pose = (*slam_graph_)[source (e, *slam_graph_)].pose_;
  Eigen::Vector6f target_pose = (*slam_graph_)[target (e, *slam_graph_)].pose_;
  pcl::CorrespondencesPtr corrs = (*slam_graph_)[e].corrs_;

  // Build the average and difference vectors for all correspondences
  std::vector<Eigen::Vector3f, Eigen::aligned_allocator<Eigen::Vector3f> > corrs_aver (corrs->size ());
  std::vector<Eigen::Vector3f, Eigen::aligned_allocator<Eigen::Vector3f> > corrs_diff (corrs->size ());
  int oci = 0;  // oci = output correspondence iterator
  for (int ici = 0; ici != static_cast<int> (corrs->size ()); ++ici)  // ici = input correspondence iterator
  {
    // Compound the point pair onto the current pose
    Eigen::Vector3f source_compounded = pcl::getTransformation (source_pose (0), source_pose (1), source_pose (2), source_pose (3), source_pose (4), source_pose (5)) * source_cloud->points[(*corrs)[ici].index_query].getVector3fMap ();
    Eigen::Vector3f target_compounded = pcl::getTransformation (target_pose (0), target_pose (1), target_pose (2), target_pose (3), target_pose (4), target_pose (5)) * target_cloud->points[(*corrs)[ici].index_match].getVector3fMap ();

    // NaN points can not be passed to the remaining computational pipeline
    if (!pcl_isfinite (source_compounded (0)) || !pcl_isfinite (source_compounded (1)) || !pcl_isfinite (source_compounded (2)) || !pcl_isfinite (target_compounded (0)) || !pcl_isfinite (target_compounded (1)) || !pcl_isfinite (target_compounded (2)))
      continue;

    // Compute the point pair average and difference and store for later use
    corrs_aver[oci] = 0.5 * (source_compounded + target_compounded);
    corrs_diff[oci] = source_compounded - target_compounded;
    oci++;
  }
  corrs_aver.resize (oci);
  corrs_diff.resize (oci);

  // Need enough valid correspondences to get a good triangulation
  if (oci < 3)
  {
    PCL_WARN("[pcl::registration::LUM::compute] The correspondences between vertex %d and %d do not contain enough valid correspondences to be considered for computation.\n", source (e, *slam_graph_), target (e, *slam_graph_));
    (*slam_graph_)[e].cinv_ = Eigen::Matrix6f::Zero ();
    (*slam_graph_)[e].cinvd_ = Eigen::Vector6f::Zero ();
    return;
  }

  // Build the matrices M'M and M'Z
  Eigen::Matrix6f MM = Eigen::Matrix6f::Zero ();
  Eigen::Vector6f MZ = Eigen::Vector6f::Zero ();
  for (int ci = 0; ci != oci; ++ci)  // ci = correspondence iterator
  {
    // Fast computation of summation elements of M'M
    MM (0, 4) -= corrs_aver[ci] (1);
    MM (0, 5) += corrs_aver[ci] (2);
    MM (1, 3) -= corrs_aver[ci] (2);
    MM (1, 4) += corrs_aver[ci] (0);
    MM (2, 3) += corrs_aver[ci] (1);
    MM (2, 5) -= corrs_aver[ci] (0);
    MM (3, 4) -= corrs_aver[ci] (0) * corrs_aver[ci] (2);
    MM (3, 5) -= corrs_aver[ci] (0) * corrs_aver[ci] (1);
    MM (4, 5) -= corrs_aver[ci] (1) * corrs_aver[ci] (2);
    MM (3, 3) += corrs_aver[ci] (1) * corrs_aver[ci] (1) + corrs_aver[ci] (2) * corrs_aver[ci] (2);
    MM (4, 4) += corrs_aver[ci] (0) * corrs_aver[ci] (0) + corrs_aver[ci] (1) * corrs_aver[ci] (1);
    MM (5, 5) += corrs_aver[ci] (0) * corrs_aver[ci] (0) + corrs_aver[ci] (2) * corrs_aver[ci] (2);

    // Fast computation of M'Z
    MZ (0) += corrs_diff[ci] (0);
    MZ (1) += corrs_diff[ci] (1);
    MZ (2) += corrs_diff[ci] (2);
    MZ (3) += corrs_aver[ci] (1) * corrs_diff[ci] (2) - corrs_aver[ci] (2) * corrs_diff[ci] (1);
    MZ (4) += corrs_aver[ci] (0) * corrs_diff[ci] (1) - corrs_aver[ci] (1) * corrs_diff[ci] (0);
    MZ (5) += corrs_aver[ci] (2) * corrs_diff[ci] (0) - corrs_aver[ci] (0) * corrs_diff[ci] (2);
  }
  // Remaining elements of M'M
  MM (0, 0) = MM (1, 1) = MM (2, 2) = static_cast<float> (oci);
  MM (4, 0) = MM (0, 4);
  MM (5, 0) = MM (0, 5);
  MM (3, 1) = MM (1, 3);
  MM (4, 1) = MM (1, 4);
  MM (3, 2) = MM (2, 3);
  MM (5, 2) = MM (2, 5);
  MM (4, 3) = MM (3, 4);
  MM (5, 3) = MM (3, 5);
  MM (5, 4) = MM (4, 5);

  // Compute pose difference estimation
  Eigen::Vector6f D = static_cast<Eigen::Vector6f> (MM.inverse () * MZ);

  // Compute s^2
  float ss = 0.0f;
  for (int ci = 0; ci != oci; ++ci)  // ci = correspondence iterator
    ss += static_cast<float> (pow (corrs_diff[ci] (0) - (D (0) + corrs_aver[ci] (2) * D (5) - corrs_aver[ci] (1) * D (4)), 2.0f)
                            + pow (corrs_diff[ci] (1) - (D (1) + corrs_aver[ci] (0) * D (4) - corrs_aver[ci] (2) * D (3)), 2.0f)
                            + pow (corrs_diff[ci] (2) - (D (2) + corrs_aver[ci] (1) * D (3) - corrs_aver[ci] (0) * D (5)), 2.0f));

  // When reaching the limitations of computation due to linearization
  if (ss < 0.0000000000001 || !pcl_isfinite (ss))
  {
    (*slam_graph_)[e].cinv_ = Eigen::Matrix6f::Zero ();
    (*slam_graph_)[e].cinvd_ = Eigen::Vector6f::Zero ();
    return;
  }

  // Store the results in the slam graph
  (*slam_graph_)[e].cinv_ = MM * (1.0f / ss);
  (*slam_graph_)[e].cinvd_ = MZ * (1.0f / ss);
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename PointT> inline Eigen::Matrix6f
pcl::registration::LUM<PointT>::incidenceCorrection (const Eigen::Vector6f &pose)
{
  Eigen::Matrix6f out = Eigen::Matrix6f::Identity ();
  float cx = cosf (pose (3)), sx = sinf (pose (3)), cy = cosf (pose (4)), sy = sinf (pose (4));
  out (0, 4) = pose (1) * sx - pose (2) * cx;
  out (0, 5) = pose (1) * cx * cy + pose (2) * sx * cy;
  out (1, 3) = pose (2);
  out (1, 4) = -pose (0) * sx;
  out (1, 5) = -pose (0) * cx * cy + pose (2) * sy;
  out (2, 3) = -pose (1);
  out (2, 4) = pose (0) * cx;
  out (2, 5) = -pose (0) * sx * cy - pose (1) * sy;
  out (3, 5) = sy;
  out (4, 4) = sx;
  out (4, 5) = cx * cy;
  out (5, 4) = cx;
  out (5, 5) = -sx * cy;
  return (out);
}

#define PCL_INSTANTIATE_LUM(T) template class PCL_EXPORTS pcl::registration::LUM<T>;

#endif  // PCL_REGISTRATION_IMPL_LUM_HPP_

