// +-------------------------------- 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 "NerianGrabber.h"
#include "NerianDialog.h"

using namespace visiontransfer;

NerianGrabber::NerianGrabber() : image_transfer(NULL), started(false), locked(false), frame_counter(0), max_distance(4.f)
{
  try
  {
    // Search for Nerian stereo devices
    DeviceEnumeration deviceEnum;
    devices = deviceEnum.discoverDevices();
    if (devices.size() == 0)
      throw(std::runtime_error("No Nerian device found"));

    // Create an async transfer object that receives data from the first detected device
    image_transfer = new AsyncTransfer(devices[0]);
  }
  catch (std::exception& e)
  {
    ::MessageBox(NULL, CString(e.what()), _T("Initialize"), MB_OK | MB_ICONSTOP);
  }
}

NerianGrabber::~NerianGrabber()
{
  // Clean up
  delete image_transfer;
}

void NerianGrabber::Start(int num_frame)
{
  started = true;
}

long long int NerianGrabber::WaitFrame(int timeout_ms)
{
  try
  {
    // Get image
    while (!image_transfer->collectReceivedImageSet(image_set, 0.1 /*timeout*/))
    {
      // Keep on trying until reception is successful
    }

    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 NerianGrabber::ReleaseFrame()
{
  locked = false;
}

void NerianGrabber::Stop()
{
  started = false;
}

int NerianGrabber::GetBitDepth()
{
  return 16;
}

void NerianGrabber::GetDepthMap(Easy3D::EDepthMap& map)
{
  if (!locked)
    throw(std::runtime_error("Frame not available"));

  try
  {

    // check format
    int index = image_set.getIndexOf(ImageSet::IMAGE_DISPARITY);
    ImageSet::ImageFormat format = image_set.getPixelFormat(index);
    if (format != ImageSet::FORMAT_12_BIT_MONO)
      throw(std::runtime_error("Unsupported image format"));

    // get data
    int width = image_set.getWidth();
    int height = image_set.getHeight();
    int stride = image_set.getRowStride(ImageSet::IMAGE_DISPARITY);
    const uint8_t* pixel_data = image_set.getPixelData(index);

    // setup a depth map
    Easy3D::EDepthMap16& map16 = dynamic_cast<Easy3D::EDepthMap16&>(map);
    map16.SetSize(width, height);
    uint16_t undef = map16.GetUndefinedValue().Value;
    map16.SetZResolution(1.f / 256);  // the source is 8.4 fixed point format, shifted to 8.8 in the depth map

    // Copy disparity data to EDepthMap8
    for (int y = 0; y < height; ++y)
    {
      const uint16_t* src = (const uint16_t*)(pixel_data + y * stride);
      uint16_t* dst = (uint16_t*)map16.GetBufferPtr(0, y);
      for (int x = 0; x < width; ++x)
      {
        if (src[x] >= 4095)     // undefined point
          dst[x] = undef;
        else
          dst[x] = src[x] << 4; // from 12bits to 16bits values
      }
    }
  }
  catch (const EException& exc)
  {
    ::MessageBox(NULL, CString(exc.what()), _T("Get Depth Map"), MB_OK | MB_DEFAULT_DESKTOP_ONLY);
  }
}

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

  try
  {
    int width = image_set.getWidth();
    int height = image_set.getHeight();
    int npix = width * height;
    float* xyzptr = recon_3d.createPointMap(image_set);

    // Push valid 3D points to an EPointCloud
    // Z origin is at max distance
    std::vector<Easy3D::E3DPoint> pts;
    pts.reserve(npix);
    pc.Clear();

    ImageSet::ImageType colImg = (image_set.hasImageType(ImageSet::IMAGE_COLOR)) ?
      ImageSet::IMAGE_COLOR : ImageSet::IMAGE_LEFT;
    bool colorEnabled = (image_set.hasImageType(colImg));

    // Convert the color information if enabled
    if (colorEnabled)
    {
      std::vector<EC24A> colors;
      colors.reserve(npix);

      unsigned char* pixel = image_set.getPixelData(colImg);
      ImageSet::ImageFormat image_format = image_set.getPixelFormat(colImg);
      int colorOffset;
      switch (image_format)
      {
      case ImageSet::FORMAT_8_BIT_MONO:
        colorOffset = 1;
        break;
      case ImageSet::FORMAT_12_BIT_MONO:
        colorOffset = 2;
        break;
      case ImageSet::FORMAT_8_BIT_RGB:
        colorOffset = 3;
        break;
      default: throw std::runtime_error("Illegal pixel format");
      }

      for (int i = 0; i < npix; ++i, xyzptr += 4, pixel += colorOffset)
      {
        if (xyzptr[2] < max_distance)
        {
          pts.emplace_back(xyzptr[0], xyzptr[1], max_distance - xyzptr[2]);

          switch (image_format)
          {
          case ImageSet::FORMAT_8_BIT_MONO:
            colors.emplace_back(*pixel, *pixel, *pixel, 255u);
            break;
          case ImageSet::FORMAT_12_BIT_MONO:
          {
            uint8_t colorBW8 = static_cast<uint8_t>(*reinterpret_cast<uint16_t*>(pixel) >> 4);
            colors.emplace_back(colorBW8, colorBW8, colorBW8, 255u);
            break;
          }
          case ImageSet::FORMAT_8_BIT_RGB:
            colors.emplace_back(*pixel, pixel[1], pixel[2], 255u);
            break;
          }
        }
      }

      pc.AddPoints(pts);
      pc.FillAttributeBufferEC24A(static_cast<int>(Easy3D::E3DAttribute_Color), colors);
    }
    else
    {
      for (int i = 0; i < npix; ++i, xyzptr += 4)
        if (xyzptr[2] < max_distance)
          pts.emplace_back(xyzptr[0], xyzptr[1], max_distance - xyzptr[2]);

      pc.AddPoints(pts);
    }

    // for zmap generation, color not needed
    point_cloud.Clear();
    point_cloud.AddPoints(pts);
  }
  catch (const EException& exc)
  {
    ::MessageBox(NULL, CString(exc.what()), _T("Get Point Cloud"), MB_OK | MB_DEFAULT_DESKTOP_ONLY);
  }
}

void NerianGrabber::GetZMap(Easy3D::EZMap& zmap)
{
  // Nerian SDK cannot generate a ZMap
  // Use Easy3D::EPointCloudToZMapConverter instead

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

void NerianGrabber::InvokeParameterPanel()
{
  NerianDialog param_dialog(max_distance);

  if (param_dialog.DoModal() == IDOK)
  {
    max_distance = param_dialog.max_distance;
  }
}
