// -*- mode: C++ -*-
/*********************************************************************
 * Software License Agreement (BSD License)
 *
 *  Copyright (c) 2013, Yuto Inagaki and JSK Lab
 *  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/o2r other materials provided
 *     with the distribution.
 *   * Neither the name of the Willow Garage 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 JSK_PCL_ROS_PARTICLE_FILTER_TRACKING_H_
#define JSK_PCL_ROS_PARTICLE_FILTER_TRACKING_H_

// ros
#include <ros/ros.h>
#include <ros/names.h>
#include <sensor_msgs/PointCloud2.h>
#include <tf/transform_broadcaster.h>
#include <tf_conversions/tf_eigen.h>
// pcl
#include <pcl_ros/pcl_nodelet.h>
#include <pcl/point_types.h>
#include <pcl/common/centroid.h>
#include <pcl/common/transforms.h>

#include <pcl/search/pcl_search.h>
#include <pcl/common/transforms.h>

#include <pcl/tracking/tracking.h>
#include <pcl/tracking/particle_filter.h>
#include <pcl/tracking/kld_adaptive_particle_filter_omp.h>
#include <pcl/tracking/particle_filter_omp.h>
#include <pcl/tracking/coherence.h>
#include <pcl/tracking/distance_coherence.h>
#include <pcl/tracking/hsv_color_coherence.h>
#include <pcl/tracking/approx_nearest_pair_point_cloud_coherence.h>
#include <pcl/tracking/nearest_pair_point_cloud_coherence.h>

#include <jsk_pcl_ros/SetPointCloud2.h>
#include <jsk_pcl_ros/ParticleFilterTrackingConfig.h>
#include <jsk_recognition_msgs/BoundingBox.h>

#include <dynamic_reconfigure/server.h>
#include <message_filters/subscriber.h>
#include <message_filters/time_synchronizer.h>
#include <message_filters/synchronizer.h>

#include <pcl/tracking/particle_filter.h>
#include <pcl/tracking/impl/tracking.hpp>
#include <pcl/tracking/impl/particle_filter.hpp>

// This namespace follows PCL coding style
namespace pcl
{
  namespace tracking
  {
    // hack pcl::tracking
    // original tracker assumes that the number of reference points is smaller
    // than the number of input.
    // ReversedParticleFilterTracker assumues that the number of reference
    // points is greater than the number of input.
    // So we need to change:
    //   1) transform input pointcloud during weight phase with inverse of each particles
    //      particle should keep meaning the pose of reference for simplicity.
    template <typename PointInT, typename StateT>
    class ReversedParticleFilterTracker: public ParticleFilterTracker<PointInT, StateT>
    {
    public:
      using Tracker<PointInT, StateT>::tracker_name_;
      using Tracker<PointInT, StateT>::search_;
      using Tracker<PointInT, StateT>::input_;
      using Tracker<PointInT, StateT>::indices_;
      using Tracker<PointInT, StateT>::getClassName;
      using ParticleFilterTracker<PointInT, StateT>::particles_;
      using ParticleFilterTracker<PointInT, StateT>::change_detector_;
      using ParticleFilterTracker<PointInT, StateT>::change_counter_;
      using ParticleFilterTracker<PointInT, StateT>::change_detector_resolution_;
      using ParticleFilterTracker<PointInT, StateT>::change_detector_interval_;
      using ParticleFilterTracker<PointInT, StateT>::use_change_detector_;
      using ParticleFilterTracker<PointInT, StateT>::alpha_;
      using ParticleFilterTracker<PointInT, StateT>::changed_;
      using ParticleFilterTracker<PointInT, StateT>::ref_;
      using ParticleFilterTracker<PointInT, StateT>::coherence_;
      using ParticleFilterTracker<PointInT, StateT>::use_normal_;
      using ParticleFilterTracker<PointInT, StateT>::particle_num_;
      using ParticleFilterTracker<PointInT, StateT>::change_detector_filter_;
      using ParticleFilterTracker<PointInT, StateT>::transed_reference_vector_;
      //using ParticleFilterTracker<PointInT, StateT>::calcLikelihood;
      using ParticleFilterTracker<PointInT, StateT>::normalizeWeight;
      using ParticleFilterTracker<PointInT, StateT>::initParticles;
      using ParticleFilterTracker<PointInT, StateT>::normalizeParticleWeight;
      using ParticleFilterTracker<PointInT, StateT>::calcBoundingBox;

      typedef Tracker<PointInT, StateT> BaseClass;
      
      typedef typename Tracker<PointInT, StateT>::PointCloudIn PointCloudIn;
      typedef typename PointCloudIn::Ptr PointCloudInPtr;
      typedef typename PointCloudIn::ConstPtr PointCloudInConstPtr;

      typedef typename Tracker<PointInT, StateT>::PointCloudState PointCloudState;
      typedef typename PointCloudState::Ptr PointCloudStatePtr;
      typedef typename PointCloudState::ConstPtr PointCloudStateConstPtr;

      typedef PointCoherence<PointInT> Coherence;
      typedef boost::shared_ptr< Coherence > CoherencePtr;
      typedef boost::shared_ptr< const Coherence > CoherenceConstPtr;

      typedef PointCloudCoherence<PointInT> CloudCoherence;
      typedef boost::shared_ptr< CloudCoherence > CloudCoherencePtr;
      typedef boost::shared_ptr< const CloudCoherence > CloudCoherenceConstPtr;

      // call initCompute only if reference pointcloud is updated
      inline void
      setReferenceCloud (const PointCloudInConstPtr &ref)
      {
        ref_ = ref;
        if (coherence_) {
          coherence_->setTargetCloud (ref_);
          coherence_->initCompute ();
        }
        else {
          PCL_ERROR("coherence_ is not yet available!");
        }
      }


    protected:
      std::vector<PointCloudInPtr> transed_input_vector_;
      
      virtual bool initCompute()
      {
        if (!Tracker<PointInT, StateT>::initCompute ())
        {
          PCL_ERROR ("[pcl::%s::initCompute] Init failed.\n", getClassName ().c_str ());
          return (false);
        }
        
        // allocate pointclouds in transed_input_vector_ instead of transed_reference_vector_
        if (transed_input_vector_.empty ())
        {
          //std::cout << "initializing " << particle_num_ << " input" << std::endl;
          // only one time allocation
          transed_input_vector_.resize (particle_num_ + 1);
          for (int i = 0; i < particle_num_ + 1; i++)
          {
            transed_input_vector_[i] = PointCloudInPtr (new PointCloudIn ());
          }
        }

        // set reference instead of input
        // if (coherence_) {
        //   coherence_->setTargetCloud (ref_);
        // }
          
        if (!change_detector_)
          change_detector_ = boost::shared_ptr<pcl::octree::OctreePointCloudChangeDetector<PointInT> >(new pcl::octree::OctreePointCloudChangeDetector<PointInT> (change_detector_resolution_));
          
        if (!particles_ || particles_->points.empty ())
          initParticles (true);
        return (true);
      }

      // only is computation without normal supported
      void computeTransformedPointCloudWithoutNormal
      (const StateT& hypothesis, PointCloudIn &cloud)
        {
          const Eigen::Affine3f trans = this->toEigenMatrix (hypothesis);
          // destructively assigns to cloud
          pcl::transformPointCloud<PointInT> (*input_, cloud, trans);
        }

      
      virtual void weight()
      {
        changed_ = true;
        if (!use_normal_)
        {
          for (size_t i = 0; i < particles_->points.size (); i++)
          {
            //std::cout << "processing " << i << " particle: " << particles_->points[i].weight << std::endl;
            // compute `inverse` of particle
            StateT inverse_particle;
            Eigen::Affine3f trans = particles_->points[i].toEigenMatrix();
            Eigen::Affine3f inverse_trans = trans.inverse();
            inverse_particle = StateT::toState(inverse_trans);
            computeTransformedPointCloudWithoutNormal (inverse_particle, *transed_input_vector_[i]);
            IndicesPtr indices;
            coherence_->compute (transed_input_vector_[i], indices, particles_->points[i].weight);
          }
        }
        else
        {
          for (size_t i = 0; i < particles_->points.size (); i++)
          {
            StateT inverse_particle;
            Eigen::Affine3f trans = particles_->points[i].toEigenMatrix();
            Eigen::Affine3f inverse_trans = trans.inverse();
            inverse_particle = StateT::toState(inverse_trans);
            IndicesPtr indices (new std::vector<int>);
            this->computeTransformedPointCloudWithNormal (inverse_particle, *indices, *transed_input_vector_[i]);
            coherence_->compute (transed_input_vector_[i], indices, particles_->points[i].weight);
          }
        }
  
        normalizeWeight ();
        
      }
    private:
    };
    
    template <typename PointInT, typename StateT>
    class ReversedParticleFilterOMPTracker: public ReversedParticleFilterTracker<PointInT, StateT>
    {
    public:
      using Tracker<PointInT, StateT>::tracker_name_;
      using Tracker<PointInT, StateT>::search_;
      using Tracker<PointInT, StateT>::input_;
      using Tracker<PointInT, StateT>::indices_;
      using Tracker<PointInT, StateT>::getClassName;
      using ParticleFilterTracker<PointInT, StateT>::particles_;
      using ParticleFilterTracker<PointInT, StateT>::change_detector_;
      using ParticleFilterTracker<PointInT, StateT>::change_counter_;
      using ParticleFilterTracker<PointInT, StateT>::change_detector_resolution_;
      using ParticleFilterTracker<PointInT, StateT>::change_detector_interval_;
      using ParticleFilterTracker<PointInT, StateT>::use_change_detector_;
      using ParticleFilterTracker<PointInT, StateT>::alpha_;
      using ParticleFilterTracker<PointInT, StateT>::changed_;
      using ParticleFilterTracker<PointInT, StateT>::ref_;
      using ParticleFilterTracker<PointInT, StateT>::coherence_;
      using ParticleFilterTracker<PointInT, StateT>::use_normal_;
      using ParticleFilterTracker<PointInT, StateT>::particle_num_;
      using ParticleFilterTracker<PointInT, StateT>::change_detector_filter_;
      using ParticleFilterTracker<PointInT, StateT>::transed_reference_vector_;
      //using ParticleFilterTracker<PointInT, StateT>::calcLikelihood;
      using ParticleFilterTracker<PointInT, StateT>::normalizeWeight;
      using ParticleFilterTracker<PointInT, StateT>::initParticles;
      using ParticleFilterTracker<PointInT, StateT>::normalizeParticleWeight;
      using ParticleFilterTracker<PointInT, StateT>::calcBoundingBox;

      typedef Tracker<PointInT, StateT> BaseClass;
      
      typedef typename Tracker<PointInT, StateT>::PointCloudIn PointCloudIn;
      typedef typename PointCloudIn::Ptr PointCloudInPtr;
      typedef typename PointCloudIn::ConstPtr PointCloudInConstPtr;

      typedef typename Tracker<PointInT, StateT>::PointCloudState PointCloudState;
      typedef typename PointCloudState::Ptr PointCloudStatePtr;
      typedef typename PointCloudState::ConstPtr PointCloudStateConstPtr;

      typedef PointCoherence<PointInT> Coherence;
      typedef boost::shared_ptr< Coherence > CoherencePtr;
      typedef boost::shared_ptr< const Coherence > CoherenceConstPtr;

      typedef PointCloudCoherence<PointInT> CloudCoherence;
      typedef boost::shared_ptr< CloudCoherence > CloudCoherencePtr;
      typedef boost::shared_ptr< const CloudCoherence > CloudCoherenceConstPtr;
      using ReversedParticleFilterTracker<PointInT, StateT>::transed_input_vector_;
    protected:

      unsigned int threads_;
      
    public:
      ReversedParticleFilterOMPTracker (unsigned int nr_threads = 0)
        : ReversedParticleFilterTracker<PointInT, StateT> ()
        , threads_ (nr_threads)
      {
        tracker_name_ = "ReversedParticleFilterOMPTracker";
      }

      inline void
      setNumberOfThreads (unsigned int nr_threads = 0) { threads_ = nr_threads; }

    protected:
      
      virtual void weight()
      {
        changed_ = true;
        if (!use_normal_)
        {
#ifdef _OPENMP
#pragma omp parallel for num_threads(threads_)
#endif
          for (size_t i = 0; i < particles_->points.size (); i++)
          {
            //std::cout << "processing " << i << " particle: " << particles_->points[i].weight << std::endl;
            // compute `inverse` of particle
            StateT inverse_particle;
            Eigen::Affine3f trans = particles_->points[i].toEigenMatrix();
            Eigen::Affine3f inverse_trans = trans.inverse();
            inverse_particle = StateT::toState(inverse_trans);
            this->computeTransformedPointCloudWithoutNormal (inverse_particle, *transed_input_vector_[i]);
            //computeTransformedPointCloudWithoutNormal (particles_->points[i], *transed_input_vector_[i]);
            IndicesPtr indices;
            coherence_->compute (transed_input_vector_[i], indices, particles_->points[i].weight);
            //std::cout << "processing " << i << " particle: " << particles_->points[i].weight << std::endl;
            //std::cout << inverse_particle << std::endl;
          }
        }
        else
        {
#ifdef _OPENMP
#pragma omp parallel for num_threads(threads_)
#endif
          for (size_t i = 0; i < particles_->points.size (); i++)
          {
            StateT inverse_particle;
            Eigen::Affine3f trans = particles_->points[i].toEigenMatrix();
            Eigen::Affine3f inverse_trans = trans.inverse();
            inverse_particle = StateT::toState(inverse_trans);
            IndicesPtr indices (new std::vector<int>);
            this->computeTransformedPointCloudWithNormal (inverse_particle, *indices, *transed_input_vector_[i]);
            coherence_->compute (transed_input_vector_[i], indices, particles_->points[i].weight);
          }
        }
  
        normalizeWeight ();
        for (size_t i = 0; i < particles_->points.size (); i++)
        {
          //std::cout << "normalized " << i << " particle: " << particles_->points[i].weight << std::endl;
        }
      }
    private:
    };
    
  }
  
}

using namespace pcl::tracking;
namespace jsk_pcl_ros
{
  class ParticleFilterTracking: public pcl_ros::PCLNodelet
  {
  public:
    typedef pcl::PointXYZRGBA PointT;
    typedef ParticleFilterTrackingConfig Config;
    typedef message_filters::sync_policies::ExactTime<
      sensor_msgs::PointCloud2,
      jsk_recognition_msgs::BoundingBox > SyncPolicy;
    typedef ParticleFilterTracker<PointT, ParticleXYZRPY>::PointCloudStatePtr
    PointCloudStatePtr;
  protected:
    pcl::PointCloud<PointT>::Ptr cloud_pass_;
    pcl::PointCloud<PointT>::Ptr cloud_pass_downsampled_;
    pcl::PointCloud<PointT>::Ptr target_cloud_;

    //boost::shared_ptr<ParticleFilterTracker<PointT, ParticleXYZRPY> > tracker_;
    boost::shared_ptr<KLDAdaptiveParticleFilterOMPTracker<PointT, ParticleXYZRPY> > tracker_;
    boost::shared_ptr<ReversedParticleFilterOMPTracker<PointT, ParticleXYZRPY> > reversed_tracker_;
    //boost::shared_ptr<ReversedParticleFilterTracker<PointT, ParticleXYZRPY> > reversed_tracker_;
    boost::mutex mtx_;
    bool new_cloud_;
    bool track_target_set_;
    bool align_box_;
    bool change_frame_;
    std::string frame_id_;
    std::string base_frame_id_;
    std::string track_target_name_;
    ros::Time stamp_;
    tf::Transform reference_transform_;

    ros::Subscriber sub_;
    ros::Subscriber sub_update_model_;

    message_filters::Subscriber<sensor_msgs::PointCloud2> sub_input_;
    message_filters::Subscriber<jsk_recognition_msgs::BoundingBox> sub_box_;
    boost::shared_ptr<message_filters::Synchronizer<SyncPolicy> >sync_;
    ros::Publisher particle_publisher_;
    ros::Publisher track_result_publisher_;
    ros::Publisher pose_stamped_publisher_;
    ros::ServiceServer renew_model_srv_;
    boost::shared_ptr <dynamic_reconfigure::Server<Config> > srv_;
    tf::TransformListener listener_;
    
    ////////////////////////////////////////////////////////
    // parameters
    ////////////////////////////////////////////////////////
    int max_particle_num_;
    double delta_;
    double epsilon_;
    int iteration_num_;
    double resample_likelihood_thr_;
    ParticleXYZRPY bin_size_;
    std::vector<double> default_step_covariance_;
    bool reversed_;
    bool not_use_reference_centroid_;
    bool not_publish_tf_;
    virtual void config_callback(Config &config, uint32_t level);
    virtual void publish_particles();
    virtual void publish_result();
    virtual std::string reference_frame_id();
    virtual void reset_tracking_target_model(
      const pcl::PointCloud<PointT>::ConstPtr &new_target_cloud);
    virtual tf::Transform change_pointcloud_frame(
      pcl::PointCloud<PointT>::Ptr cloud);

    virtual void cloud_cb(const sensor_msgs::PointCloud2 &pc);
    virtual bool renew_model_cb(
      jsk_pcl_ros::SetPointCloud2::Request &req,
      jsk_pcl_ros::SetPointCloud2::Response &response);
    virtual void renew_model_with_box_topic_cb(
      const sensor_msgs::PointCloud2::ConstPtr &pc_ptr,
      const jsk_recognition_msgs::BoundingBox::ConstPtr &bb_ptr);
    virtual void renew_model_topic_cb(const sensor_msgs::PointCloud2 &pc);

    ////////////////////////////////////////////////////////
    // Wrap particle filter methods
    ////////////////////////////////////////////////////////
    virtual void tracker_set_trans(const Eigen::Affine3f& trans);
    virtual void tracker_set_step_noise_covariance(
      const std::vector<double>& covariance);
    virtual void tracker_set_initial_noise_covariance(
      const std::vector<double>& covariance);
    virtual void tracker_set_initial_noise_mean(
      const std::vector<double>& mean);
    virtual void tracker_set_iteration_num(const int num);
    virtual void tracker_set_particle_num(const int num);
    virtual void tracker_set_resample_likelihood_thr(double thr);
    virtual void tracker_set_use_normal(bool use_normal);
    virtual void tracker_set_cloud_coherence(
      ApproxNearestPairPointCloudCoherence<PointT>::Ptr coherence);
    virtual void tracker_set_maximum_particle_num(int num);
    virtual void tracker_set_delta(double delta);
    virtual void tracker_set_epsilon(double epsilon);
    virtual void tracker_set_bin_size(const ParticleXYZRPY bin_size);
    virtual ParticleFilterTracker<PointT, ParticleXYZRPY>::PointCloudStatePtr
    tracker_get_particles();
    virtual ParticleXYZRPY tracker_get_result();
    virtual Eigen::Affine3f tracker_to_eigen_matrix(
      const ParticleXYZRPY& result);
    virtual pcl::PointCloud<PointT>::ConstPtr tracker_get_reference_cloud();
    virtual void tracker_set_reference_cloud(pcl::PointCloud<PointT>::Ptr ref);
    virtual void tracker_reset_tracking();
    virtual void tracker_set_input_cloud(pcl::PointCloud<PointT>::Ptr input);
    virtual void tracker_compute();
  private:
    virtual void onInit();

  };
}

#endif
