/*
 * Software License Agreement (BSD License)
 *
 *  Point Cloud Library (PCL) - www.pointclouds.org
 *  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 Willow Garage, Inc. 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.
 *
 */

#pragma once

#include <pcl/common/distances.h>
#include <pcl/common/io.h> // for getFieldIndex
#include <pcl/common/point_tests.h> // for pcl::isFinite
#include <pcl/search/organized.h>
#include <pcl/search/kdtree.h>


namespace pcl
{

template <>
float squaredEuclideanDistance (const pcl::segmentation::grabcut::Color &c1,
                                const pcl::segmentation::grabcut::Color &c2)
{
  return ((c1.r-c2.r)*(c1.r-c2.r)+(c1.g-c2.g)*(c1.g-c2.g)+(c1.b-c2.b)*(c1.b-c2.b));
}

namespace segmentation
{

template <typename PointT>
grabcut::Color::Color (const PointT& p)
{
  r = static_cast<float> (p.r) / 255.0;
  g = static_cast<float> (p.g) / 255.0;
  b = static_cast<float> (p.b) / 255.0;
}

template <typename PointT>
grabcut::Color::operator PointT () const
{
  PointT p;
  p.r = static_cast<std::uint32_t> (r * 255);
  p.g = static_cast<std::uint32_t> (g * 255);
  p.b = static_cast<std::uint32_t> (b * 255);
  return (p);
}

} // namespace segmentation

template <typename PointT> void
GrabCut<PointT>::setInputCloud (const PointCloudConstPtr &cloud)
{
  input_ = cloud;
}

template <typename PointT> bool
GrabCut<PointT>::initCompute ()
{
  using namespace pcl::segmentation::grabcut;
  if (!pcl::PCLBase<PointT>::initCompute ())
  {
    PCL_ERROR ("[pcl::GrabCut::initCompute ()] Init failed!\n");
    return (false);
  }

  std::vector<pcl::PCLPointField> in_fields_;
  if ((pcl::getFieldIndex<PointT> ("rgb", in_fields_) == -1) &&
      (pcl::getFieldIndex<PointT> ("rgba", in_fields_) == -1))
  {
    PCL_ERROR ("[pcl::GrabCut::initCompute ()] No RGB data available, aborting!\n");
    return (false);
  }

  // Initialize the working image
  image_.reset (new Image (input_->width, input_->height));
  for (std::size_t i = 0; i < input_->size (); ++i)
  {
    (*image_) [i] = Color ((*input_)[i]);
  }
  width_ = image_->width;
  height_ = image_->height;

  // Initialize the spatial locator
  if (!tree_ && !input_->isOrganized ())
  {
    tree_.reset (new pcl::search::KdTree<PointT> (true));
    tree_->setInputCloud (input_);
  }

  const std::size_t indices_size = indices_->size ();
  trimap_ = std::vector<segmentation::grabcut::TrimapValue> (indices_size, TrimapUnknown);
  hard_segmentation_ = std::vector<segmentation::grabcut::SegmentationValue> (indices_size,
                                                                              SegmentationBackground);
  GMM_component_.resize (indices_size);
  n_links_.resize (indices_size);

  //  soft_segmentation_ = 0;		// Not yet implemented
  foreground_GMM_.resize (K_);
  background_GMM_.resize (K_);

  //set some constants
  computeL ();

  if (image_->isOrganized ())
  {
    computeBetaOrganized ();
    computeNLinksOrganized ();
  }
  else
  {
    computeBetaNonOrganized ();
    computeNLinksNonOrganized ();
  }

  initialized_ = false;
  return (true);
}

template <typename PointT> void
GrabCut<PointT>::addEdge (vertex_descriptor v1, vertex_descriptor v2, float capacity, float rev_capacity)
{
  graph_.addEdge (v1, v2, capacity, rev_capacity);
}

template <typename PointT> void
GrabCut<PointT>::setTerminalWeights (vertex_descriptor v, float source_capacity, float sink_capacity)
{
  graph_.addSourceEdge (v, source_capacity);
  graph_.addTargetEdge (v, sink_capacity);
}

template <typename PointT> void
GrabCut<PointT>::setBackgroundPointsIndices (const PointIndicesConstPtr &indices)
{
  using namespace pcl::segmentation::grabcut;
  if (!initCompute ())
    return;

  std::fill (trimap_.begin (), trimap_.end (), TrimapBackground);
  std::fill (hard_segmentation_.begin (), hard_segmentation_.end (), SegmentationBackground);
  for (const auto &index : indices->indices)
  {
    trimap_[index] = TrimapUnknown;
    hard_segmentation_[index] = SegmentationForeground;
  }

  if (!initialized_)
  {
    fitGMMs ();
    initialized_ = true;
  }
}

template <typename PointT> void
GrabCut<PointT>::fitGMMs ()
{
  // Step 3: Build GMMs using Orchard-Bouman clustering algorithm
  buildGMMs (*image_, *indices_, hard_segmentation_, GMM_component_, background_GMM_, foreground_GMM_);

  // Initialize the graph for graphcut (do this here so that the T-Link debugging image will be initialized)
  initGraph ();
}

template <typename PointT> int
GrabCut<PointT>::refineOnce ()
{
  // Steps 4 and 5: Learn new GMMs from current segmentation
  learnGMMs (*image_, *indices_, hard_segmentation_, GMM_component_, background_GMM_, foreground_GMM_);

  // Step 6: Run GraphCut and update segmentation
  initGraph ();

  float flow = graph_.solve ();

  int changed = updateHardSegmentation ();
  PCL_INFO ("%d pixels changed segmentation (max flow = %f)\n", changed, flow);

  return (changed);
}

template <typename PointT> void
GrabCut<PointT>::refine ()
{
  std::size_t changed = indices_->size ();

  while (changed)
    changed = refineOnce ();
}

template <typename PointT> int
GrabCut<PointT>::updateHardSegmentation ()
{
  using namespace pcl::segmentation::grabcut;

  int changed = 0;

  const int number_of_indices = static_cast<int> (indices_->size ());
  for (int i_point = 0; i_point < number_of_indices; ++i_point)
  {
    SegmentationValue old_value = hard_segmentation_ [i_point];

    if (trimap_ [i_point] == TrimapBackground)
      hard_segmentation_ [i_point] = SegmentationBackground;
    else
      if (trimap_ [i_point] == TrimapForeground)
        hard_segmentation_ [i_point] = SegmentationForeground;
      else	// TrimapUnknown
      {
        if (isSource (graph_nodes_[i_point]))
          hard_segmentation_ [i_point] = SegmentationForeground;
        else
          hard_segmentation_ [i_point] = SegmentationBackground;
      }

    if (old_value != hard_segmentation_ [i_point])
      ++changed;
  }
  return (changed);
}

template <typename PointT> void
GrabCut<PointT>::setTrimap (const PointIndicesConstPtr &indices, segmentation::grabcut::TrimapValue t)
{
  using namespace pcl::segmentation::grabcut;
  for (const auto &index : indices->indices)
    trimap_[index] = t;

  // Immediately set the hard segmentation as well so that the display will update.
  if (t == TrimapForeground)
    for (const auto &index : indices->indices)
      hard_segmentation_[index] = SegmentationForeground;
  else
    if (t == TrimapBackground)
      for (const auto &index : indices->indices)
        hard_segmentation_[index] = SegmentationBackground;
}

template <typename PointT> void
GrabCut<PointT>::initGraph ()
{
  using namespace pcl::segmentation::grabcut;
  const int number_of_indices = static_cast<int> (indices_->size ());
  // Set up the graph (it can only be used once, so we have to recreate it each time the graph is updated)
  graph_.clear ();
  graph_nodes_.clear ();
  graph_nodes_.resize (indices_->size ());
  int start = graph_.addNodes (indices_->size ());
  for (std::size_t idx = 0; idx < indices_->size (); ++idx)
  {
    graph_nodes_[idx] = start;
    ++start;
  }

  // Set T-Link weights
  for (int i_point = 0; i_point < number_of_indices; ++i_point)
  {
    int point_index = (*indices_) [i_point];
    float back, fore;

    switch (trimap_[point_index])
    {
      case TrimapUnknown :
      {
        fore = static_cast<float> (-std::log (background_GMM_.probabilityDensity ((*image_)[point_index])));
        back = static_cast<float> (-std::log (foreground_GMM_.probabilityDensity ((*image_)[point_index])));
        break;
      }
      case TrimapBackground :
      {
        fore = 0;
        back = L_;
        break;
      }
      default :
      {
        fore = L_;
        back = 0;
      }
    }

    setTerminalWeights (graph_nodes_[i_point], fore, back);
  }

  // Set N-Link weights from precomputed values
  for (int i_point = 0; i_point < number_of_indices; ++i_point)
  {
    const NLinks &n_link = n_links_[i_point];
    if (n_link.nb_links > 0)
    {
      const auto point_index = (*indices_) [i_point];
      std::vector<float>::const_iterator weights_it  = n_link.weights.begin ();
      for (auto indices_it = n_link.indices.cbegin (); indices_it != n_link.indices.cend (); ++indices_it, ++weights_it)
      {
        if ((*indices_it != point_index) && (*indices_it != UNAVAILABLE))
        {
          addEdge (graph_nodes_[i_point], graph_nodes_[*indices_it], *weights_it, *weights_it);
        }
      }
    }
  }
}

template <typename PointT> void
GrabCut<PointT>::computeNLinksNonOrganized ()
{
  const int number_of_indices = static_cast<int> (indices_->size ());
  for (int i_point = 0; i_point < number_of_indices; ++i_point)
  {
    NLinks &n_link = n_links_[i_point];
    if (n_link.nb_links > 0)
    {
      const auto point_index = (*indices_) [i_point];
      auto dists_it = n_link.dists.cbegin ();
      auto weights_it = n_link.weights.begin ();
      for (auto indices_it = n_link.indices.cbegin (); indices_it != n_link.indices.cend (); ++indices_it, ++dists_it, ++weights_it)
      {
        if (*indices_it != point_index)
        {
          // We saved the color distance previously at the computeBeta stage for optimization purpose
          float color_distance = *weights_it;
          // Set the real weight
          *weights_it = static_cast<float> (lambda_ * std::exp (-beta_ * color_distance) / sqrt (*dists_it));
        }
      }
    }
  }
}

template <typename PointT> void
GrabCut<PointT>::computeNLinksOrganized ()
{
	for( unsigned int y = 0; y < image_->height; ++y )
	{
    for( unsigned int x = 0; x < image_->width; ++x )
    {
      // We saved the color and euclidean distance previously at the computeBeta stage for
      // optimization purpose but here we compute the real weight
      std::size_t point_index = y * input_->width + x;
      NLinks &links = n_links_[point_index];

      if( x > 0 && y < image_->height-1 )
        links.weights[0] = lambda_ * std::exp (-beta_ * links.weights[0]) / links.dists[0];

      if( y < image_->height-1 )
        links.weights[1] = lambda_ * std::exp (-beta_ * links.weights[1]) / links.dists[1];

      if( x < image_->width-1 && y < image_->height-1 )
        links.weights[2] = lambda_ * std::exp (-beta_ * links.weights[2]) / links.dists[2];

      if( x < image_->width-1 )
        links.weights[3] = lambda_ * std::exp (-beta_ * links.weights[3]) / links.dists[3];
    }
	}
}

template <typename PointT> void
GrabCut<PointT>::computeBetaNonOrganized ()
{
  float result = 0;
  std::size_t edges = 0;

  const int number_of_indices = static_cast<int> (indices_->size ());

  for (int i_point = 0; i_point < number_of_indices; i_point++)
  {
    const auto point_index = (*indices_)[i_point];
    const PointT& point = input_->points [point_index];
    if (pcl::isFinite (point))
    {
      NLinks &links = n_links_[i_point];
      int found = tree_->nearestKSearch (point, nb_neighbours_, links.indices, links.dists);
      if (found > 1)
      {
        links.nb_links = found - 1;
        links.weights.reserve (links.nb_links);
        for (const auto& nn_index : links.indices)
        {
          if (nn_index != point_index)
          {
            float color_distance = squaredEuclideanDistance ((*image_)[point_index], (*image_)[nn_index]);
            links.weights.push_back (color_distance);
            result+= color_distance;
            ++edges;
          }
          else
            links.weights.push_back (0.f);
        }
      }
    }
  }

  beta_ = 1e5 / (2*result / edges);
}

template <typename PointT> void
GrabCut<PointT>::computeBetaOrganized ()
{
  float result = 0;
  std::size_t edges = 0;

  for (unsigned int y = 0; y < input_->height; ++y)
  {
    for (unsigned int x = 0; x < input_->width; ++x)
    {
      std::size_t point_index = y * input_->width + x;
      NLinks &links = n_links_[point_index];
      links.nb_links = 4;
      links.weights.resize (links.nb_links, 0);
      links.dists.resize (links.nb_links, 0);
      links.indices.resize (links.nb_links, -1);

      if (x > 0 && y < input_->height-1)
      {
        std::size_t upleft = (y+1)  * input_->width + x - 1;
        links.indices[0] = upleft;
        links.dists[0] = std::sqrt (2.f);
        float color_dist =  squaredEuclideanDistance ((*image_)[point_index],
                                                      (*image_)[upleft]);
        links.weights[0] = color_dist;
        result+= color_dist;
        edges++;
      }

      if (y < input_->height-1)
      {
        std::size_t up = (y+1) * input_->width + x;
        links.indices[1] = up;
        links.dists[1] = 1;
        float color_dist =  squaredEuclideanDistance ((*image_)[point_index],
                                                      (*image_)[up]);
        links.weights[1] = color_dist;
        result+= color_dist;
        edges++;
      }

      if (x < input_->width-1 && y < input_->height-1)
      {
        std::size_t upright = (y+1) * input_->width + x + 1;
        links.indices[2] = upright;
        links.dists[2] = std::sqrt (2.f);
        float color_dist =  squaredEuclideanDistance ((*image_)[point_index],
                                                      image_->points [upright]);
        links.weights[2] = color_dist;
        result+= color_dist;
        edges++;
      }

      if (x < input_->width-1)
      {
        std::size_t right = y * input_->width + x + 1;
        links.indices[3] = right;
        links.dists[3] = 1;
        float color_dist =  squaredEuclideanDistance ((*image_)[point_index],
                                                      (*image_)[right]);
        links.weights[3] = color_dist;
        result+= color_dist;
        edges++;
      }
    }
  }

  beta_ = 1e5 / (2*result / edges);
}

template <typename PointT> void
GrabCut<PointT>::computeL ()
{
  L_ = 8*lambda_ + 1;
}

template <typename PointT> void
GrabCut<PointT>::extract (std::vector<pcl::PointIndices>& clusters)
{
  using namespace pcl::segmentation::grabcut;
  clusters.clear ();
  clusters.resize (2);
  clusters[0].indices.reserve (indices_->size ());
  clusters[1].indices.reserve (indices_->size ());
  refine ();
  assert (hard_segmentation_.size () == indices_->size ());
  const int indices_size = static_cast<int> (indices_->size ());
  for (int i = 0; i < indices_size; ++i)
    if (hard_segmentation_[i] == SegmentationForeground)
      clusters[1].indices.push_back (i);
    else
      clusters[0].indices.push_back (i);
}

} // namespace pcl