// +-------------------------------- DISCLAIMER ---------------------------------+
// |                                                                             |
// | This program is provided to you free of charge "as such".                   |
// | You should not assume that this program is error-free or suitable for any   |
// | purpose whatsoever.                                                         |
// |                                                                             |
// | EURESYS does not give any representation, warranty or undertaking that this |
// | program is free of any defect or error or suitable for any purpose. EURESYS |
// | shall not be liable, in contract, in torts or otherwise, for any damages,   |
// | loss, costs, expenses or other claims for compensation, including those     |
// | asserted by third parties, arising out of or in connection with the use of  |
// | this program.                                                               |
// +-----------------------------------------------------------------------------+

#include "../Common/stdafx.h"
#include "../Common/Resource.h"
#include "AzureKinectGrabber.h"
#include "AzureKinectDialog.h"

AzureKinectGrabber::AzureKinectGrabber() : started(false), locked(false), frame_counter(0),
device(NULL), config(K4A_DEVICE_CONFIG_INIT_DISABLE_ALL), capture(NULL), depth_mode(K4A_DEPTH_MODE_NFOV_UNBINNED), frame_rate(K4A_FRAMES_PER_SECOND_5),
depth_image(NULL), xy_table(NULL)
{
  try
  {
    uint32_t device_count = k4a_device_get_installed_count();
    if (device_count == 0)
      throw(std::runtime_error("No Azure Kinect device detected"));

    if (k4a_device_open(K4A_DEVICE_DEFAULT, &device) != K4A_RESULT_SUCCEEDED)
      throw(std::runtime_error("Failed to open device"));

    ConfigureDevice();
  }
  catch (std::exception& e)
  {
    ::MessageBox(NULL, CString(e.what()), _T("Initialize"), MB_OK | MB_ICONSTOP);
  }
}

AzureKinectGrabber::~AzureKinectGrabber()
{
  k4a_image_release(xy_table);

  if (device != NULL)
    k4a_device_close(device);
}

void AzureKinectGrabber::Start(int num_frame)
{
  if (k4a_device_start_cameras(device, &config) != K4A_RESULT_SUCCEEDED)
    throw(std::runtime_error("Failed to start the camera"));
  started = true;
}

long long int AzureKinectGrabber::WaitFrame(int timeout_ms)
{
  try
  {
    // Get a capture
    switch (k4a_device_get_capture(device, &capture, timeout_ms))
    {
    case K4A_WAIT_RESULT_SUCCEEDED:
      break;
    case K4A_WAIT_RESULT_TIMEOUT:
      throw(std::runtime_error("Timed out waiting for a capture"));
    case K4A_WAIT_RESULT_FAILED:
      throw(std::runtime_error("Failed to read a capture"));
    }

    depth_image = k4a_capture_get_depth_image(capture);
    if (depth_image == 0)
      throw(std::runtime_error("Failed to get depth image from capture"));

    locked = true;
    frame_counter++;

    return frame_counter;
  }
  catch (std::exception& e)
  {
    ::MessageBox(NULL, CString(e.what()), _T("WaitFrame"), MB_OK | MB_ICONSTOP);
  }
  return 0;
}

void AzureKinectGrabber::ReleaseFrame()
{
  try
  {
    k4a_image_release(depth_image);
    k4a_capture_release(capture);

    locked = false;
  }
  catch (std::exception& e)
  {
    ::MessageBox(NULL, CString(e.what()), _T("ReleaseFrame"), MB_OK | MB_ICONSTOP);
  }
}

void AzureKinectGrabber::Stop()
{
  k4a_device_stop_cameras(device);
  started = false;
}

int AzureKinectGrabber::GetBitDepth()
{
  return 16;  // Depth values are always in 16 bits format
}

void AzureKinectGrabber::GetDepthMap(Easy3D::EDepthMap& map)
{
  if (!locked)
    throw(std::runtime_error("Frame not available"));
  if (map.GetType() != EImageType_Depth16)
    throw(std::runtime_error("Incompatible Depth Map format"));

  // copy depth_image to map
  int width = k4a_image_get_width_pixels(depth_image);
  int height = k4a_image_get_height_pixels(depth_image);
  map.SetSize(width, height);

  uint16_t *depth_data = (uint16_t *)k4a_image_get_buffer(depth_image);

  for (int y = 0; y < height; depth_data += width, ++y)
  {
    void* dst = map.GetBufferPtr(0, y);
    memcpy(dst, depth_data, 2 * width);
  }
}

void AzureKinectGrabber::GetPointCloud(Easy3D::EPointCloud& pc)
{
  if (!locked)
    throw(std::runtime_error("Frame not available"));

  int width = k4a_image_get_width_pixels(depth_image);
  int height = k4a_image_get_height_pixels(depth_image);

  // convert depth image to point cloud using pre computed table
  uint16_t *depth_data = (uint16_t *)k4a_image_get_buffer(depth_image);
  k4a_float2_t *xy_table_data = (k4a_float2_t *)k4a_image_get_buffer(xy_table);

  std::vector<Easy3D::E3DPoint> pts;
  pts.reserve(width * height);
  Easy3D::E3DPoint p;
  float z_max = -1.f;
  for (int i = 0; i < width * height; ++i)
  {
    if (depth_data[i] != 0 && !isnan(xy_table_data[i].xy.x) && !isnan(xy_table_data[i].xy.y))
    {
      p.X = xy_table_data[i].xy.x * (float)depth_data[i];
      p.Y = xy_table_data[i].xy.y * (float)depth_data[i];
      p.Z = (float)depth_data[i];
      if (p.Z > z_max)
        z_max = p.Z;
      pts.push_back(p);
    }
  }

  // invert Y and Z axis to get the origin at the furthest point, going up to the camera
  for (size_t i = 0; i < pts.size(); ++i)
  {
    pts[i].Y = -pts[i].Y;
    pts[i].Z = z_max - pts[i].Z;
  }

  point_cloud.Clear();
  point_cloud.AddPoints(pts);

  pc = point_cloud;
}

void AzureKinectGrabber::GetZMap(Easy3D::EZMap& zmap)
{
  // Azure Kinect SDK cannot generate a ZMap
  // Use Easy3D::EPointCloudToZMapConverter instead

  Easy3D::EPointCloudToZMapConverter converter;
  converter.Convert(point_cloud, &zmap);
}

// code snippet from Azure Kinect SDK example "fastpointcloud"
static void create_xy_table(const k4a_calibration_t *calibration, k4a_image_t xy_table)
{
  k4a_float2_t *table_data = (k4a_float2_t *)(void *)k4a_image_get_buffer(xy_table);

  int width = calibration->depth_camera_calibration.resolution_width;
  int height = calibration->depth_camera_calibration.resolution_height;

  k4a_float2_t p;
  k4a_float3_t ray;
  int valid;

  for (int y = 0, idx = 0; y < height; y++)
  {
    p.xy.y = (float)y;
    for (int x = 0; x < width; x++, idx++)
    {
      p.xy.x = (float)x;

      k4a_calibration_2d_to_3d(
        calibration, &p, 1.f, K4A_CALIBRATION_TYPE_DEPTH, K4A_CALIBRATION_TYPE_DEPTH, &ray, &valid);

      if (valid)
      {
        table_data[idx].xy.x = ray.xyz.x;
        table_data[idx].xy.y = ray.xyz.y;
      }
      else
      {
        table_data[idx].xy.x = nanf("");
        table_data[idx].xy.y = nanf("");
      }
    }
  }
}

void AzureKinectGrabber::ConfigureDevice()
{
  config.depth_mode = depth_mode;
  config.camera_fps = frame_rate;

  if (k4a_device_get_calibration(device, config.depth_mode, config.color_resolution, &calibration) != K4A_RESULT_SUCCEEDED)
    throw(std::runtime_error("Failed to get calibration"));

  k4a_image_create(K4A_IMAGE_FORMAT_CUSTOM,
    calibration.depth_camera_calibration.resolution_width,
    calibration.depth_camera_calibration.resolution_height,
    calibration.depth_camera_calibration.resolution_width * (int)sizeof(k4a_float2_t),
    &xy_table);

  create_xy_table(&calibration, xy_table);
}

void AzureKinectGrabber::InvokeParameterPanel()
{
  AzureKinectParamDialog param_dialog(depth_mode, frame_rate);

  if (param_dialog.DoModal() == IDOK)
  {
    depth_mode = (k4a_depth_mode_t)param_dialog.depth_mode;
    frame_rate = (k4a_fps_t)param_dialog.frame_rate;

    ConfigureDevice();
  }
}
