#include <gtest/gtest.h>

#include <cmath>
#include <cstdint>

#include <robot_depth_image_proc/point_cloud_xyz.h>
#include <robot_sensor_msgs/image_encodings.h>
#include <robot_sensor_msgs/point_cloud2_iterator.h>

namespace
{

robot_sensor_msgs::CameraInfo makeCameraInfo(uint32_t width, uint32_t height)
{
  const double cx = (static_cast<double>(width) - 1.0) * 0.5;
  const double cy = (static_cast<double>(height) - 1.0) * 0.5;

  robot_sensor_msgs::CameraInfo info;
  info.height = height;
  info.width = width;
  info.distortion_model = "plumb_bob";
  info.D = {0.0, 0.0, 0.0, 0.0, 0.0};
  info.K = {
      525.0, 0.0, cx,
      0.0, 525.0, cy,
      0.0, 0.0, 1.0};
  info.R = {
      1.0, 0.0, 0.0,
      0.0, 1.0, 0.0,
      0.0, 0.0, 1.0};
  info.P = {
      525.0, 0.0, cx, 0.0,
      0.0, 525.0, cy, 0.0,
      0.0, 0.0, 1.0, 0.0};
  return info;
}

robot_sensor_msgs::Image makeFlatDepthImage(
    uint32_t width,
    uint32_t height,
    uint16_t depth_mm)
{
  robot_sensor_msgs::Image image;
  image.height = height;
  image.width = width;
  image.encoding = robot_sensor_msgs::image_encodings::TYPE_16UC1;
  image.is_bigendian = false;
  image.step = width * sizeof(uint16_t);
  image.data.resize(static_cast<size_t>(height) * image.step);

  auto* depth = reinterpret_cast<uint16_t*>(image.data.data());
  for (size_t i = 0; i < width * height; ++i)
  {
    depth[i] = depth_mm;
  }
  return image;
}

}  // namespace

TEST(PointCloudXyz, ConvertsFlatDepthImage)
{
  const uint32_t width = 3;
  const uint32_t height = 3;
  const uint16_t depth_mm = 2000;

  const robot_sensor_msgs::Image depth = makeFlatDepthImage(width, height, depth_mm);
  const robot_sensor_msgs::CameraInfo info = makeCameraInfo(width, height);

  const robot_sensor_msgs::PointCloud2 cloud =
      depth_image_proc::convertDepthToPointCloud(depth, info);

  EXPECT_EQ(cloud.width, width);
  EXPECT_EQ(cloud.height, height);
  ASSERT_FALSE(cloud.data.empty());

  robot_sensor_msgs::PointCloud2ConstIterator<float> iter_x(cloud, "x");
  robot_sensor_msgs::PointCloud2ConstIterator<float> iter_y(cloud, "y");
  robot_sensor_msgs::PointCloud2ConstIterator<float> iter_z(cloud, "z");

  const float expected_z = static_cast<float>(depth_mm) * 0.001f;
  const size_t center_index = (height / 2) * width + (width / 2);
  size_t index = 0;

  for (size_t i = 0; i < width * height; ++i, ++iter_x, ++iter_y, ++iter_z, ++index)
  {
    EXPECT_NEAR(*iter_z, expected_z, 1e-3f);

    if (index == center_index)
    {
      EXPECT_NEAR(*iter_x, 0.0f, 1e-3f);
      EXPECT_NEAR(*iter_y, 0.0f, 1e-3f);
    }
  }
}

TEST(PointCloudXyz, RejectsUnsupportedEncoding)
{
  robot_sensor_msgs::Image depth = makeFlatDepthImage(4, 4, 1500);
  depth.encoding = robot_sensor_msgs::image_encodings::RGB8;

  const robot_sensor_msgs::CameraInfo info = makeCameraInfo(4, 4);
  const robot_sensor_msgs::PointCloud2 cloud =
      depth_image_proc::convertDepthToPointCloud(depth, info);

  EXPECT_EQ(cloud.width, 0u);
  EXPECT_EQ(cloud.height, 0u);
}

int main(int argc, char** argv)
{
  testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}