/*
 * Copyright (c) 20013, Willow Garage, 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 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.
 *
 *
 *  Created on: Jan 22, 2013
 *      Author: jkammerl
 */

#include <sensor_msgs/CameraInfo.h>
#include <sensor_msgs/image_encodings.h>

#include <string.h>
#include <sstream>

#include "depth_cloud_mld.h"

namespace enc = sensor_msgs::image_encodings;

#define POINT_STEP (sizeof(float) * 4)

namespace rviz
{
// Encapsulate differences between processing float and uint16_t depths
template <typename T>
struct DepthTraits
{
};

template <>
struct DepthTraits<uint16_t>
{
  static inline bool valid(float depth)
  {
    return depth != 0.0;
  }
  static inline float toMeters(uint16_t depth)
  {
    return depth * 0.001f;
  } // originally mm
};

template <>
struct DepthTraits<float>
{
  static inline bool valid(float depth)
  {
    return std::isfinite(depth);
  }
  static inline float toMeters(float depth)
  {
    return depth;
  }
};


struct RGBA
{
  uint8_t red;
  uint8_t green;
  uint8_t blue;
  uint8_t alpha;
};


void MultiLayerDepth::initializeConversion(const sensor_msgs::ImageConstPtr& depth_msg,
                                           sensor_msgs::CameraInfoConstPtr& camera_info_msg)
{
  if (!depth_msg || !camera_info_msg)
  {
    std::string error_msg("Waiting for CameraInfo message..");
    throw(MultiLayerDepthException(error_msg));
  }

  // do some sanity checks
  unsigned binning_x = camera_info_msg->binning_x > 1 ? camera_info_msg->binning_x : 1;
  unsigned binning_y = camera_info_msg->binning_y > 1 ? camera_info_msg->binning_y : 1;

  unsigned roi_width =
      camera_info_msg->roi.width > 0 ? camera_info_msg->roi.width : camera_info_msg->width;
  unsigned roi_height =
      camera_info_msg->roi.height > 0 ? camera_info_msg->roi.height : camera_info_msg->height;

  unsigned expected_width = roi_width / binning_x;
  unsigned expected_height = roi_height / binning_y;

  if (expected_width != depth_msg->width || expected_height != depth_msg->height)
  {
    std::ostringstream s;
    s << "Depth image size and camera info don't match: ";
    s << depth_msg->width << " x " << depth_msg->height;
    s << " vs " << expected_width << " x " << expected_height;
    s << "(binning: " << binning_x << " x " << binning_y;
    s << ", ROI size: " << roi_width << " x " << roi_height << ")";
    throw(MultiLayerDepthException(s.str()));
  }

  int width = depth_msg->width;
  int height = depth_msg->height;

  std::size_t size = width * height;

  if (size != shadow_depth_.size())
  {
    // Allocate memory for shadow processing
    shadow_depth_.resize(size, 0.0f);
    shadow_timestamp_.resize(size, 0.0);
    shadow_buffer_.resize(size * POINT_STEP, 0);

    // Procompute 3D projection matrix
    //
    // The following computation of center_x,y and fx,fy duplicates
    // code in the image_geometry package, but this avoids dependency
    // on OpenCV, which simplifies releasing rviz.

    // Combine unit conversion (if necessary) with scaling by focal length for computing (X,Y)
    double scale_x = camera_info_msg->binning_x > 1 ? (1.0 / camera_info_msg->binning_x) : 1.0;
    double scale_y = camera_info_msg->binning_y > 1 ? (1.0 / camera_info_msg->binning_y) : 1.0;

    // Use correct principal point from calibration
    float center_x = (camera_info_msg->P[2] - camera_info_msg->roi.x_offset) * scale_x;
    float center_y = (camera_info_msg->P[6] - camera_info_msg->roi.y_offset) * scale_y;

    double fx = camera_info_msg->P[0] * scale_x;
    double fy = camera_info_msg->P[5] * scale_y;

    float constant_x = 1.0f / fx;
    float constant_y = 1.0f / fy;

    projection_map_x_.resize(width);
    projection_map_y_.resize(height);
    std::vector<float>::iterator projX = projection_map_x_.begin();
    std::vector<float>::iterator projY = projection_map_y_.begin();

    // precompute 3D projection matrix
    for (int v = 0; v < height; ++v, ++projY)
      *projY = (v - center_y) * constant_y;

    for (int u = 0; u < width; ++u, ++projX)
      *projX = (u - center_x) * constant_x;

    // reset shadow vectors
    reset();
  }
}


template <typename T>
sensor_msgs::PointCloud2Ptr
MultiLayerDepth::generatePointCloudSL(const sensor_msgs::ImageConstPtr& depth_msg,
                                      std::vector<uint32_t>& rgba_color_raw)
{
  int width = depth_msg->width;
  int height = depth_msg->height;

  sensor_msgs::PointCloud2Ptr cloud_msg = initPointCloud();
  cloud_msg->data.resize(height * width * cloud_msg->point_step);

  uint32_t* color_img_ptr = nullptr;

  if (!rgba_color_raw.empty())
    color_img_ptr = &rgba_color_raw[0];

  ////////////////////////////////////////////////
  // depth map to point cloud conversion
  ////////////////////////////////////////////////

  float* cloud_data_ptr = reinterpret_cast<float*>(&cloud_msg->data[0]);

  std::size_t point_count = 0;
  std::size_t point_idx = 0;

  const T* depth_img_ptr = (T*)&depth_msg->data[0];

  std::vector<float>::iterator proj_x;
  std::vector<float>::const_iterator proj_x_end = projection_map_x_.end();

  std::vector<float>::iterator proj_y;
  std::vector<float>::const_iterator proj_y_end = projection_map_y_.end();

  // iterate over projection matrix
  for (proj_y = projection_map_y_.begin(); proj_y != proj_y_end; ++proj_y)
  {
    for (proj_x = projection_map_x_.begin(); proj_x != proj_x_end; ++proj_x, ++point_idx, ++depth_img_ptr)
    {
      T depth_raw = *depth_img_ptr;
      if (DepthTraits<T>::valid(depth_raw))
      {
        float depth = DepthTraits<T>::toMeters(depth_raw);

        // define point color
        uint32_t color;
        if (color_img_ptr)
        {
          color = *color_img_ptr;
        }
        else
        {
          color = ((uint32_t)255 << 16 | (uint32_t)255 << 8 | (uint32_t)255);
        }

        // fill in X,Y,Z and color
        *cloud_data_ptr = (*proj_x) * depth;
        ++cloud_data_ptr;
        *cloud_data_ptr = (*proj_y) * depth;
        ++cloud_data_ptr;
        *cloud_data_ptr = depth;
        ++cloud_data_ptr;
        *cloud_data_ptr = *reinterpret_cast<float*>(&color);
        ++cloud_data_ptr;

        ++point_count;
      }

      // increase color iterator pointer
      if (color_img_ptr)
        ++color_img_ptr;
    }
  }

  finalizingPointCloud(cloud_msg, point_count);

  return cloud_msg;
}


template <typename T>
sensor_msgs::PointCloud2Ptr
MultiLayerDepth::generatePointCloudML(const sensor_msgs::ImageConstPtr& depth_msg,
                                      std::vector<uint32_t>& rgba_color_raw)
{
  int width = depth_msg->width;
  int height = depth_msg->height;

  sensor_msgs::PointCloud2Ptr cloud_msg = initPointCloud();
  cloud_msg->data.resize(height * width * cloud_msg->point_step * 2);

  uint32_t* color_img_ptr = nullptr;

  if (!rgba_color_raw.empty())
    color_img_ptr = &rgba_color_raw[0];

  ////////////////////////////////////////////////
  // depth map to point cloud conversion
  ////////////////////////////////////////////////

  float* cloud_data_ptr = reinterpret_cast<float*>(&cloud_msg->data[0]);
  uint8_t* cloud_shadow_buffer_ptr = &shadow_buffer_[0];

  const std::size_t point_step = cloud_msg->point_step;

  std::size_t point_count = 0;
  std::size_t point_idx = 0;

  double time_now = ros::Time::now().toSec();
  double time_expire = time_now - shadow_time_out_;

  const T* depth_img_ptr = (T*)&depth_msg->data[0];

  std::vector<float>::iterator proj_x;
  std::vector<float>::const_iterator proj_x_end = projection_map_x_.end();

  std::vector<float>::iterator proj_y;
  std::vector<float>::const_iterator proj_y_end = projection_map_y_.end();

  // iterate over projection matrix
  for (proj_y = projection_map_y_.begin(); proj_y != proj_y_end; ++proj_y)
  {
    for (proj_x = projection_map_x_.begin(); proj_x != proj_x_end;
         ++proj_x, ++point_idx, ++depth_img_ptr, cloud_shadow_buffer_ptr += point_step)
    {
      // lookup shadow depth
      float shadow_depth = shadow_depth_[point_idx];

      // check for time-outs
      if ((shadow_depth != 0.0f) && (shadow_timestamp_[point_idx] < time_expire))
      {
        // clear shadow pixel
        shadow_depth = shadow_depth_[point_idx] = 0.0f;
      }

      T depth_raw = *depth_img_ptr;
      if (DepthTraits<T>::valid(depth_raw))
      {
        float depth = DepthTraits<T>::toMeters(depth_raw);

        // pointer to current point data
        float* cloud_data_pixel_ptr = cloud_data_ptr;

        // define point color
        uint32_t color;
        if (color_img_ptr)
        {
          color = *color_img_ptr;
        }
        else
        {
          color = ((uint32_t)255 << 16 | (uint32_t)255 << 8 | (uint32_t)255);
        }

        // fill in X,Y,Z and color
        *cloud_data_ptr = (*proj_x) * depth;
        ++cloud_data_ptr;
        *cloud_data_ptr = (*proj_y) * depth;
        ++cloud_data_ptr;
        *cloud_data_ptr = depth;
        ++cloud_data_ptr;
        *cloud_data_ptr = *reinterpret_cast<float*>(&color);
        ++cloud_data_ptr;

        ++point_count;

        // if shadow point exists -> display it
        if (depth < shadow_depth - shadow_distance_)
        {
          // copy point data from shadow buffer to point cloud
          memcpy(cloud_data_ptr, cloud_shadow_buffer_ptr, point_step);
          cloud_data_ptr += 4;
          ++point_count;
        }
        else
        {
          // save a copy of current point to shadow buffer
          memcpy(cloud_shadow_buffer_ptr, cloud_data_pixel_ptr, point_step);

          // reduce color intensity in shadow buffer
          RGBA* color = reinterpret_cast<RGBA*>(cloud_shadow_buffer_ptr + sizeof(float) * 3);
          color->red /= 2;
          color->green /= 2;
          color->blue /= 2;

          // update shadow depth & time out
          shadow_depth_[point_idx] = depth;
          shadow_timestamp_[point_idx] = time_now;
        }
      }
      else
      {
        // current depth pixel is invalid -> check shadow buffer
        if (shadow_depth != 0)
        {
          // copy shadow point to point cloud
          memcpy(cloud_data_ptr, cloud_shadow_buffer_ptr, point_step);
          cloud_data_ptr += 4;
          ++point_count;
        }
      }

      // increase color iterator pointer
      if (color_img_ptr)
        ++color_img_ptr;
    }
  }

  finalizingPointCloud(cloud_msg, point_count);

  return cloud_msg;
}


template <typename T>
void MultiLayerDepth::convertColor(const sensor_msgs::ImageConstPtr& color_msg,
                                   std::vector<uint32_t>& rgba_color_raw)
{
  size_t i;
  size_t num_pixel = color_msg->width * color_msg->height;

  // query image properties
  int num_channels = enc::numChannels(color_msg->encoding);

  bool rgb_encoding = false;
  if (color_msg->encoding.find("rgb") != std::string::npos)
    rgb_encoding = true;

  bool has_alpha = enc::hasAlpha(color_msg->encoding);

  // prepare output vector
  rgba_color_raw.clear();
  rgba_color_raw.reserve(num_pixel);

  uint8_t* img_ptr = (uint8_t*)&color_msg->data[sizeof(T) - 1]; // pointer to most significant byte

  // color conversion
  switch (num_channels)
  {
  case 1:
    // grayscale image
    for (i = 0; i < num_pixel; ++i)
    {
      uint8_t gray_value = *img_ptr;
      img_ptr += sizeof(T);

      rgba_color_raw.push_back((uint32_t)gray_value << 16 | (uint32_t)gray_value << 8 |
                               (uint32_t)gray_value);
    }
    break;
  case 3:
  case 4:
    // rgb/bgr encoding
    for (i = 0; i < num_pixel; ++i)
    {
      uint8_t color1 = *((uint8_t*)img_ptr);
      img_ptr += sizeof(T);
      uint8_t color2 = *((uint8_t*)img_ptr);
      img_ptr += sizeof(T);
      uint8_t color3 = *((uint8_t*)img_ptr);
      img_ptr += sizeof(T);

      if (has_alpha)
        img_ptr += sizeof(T); // skip alpha values

      if (rgb_encoding)
      {
        // rgb encoding
        rgba_color_raw.push_back((uint32_t)color1 << 16 | (uint32_t)color2 << 8 | (uint32_t)color3 << 0);
      }
      else
      {
        // bgr encoding
        rgba_color_raw.push_back((uint32_t)color3 << 16 | (uint32_t)color2 << 8 | (uint32_t)color1 << 0);
      }
    }
    break;
  default:
    break;
  }
}

sensor_msgs::PointCloud2Ptr
MultiLayerDepth::generatePointCloudFromDepth(const sensor_msgs::ImageConstPtr& depth_msg,
                                             const sensor_msgs::ImageConstPtr& color_msg,
                                             sensor_msgs::CameraInfoConstPtr camera_info_msg)
{
  // Add data to multi depth image
  sensor_msgs::PointCloud2Ptr point_cloud_out;

  // Bit depth of image encoding
  int bitDepth = enc::bitDepth(depth_msg->encoding);
  int numChannels = enc::numChannels(depth_msg->encoding);

  if (!camera_info_msg)
  {
    throw MultiLayerDepthException("Camera info missing!");
  }

  // precompute projection matrix and initialize shadow buffer
  initializeConversion(depth_msg, camera_info_msg);

  std::vector<uint32_t> rgba_color_raw_;

  if (color_msg)
  {
    if (depth_msg->width != color_msg->width || depth_msg->height != color_msg->height)
    {
      std::stringstream error_msg;
      error_msg << "Depth image resolution (" << (int)depth_msg->width << "x" << (int)depth_msg->height
                << ") "
                   "does not match color image resolution ("
                << (int)color_msg->width << "x" << (int)color_msg->height << ")";
      throw(MultiLayerDepthException(error_msg.str()));
    }

    // convert color coding to 8-bit rgb data
    switch (enc::bitDepth(color_msg->encoding))
    {
    case 8:
      convertColor<uint8_t>(color_msg, rgba_color_raw_);
      break;
    case 16:
      convertColor<uint16_t>(color_msg, rgba_color_raw_);
      break;
    default:
      std::string error_msg("Color image has invalid bit depth");
      throw(MultiLayerDepthException(error_msg));
      break;
    }
  }

  if (!occlusion_compensation_)
  {
    // generate single layer depth cloud

    if ((bitDepth == 32) && (numChannels == 1))
    {
      // floating point encoded depth map
      point_cloud_out = generatePointCloudSL<float>(depth_msg, rgba_color_raw_);
    }
    else if ((bitDepth == 16) && (numChannels == 1))
    {
      // 32bit integer encoded depth map
      point_cloud_out = generatePointCloudSL<uint16_t>(depth_msg, rgba_color_raw_);
    }
  }
  else
  {
    // generate two layered depth cloud (depth+shadow)

    if ((bitDepth == 32) && (numChannels == 1))
    {
      // floating point encoded depth map
      point_cloud_out = generatePointCloudML<float>(depth_msg, rgba_color_raw_);
    }
    else if ((bitDepth == 16) && (numChannels == 1))
    {
      // 32bit integer encoded depth map
      point_cloud_out = generatePointCloudML<uint16_t>(depth_msg, rgba_color_raw_);
    }
  }

  if (!point_cloud_out)
  {
    std::string error_msg("Depth image has invalid format (only 16 bit and float are supported)!");
    throw(MultiLayerDepthException(error_msg));
  }

  return point_cloud_out;
}

sensor_msgs::PointCloud2Ptr MultiLayerDepth::initPointCloud()
{
  sensor_msgs::PointCloud2Ptr point_cloud_out =
      sensor_msgs::PointCloud2Ptr(new sensor_msgs::PointCloud2());

  point_cloud_out->fields.resize(4);
  std::size_t point_offset = 0;

  point_cloud_out->fields[0].name = "x";
  point_cloud_out->fields[0].datatype = sensor_msgs::PointField::FLOAT32;
  point_cloud_out->fields[0].count = 1;
  point_cloud_out->fields[0].offset = point_offset;
  point_offset += sizeof(float);

  point_cloud_out->fields[1].name = "y";
  point_cloud_out->fields[1].datatype = sensor_msgs::PointField::FLOAT32;
  point_cloud_out->fields[1].count = 1;
  point_cloud_out->fields[1].offset = point_offset;
  point_offset += sizeof(float);

  point_cloud_out->fields[2].name = "z";
  point_cloud_out->fields[2].datatype = sensor_msgs::PointField::FLOAT32;
  point_cloud_out->fields[2].count = 1;
  point_cloud_out->fields[2].offset = point_offset;
  point_offset += sizeof(float);

  point_cloud_out->fields[3].name = "rgb";
  point_cloud_out->fields[3].datatype = sensor_msgs::PointField::FLOAT32;
  point_cloud_out->fields[3].count = 1;
  point_cloud_out->fields[3].offset = point_offset;
  point_offset += sizeof(float);

  point_cloud_out->point_step = point_offset;

  point_cloud_out->is_bigendian = false;
  point_cloud_out->is_dense = false;

  return point_cloud_out;
}

void MultiLayerDepth::finalizingPointCloud(sensor_msgs::PointCloud2Ptr& point_cloud, std::size_t size)
{
  point_cloud->width = size;
  point_cloud->height = 1;
  point_cloud->data.resize(point_cloud->height * point_cloud->width * point_cloud->point_step);
  point_cloud->row_step = point_cloud->point_step * point_cloud->width;
}

} // namespace rviz
