#include <gtest/gtest.h>

#include <cmath>
#include <cstdint>

#include <robot_depth_image_proc/depth_frame_filter.h>
#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);
}

TEST(PointCloudXyzFiltered, DecimatesAndClipsRange)
{
  const uint32_t width = 8;
  const uint32_t height = 8;
  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);

  depth_image_proc::DepthFilterConfig config;
  config.decimation = 2;
  config.range_min = 0.3;
  config.range_max = 4.0;
  config.speckle_min_neighbors = 0;

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

  EXPECT_EQ(cloud.height, 1u);
  EXPECT_EQ(cloud.width, (width / 2) * (height / 2));
  EXPECT_TRUE(cloud.is_dense);

  robot_sensor_msgs::PointCloud2ConstIterator<float> iter_z(cloud, "z");
  for (size_t i = 0; i < cloud.width; ++i, ++iter_z)
  {
    EXPECT_NEAR(*iter_z, 2.0f, 1e-3f);
  }

  // Every point out of range -> empty but well-formed cloud.
  config.range_min = 3.0;
  config.range_max = 4.0;
  const robot_sensor_msgs::PointCloud2 empty_cloud =
      depth_image_proc::convertDepthToPointCloudFiltered(depth, info, config);
  EXPECT_EQ(empty_cloud.width, 0u);
  EXPECT_FALSE(empty_cloud.fields.empty());
}

TEST(PointCloudXyzFiltered, RemovesSpecklePoint)
{
  const uint32_t width = 9;
  const uint32_t height = 9;
  const uint16_t depth_mm = 2000;

  robot_sensor_msgs::Image depth = makeFlatDepthImage(width, height, depth_mm);
  // One isolated pixel jumps 0.5 m out of the surface: a flying pixel.
  auto* data = reinterpret_cast<uint16_t*>(depth.data.data());
  data[4 * width + 4] = 2500;

  const robot_sensor_msgs::CameraInfo info = makeCameraInfo(width, height);

  depth_image_proc::DepthFilterConfig config;
  config.decimation = 1;
  config.range_min = 0.3;
  config.range_max = 4.0;
  config.speckle_max_delta = 0.08;
  config.speckle_min_neighbors = 3;

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

  EXPECT_EQ(cloud.width, width * height - 1);

  robot_sensor_msgs::PointCloud2ConstIterator<float> iter_z(cloud, "z");
  for (size_t i = 0; i < cloud.width; ++i, ++iter_z)
  {
    EXPECT_NEAR(*iter_z, 2.0f, 1e-3f);
  }
}

TEST(PointCloudXyzFiltered, RejectsInvalidConfig)
{
  const robot_sensor_msgs::Image depth = makeFlatDepthImage(4, 4, 1500);
  const robot_sensor_msgs::CameraInfo info = makeCameraInfo(4, 4);

  depth_image_proc::DepthFilterConfig config;
  config.decimation = 0;

  const robot_sensor_msgs::PointCloud2 cloud =
      depth_image_proc::convertDepthToPointCloudFiltered(depth, info, config);
  EXPECT_TRUE(cloud.fields.empty());
}

namespace
{

size_t countValidPixels(const robot_sensor_msgs::Image& image)
{
  const auto* depth = reinterpret_cast<const uint16_t*>(image.data.data());
  size_t count = 0;
  for (size_t i = 0; i < image.width * image.height; ++i)
  {
    if (depth[i] != 0)
    {
      ++count;
    }
  }
  return count;
}

}  // namespace

TEST(DepthFrameFilter, RemovesEdgeHalo)
{
  const uint32_t width = 10;
  const uint32_t height = 10;

  // Left half at 1.0 m, right half at 2.0 m: a vertical object edge.
  robot_sensor_msgs::Image depth = makeFlatDepthImage(width, height, 1000);
  auto* data = reinterpret_cast<uint16_t*>(depth.data.data());
  for (uint32_t v = 0; v < height; ++v)
  {
    for (uint32_t u = 5; u < width; ++u)
    {
      data[v * width + u] = 2000;
    }
  }

  depth_image_proc::DepthFilterConfig config;
  config.edge_max_delta = 0.1;
  config.edge_dilation = 1;
  config.edge_window = 1;
  config.edge_invalid_border = false;
  config.temporal_min_frames = 0;

  depth_image_proc::DepthFrameFilter filter(config);
  filter.apply(depth);

  // Edge pixels (columns 4 and 5) plus a 1 px halo (columns 3 and 6) removed.
  EXPECT_EQ(countValidPixels(depth), (width - 4) * height);
  for (uint32_t v = 0; v < height; ++v)
  {
    for (uint32_t u = 3; u <= 6; ++u)
    {
      EXPECT_EQ(data[v * width + u], 0u) << "u=" << u << " v=" << v;
    }
    EXPECT_NE(data[v * width + 2], 0u);
    EXPECT_NE(data[v * width + 7], 0u);
  }
}

TEST(DepthFrameFilter, RemovesSmoothFlyingPixelRamp)
{
  const uint32_t width = 24;
  const uint32_t height = 6;

  // Foreground at 1.0 m (u 0-7), a smooth 80 mm/px ramp (u 8-15), background
  // at 1.64 m (u 16-23). Every adjacent step stays below edge_max_delta, so
  // only the wide-baseline test can see the jump.
  robot_sensor_msgs::Image depth = makeFlatDepthImage(width, height, 1000);
  auto* data = reinterpret_cast<uint16_t*>(depth.data.data());
  for (uint32_t v = 0; v < height; ++v)
  {
    for (uint32_t u = 8; u < 16; ++u)
    {
      data[v * width + u] = static_cast<uint16_t>(1000 + 80 * (u - 7));
    }
    for (uint32_t u = 16; u < width; ++u)
    {
      data[v * width + u] = 1640;
    }
  }

  depth_image_proc::DepthFilterConfig config;
  config.edge_max_delta = 0.1;
  config.edge_dilation = 0;
  config.edge_window = 4;
  config.edge_invalid_border = false;
  config.temporal_min_frames = 0;

  depth_image_proc::DepthFrameFilter filter(config);
  filter.apply(depth);

  // Spans where |z(u+4) - z(u)| > 0.1 m cover columns 5..17: the whole ramp
  // plus its shoulders goes, the flat surfaces on both sides stay.
  for (uint32_t v = 0; v < height; ++v)
  {
    for (uint32_t u = 0; u < width; ++u)
    {
      if (u >= 5 && u <= 17)
      {
        EXPECT_EQ(data[v * width + u], 0u) << "u=" << u << " v=" << v;
      }
      else
      {
        EXPECT_NE(data[v * width + u], 0u) << "u=" << u << " v=" << v;
      }
    }
  }
}

TEST(DepthFrameFilter, RemovesPixelsHuggingInvalidHoles)
{
  const uint32_t width = 10;
  const uint32_t height = 10;

  // Flat surface with a no-data band (column 5, rows 2-7), as stereo
  // matching leaves at occlusion boundaries. No valid-to-valid jump exists.
  robot_sensor_msgs::Image depth = makeFlatDepthImage(width, height, 2000);
  auto* data = reinterpret_cast<uint16_t*>(depth.data.data());
  for (uint32_t v = 2; v <= 7; ++v)
  {
    data[v * width + 5] = 0;
  }

  depth_image_proc::DepthFilterConfig config;
  config.edge_max_delta = 0.1;
  config.edge_dilation = 0;
  config.edge_window = 1;
  config.edge_invalid_border = true;
  config.temporal_min_frames = 0;

  depth_image_proc::DepthFrameFilter filter(config);
  filter.apply(depth);

  // Hole (6 px) plus its 4-connected valid border (14 px) are invalid.
  EXPECT_EQ(countValidPixels(depth), width * height - 20);
  for (uint32_t v = 2; v <= 7; ++v)
  {
    EXPECT_EQ(data[v * width + 4], 0u) << "v=" << v;
    EXPECT_EQ(data[v * width + 6], 0u) << "v=" << v;
    EXPECT_NE(data[v * width + 3], 0u) << "v=" << v;
    EXPECT_NE(data[v * width + 7], 0u) << "v=" << v;
  }
  EXPECT_EQ(data[1 * width + 5], 0u);
  EXPECT_EQ(data[8 * width + 5], 0u);
  EXPECT_NE(data[0 * width + 5], 0u);
  EXPECT_NE(data[9 * width + 5], 0u);
}

TEST(DepthFrameFilter, RejectsTransientPixels)
{
  const uint32_t width = 6;
  const uint32_t height = 6;

  depth_image_proc::DepthFilterConfig config;
  config.edge_max_delta = 0.0;
  config.temporal_max_delta = 0.06;
  config.temporal_min_frames = 2;

  depth_image_proc::DepthFrameFilter filter(config);

  // Frame 1: nothing has history yet, everything suppressed.
  robot_sensor_msgs::Image frame = makeFlatDepthImage(width, height, 2000);
  filter.apply(frame);
  EXPECT_EQ(countValidPixels(frame), 0u);

  // Frame 2: static scene is now stable and passes through.
  frame = makeFlatDepthImage(width, height, 2000);
  filter.apply(frame);
  EXPECT_EQ(countValidPixels(frame), width * height);

  // Frame 3: one pixel jumps 0.5 m (object crossing the view) -> rejected,
  // the static background stays.
  frame = makeFlatDepthImage(width, height, 2000);
  auto* data = reinterpret_cast<uint16_t*>(frame.data.data());
  data[3 * width + 3] = 2500;
  filter.apply(frame);
  EXPECT_EQ(countValidPixels(frame), width * height - 1);
  EXPECT_EQ(data[3 * width + 3], 0u);

  // Frame 4: the pixel holds its new depth -> accepted again.
  frame = makeFlatDepthImage(width, height, 2000);
  data = reinterpret_cast<uint16_t*>(frame.data.data());
  data[3 * width + 3] = 2500;
  filter.apply(frame);
  EXPECT_EQ(countValidPixels(frame), width * height);
}

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