/*
 * Software License Agreement (BSD License)
 *
 * Point Cloud Library (PCL) - www.pointclouds.org
 * Copyright (c) 2009-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.
 *
 */

#ifndef PCL_REGISTRATION_IMPL_JOINT_ICP_HPP_
#define PCL_REGISTRATION_IMPL_JOINT_ICP_HPP_

#include <pcl/registration/boost.h>
#include <pcl/correspondence.h>
#include <pcl/console/print.h>


namespace pcl
{

template <typename PointSource, typename PointTarget, typename Scalar> void
JointIterativeClosestPoint<PointSource, PointTarget, Scalar>::computeTransformation (
    PointCloudSource &output, const Matrix4 &guess)
{
  // Point clouds containing the correspondences of each point in <input, indices>
  if (sources_.size () != targets_.size () || sources_.empty () || targets_.empty ())
  {
    PCL_ERROR ("[pcl::%s::computeTransformation] Must set InputSources and InputTargets to the same, nonzero size!\n", 
        getClassName ().c_str ());
    return;
  }
  bool manual_correspondence_estimations_set = true;
  if (correspondence_estimations_.empty ())
  {
    manual_correspondence_estimations_set = false;
    correspondence_estimations_.resize (sources_.size ());
    for (std::size_t i = 0; i < sources_.size (); i++)
    {
      correspondence_estimations_[i] = correspondence_estimation_->clone ();
      KdTreeReciprocalPtr src_tree (new KdTreeReciprocal);
      KdTreePtr tgt_tree (new KdTree);
      correspondence_estimations_[i]->setSearchMethodTarget (tgt_tree);
      correspondence_estimations_[i]->setSearchMethodSource (src_tree);
    }
  }
  if (correspondence_estimations_.size () != sources_.size ())
  {
    PCL_ERROR ("[pcl::%s::computeTransform] Must set CorrespondenceEstimations to be the same size as the joint\n", 
        getClassName ().c_str ());
    return;
  }
  std::vector<PointCloudSourcePtr> inputs_transformed (sources_.size ());
  for (std::size_t i = 0; i < sources_.size (); i++)
  {
    inputs_transformed[i].reset (new PointCloudSource);
  }

  nr_iterations_ = 0;
  converged_ = false;

  // Initialise final transformation to the guessed one
  final_transformation_ = guess;

  // Make a combined transformed input and output
  std::vector<std::size_t> input_offsets (sources_.size ());
  std::vector<std::size_t> target_offsets (targets_.size ());
  PointCloudSourcePtr sources_combined (new PointCloudSource);
  PointCloudSourcePtr inputs_transformed_combined (new PointCloudSource);
  PointCloudTargetPtr targets_combined (new PointCloudTarget);
  std::size_t input_offset = 0;
  std::size_t target_offset = 0;
  for (std::size_t i = 0; i < sources_.size (); i++)
  {
    // If the guessed transformation is non identity
    if (guess != Matrix4::Identity ())
    {
       // Apply guessed transformation prior to search for neighbours
      this->transformCloud (*sources_[i], *inputs_transformed[i], guess);
    }
    else
    {
      *inputs_transformed[i] = *sources_[i];
    }
    *sources_combined += *sources_[i];
    *inputs_transformed_combined += *inputs_transformed[i];
    *targets_combined += *targets_[i];
    input_offsets[i] = input_offset;
    target_offsets[i] = target_offset;
    input_offset += inputs_transformed[i]->size ();
    target_offset += targets_[i]->size ();
  }

  transformation_ = Matrix4::Identity ();
  // Make blobs if necessary
  determineRequiredBlobData ();
  // Pass in the default target for the Correspondence Estimation/Rejection code
  for (std::size_t i = 0; i < sources_.size (); i++)
  {
    correspondence_estimations_[i]->setInputTarget (targets_[i]);
    if (correspondence_estimations_[i]->requiresTargetNormals ())
    {
      PCLPointCloud2::Ptr target_blob (new PCLPointCloud2);
      pcl::toPCLPointCloud2 (*targets_[i], *target_blob);
      correspondence_estimations_[i]->setTargetNormals (target_blob);
    }
  }

  PCLPointCloud2::Ptr targets_combined_blob (new PCLPointCloud2);
  if (!correspondence_rejectors_.empty () && need_target_blob_)
    pcl::toPCLPointCloud2 (*targets_combined, *targets_combined_blob);

  for (std::size_t i = 0; i < correspondence_rejectors_.size (); ++i)
  {
    registration::CorrespondenceRejector::Ptr& rej = correspondence_rejectors_[i];
    if (rej->requiresTargetPoints ())
      rej->setTargetPoints (targets_combined_blob);
    if (rej->requiresTargetNormals () && target_has_normals_)
      rej->setTargetNormals (targets_combined_blob);
  }

  convergence_criteria_->setMaximumIterations (max_iterations_);
  convergence_criteria_->setRelativeMSE (euclidean_fitness_epsilon_);
  convergence_criteria_->setTranslationThreshold (transformation_epsilon_);
  convergence_criteria_->setRotationThreshold (1.0 - transformation_epsilon_);

  // Repeat until convergence
  std::vector<CorrespondencesPtr> partial_correspondences_ (sources_.size ());
  for (std::size_t i = 0; i < sources_.size (); i++)
  {
    partial_correspondences_[i].reset (new pcl::Correspondences);
  }

  do
  {
    // Save the previously estimated transformation
    previous_transformation_ = transformation_;

    // Set the source each iteration, to ensure the dirty flag is updated
    correspondences_->clear ();
    for (std::size_t i = 0; i < correspondence_estimations_.size (); i++)
    {
      correspondence_estimations_[i]->setInputSource (inputs_transformed[i]);
      // Get blob data if needed
      if (correspondence_estimations_[i]->requiresSourceNormals ())
      {
        PCLPointCloud2::Ptr input_transformed_blob (new PCLPointCloud2);
        toPCLPointCloud2 (*inputs_transformed[i], *input_transformed_blob);
        correspondence_estimations_[i]->setSourceNormals (input_transformed_blob);
      }

      // Estimate correspondences on each cloud pair separately
      if (use_reciprocal_correspondence_)
      {
        correspondence_estimations_[i]->determineReciprocalCorrespondences (*partial_correspondences_[i], corr_dist_threshold_);
      }
      else
      {
        correspondence_estimations_[i]->determineCorrespondences (*partial_correspondences_[i], corr_dist_threshold_);
      }
      PCL_DEBUG ("[pcl::%s::computeTransformation] Found %d partial correspondences for cloud [%d]\n",
          getClassName ().c_str (),
          partial_correspondences_[i]->size (), i);
      for (std::size_t j = 0; j < partial_correspondences_[i]->size (); j++)
      {
        pcl::Correspondence corr = partial_correspondences_[i]->at (j);
        // Update the offsets to be for the combined clouds
        corr.index_query += input_offsets[i];
        corr.index_match += target_offsets[i];
        correspondences_->push_back (corr);
      }
    }
    PCL_DEBUG ("[pcl::%s::computeTransformation] Total correspondences: %d\n", getClassName ().c_str (), correspondences_->size ());

    PCLPointCloud2::Ptr inputs_transformed_combined_blob;
    if (need_source_blob_)
    {
      inputs_transformed_combined_blob.reset (new PCLPointCloud2);
      toPCLPointCloud2 (*inputs_transformed_combined, *inputs_transformed_combined_blob);
    }
    CorrespondencesPtr temp_correspondences (new Correspondences (*correspondences_));
    for (std::size_t i = 0; i < correspondence_rejectors_.size (); ++i)
    {
      PCL_DEBUG ("Applying a correspondence rejector method: %s.\n", correspondence_rejectors_[i]->getClassName ().c_str ());
      registration::CorrespondenceRejector::Ptr& rej = correspondence_rejectors_[i];
      PCL_DEBUG ("Applying a correspondence rejector method: %s.\n", rej->getClassName ().c_str ());
      if (rej->requiresSourcePoints ())
        rej->setSourcePoints (inputs_transformed_combined_blob);
      if (rej->requiresSourceNormals () && source_has_normals_)
        rej->setSourceNormals (inputs_transformed_combined_blob);
      rej->setInputCorrespondences (temp_correspondences);
      rej->getCorrespondences (*correspondences_);
      // Modify input for the next iteration
      if (i < correspondence_rejectors_.size () - 1)
        *temp_correspondences = *correspondences_;
    }

    int cnt = correspondences_->size ();
    // Check whether we have enough correspondences
    if (cnt < min_number_correspondences_)
    {
      PCL_ERROR ("[pcl::%s::computeTransformation] Not enough correspondences found. Relax your threshold parameters.\n", getClassName ().c_str ());
      convergence_criteria_->setConvergenceState(pcl::registration::DefaultConvergenceCriteria<Scalar>::CONVERGENCE_CRITERIA_NO_CORRESPONDENCES);
      converged_ = false;
      break;
    }

    // Estimate the transform jointly, on a combined correspondence set
    transformation_estimation_->estimateRigidTransformation (*inputs_transformed_combined, *targets_combined, *correspondences_, transformation_);

    // Transform the combined data
    this->transformCloud (*inputs_transformed_combined, *inputs_transformed_combined, transformation_);
    // And all its components
    for (std::size_t i = 0; i < sources_.size (); i++)
    {
      this->transformCloud (*inputs_transformed[i], *inputs_transformed[i], transformation_);
    }

    // Obtain the final transformation
    final_transformation_ = transformation_ * final_transformation_;

    ++nr_iterations_;

    // Update the vizualization of icp convergence
    //if (update_visualizer_ != 0)
    //  update_visualizer_(output, source_indices_good, *target_, target_indices_good );

    converged_ = static_cast<bool> ((*convergence_criteria_));
  }
  while (!converged_);

  PCL_DEBUG ("Transformation is:\n\t%5f\t%5f\t%5f\t%5f\n\t%5f\t%5f\t%5f\t%5f\n\t%5f\t%5f\t%5f\t%5f\n\t%5f\t%5f\t%5f\t%5f\n",
      final_transformation_ (0, 0), final_transformation_ (0, 1), final_transformation_ (0, 2), final_transformation_ (0, 3),
      final_transformation_ (1, 0), final_transformation_ (1, 1), final_transformation_ (1, 2), final_transformation_ (1, 3),
      final_transformation_ (2, 0), final_transformation_ (2, 1), final_transformation_ (2, 2), final_transformation_ (2, 3),
      final_transformation_ (3, 0), final_transformation_ (3, 1), final_transformation_ (3, 2), final_transformation_ (3, 3));

  // For fitness checks, etc, we'll use an aggregated cloud for now (should be evaluating independently for correctness, but this requires propagating a few virtual methods from Registration)
  IterativeClosestPoint<PointSource, PointTarget, Scalar>::setInputSource (sources_combined);
  IterativeClosestPoint<PointSource, PointTarget, Scalar>::setInputTarget (targets_combined);

  // If we automatically set the correspondence estimators, we should clear them now
  if (!manual_correspondence_estimations_set)
  {
    correspondence_estimations_.clear ();
  }


  // By definition, this method will return an empty cloud (for compliance with the ICP API).
  // We can figure out a better solution, if necessary.
  output = PointCloudSource ();
}


template <typename PointSource, typename PointTarget, typename Scalar> void
JointIterativeClosestPoint<PointSource, PointTarget, Scalar>::determineRequiredBlobData ()
{
  need_source_blob_ = false;
  need_target_blob_ = false;
  // Check estimators
  for (std::size_t i = 0; i < correspondence_estimations_.size (); i++)
  {
    CorrespondenceEstimationPtr& ce = correspondence_estimations_[i];

    need_source_blob_ |= ce->requiresSourceNormals ();
    need_target_blob_ |= ce->requiresTargetNormals ();
    // Add warnings if necessary
    if (ce->requiresSourceNormals () && !source_has_normals_)
    {
        PCL_WARN("[pcl::%s::determineRequiredBlobData] Estimator expects source normals, but we can't provide them.\n", getClassName ().c_str ());
    }
    if (ce->requiresTargetNormals () && !target_has_normals_)
    {
        PCL_WARN("[pcl::%s::determineRequiredBlobData] Estimator expects target normals, but we can't provide them.\n", getClassName ().c_str ());
    }
  }
  // Check rejectors
  for (std::size_t i = 0; i < correspondence_rejectors_.size (); i++)
  {
    registration::CorrespondenceRejector::Ptr& rej = correspondence_rejectors_[i];
    need_source_blob_ |= rej->requiresSourcePoints ();
    need_source_blob_ |= rej->requiresSourceNormals ();
    need_target_blob_ |= rej->requiresTargetPoints ();
    need_target_blob_ |= rej->requiresTargetNormals ();
    if (rej->requiresSourceNormals () && !source_has_normals_)
    {
      PCL_WARN("[pcl::%s::determineRequiredBlobData] Rejector %s expects source normals, but we can't provide them.\n", getClassName ().c_str (), rej->getClassName ().c_str ());
    }
    if (rej->requiresTargetNormals () && !target_has_normals_)
    {
      PCL_WARN("[pcl::%s::determineRequiredBlobData] Rejector %s expects target normals, but we can't provide them.\n", getClassName ().c_str (), rej->getClassName ().c_str ());
    }
  }
}

} // namespace pcl

#endif /* PCL_REGISTRATION_IMPL_JOINT_ICP_HPP_ */

