//***************************************************************************************/
// 
// File name: KinectCamera.cpp
// Location: See Matrox Example Launcher in the MIL Control Center
// 
//
// Synopsis:  Implementation of the CKinectCamera class, a class that manages
//            the Microsoft Kinect.
//
// Copyright (C) Matrox Electronic Systems Ltd., 1992-2016.
// All Rights Reserved

#include <mil.h>
#include <Windows.h>
#include "KinectCamera.h"

#if USE_REAL_CAMERA
#if (KINECT_CAMERA_VERSION == 1)

// Include the header for the Kinect. Define the static_assert function for compiler < VS2010
#pragma comment(lib, "Kinect10.lib")
#if _MSC_VER < 1600
#define static_assert(a, b)
#endif
#include "NuiApi.h"

// Constants.
static const NUI_IMAGE_RESOLUTION KINECT_COLOR_IMAGE_RESOLUTION = NUI_IMAGE_RESOLUTION_640x480;
static const NUI_IMAGE_RESOLUTION KINECT_DEPTH_IMAGE_RESOLUTION = NUI_IMAGE_RESOLUTION_640x480;
static const MIL_INT KINECT_COLOR_NB_FRAMES = 2;
static const MIL_INT KINECT_DEPTH_NB_FRAMES = 2;
static const MIL_INT KINECT_IR_NB_FRAMES = 2;

//*****************************************************************************
// Constructor.
//*****************************************************************************
CKinectCamera::CKinectCamera(MIL_ID MilSystem)
   : CKinectCameraBase(MilSystem)
   {
   // Create the sensor.
   if(FAILED(NuiCreateSensorByIndex(0, &m_pNuiSensor)))
      m_SensorStatus = KINECT_CAMERA_NOT_OK;  
   }


//*****************************************************************************
// Destructor.
//*****************************************************************************
CKinectCamera::~CKinectCamera()
   {
   // Close down the camera.
   CloseDown();
   }

void CKinectCamera::ShutdownCamera()
   {
   // Shut down the sensor.
   if(m_pNuiSensor)
      m_pNuiSensor->NuiShutdown();

   // Release the sensor. This should also close the image streams handles.
   if(m_pNuiSensor)
      {
      m_pNuiSensor->Release();
      m_pNuiSensor = NULL;
      }
   }

//*****************************************************************************
// Initialization function.
//*****************************************************************************
MIL_INT CKinectCamera::InitCamera(ColorStreamTypeEnum ColorStreamType, bool UseDepth)
   {
   if(m_SensorStatus == KINECT_CAMERA_NOT_OK)
      return m_SensorStatus;

   // Set whether to use the IR camera.
   m_ColorStreamType = ColorStreamType;

   // Open/Close the IR Emitter. 
   if(FAILED(m_pNuiSensor->NuiSetForceInfraredEmitterOff(!UseDepth)))
      return KINECT_CAMERA_NOT_OK;  

   // ReInitialize the Kinect sensor.
   if(FAILED(m_pNuiSensor->NuiInitialize(NUI_INITIALIZE_FLAG_USES_COLOR + (UseDepth ? NUI_INITIALIZE_FLAG_USES_DEPTH : 0))))
      return KINECT_CAMERA_NOT_OK;

   // Open the color/infrared streams of the Kinect.
   if(FAILED(m_pNuiSensor->NuiImageStreamOpen(m_ColorStreamType == enColor? NUI_IMAGE_TYPE_COLOR : NUI_IMAGE_TYPE_COLOR_INFRARED,
                                              KINECT_COLOR_IMAGE_RESOLUTION,
                                              0,
                                              KINECT_COLOR_NB_FRAMES,
                                              m_NextColorFrameEvent,
                                              &m_ColorStreamHandle)))
      return KINECT_CAMERA_NOT_OK;

   // Open the depth stream(do not open if in IR mode).
   if(UseDepth)
      {
      if(FAILED(m_pNuiSensor->NuiImageStreamOpen(NUI_IMAGE_TYPE_DEPTH,
                                                 KINECT_DEPTH_IMAGE_RESOLUTION,
                                                 0,
                                                 KINECT_DEPTH_NB_FRAMES,
                                                 m_NextDepthFrameEvent,
                                                 &m_DepthStreamHandle)))
            return KINECT_CAMERA_NOT_OK;

      // Put the depth steam in near mode.
      if(FAILED(NuiImageStreamSetImageFrameFlags(m_DepthStreamHandle, NUI_IMAGE_STREAM_FLAG_ENABLE_NEAR_MODE)))
         return KINECT_CAMERA_NOT_OK;
      }   
   return KINECT_CAMERA_OK;
   }

//*****************************************************************************
// Function that processes the Kinect buffer.
//*****************************************************************************
void CKinectCamera::ProcessKinectBufferColor(MIL_DIG_HOOK_FUNCTION_PTR HookHandlerPtr, void* UserDataPtr)
   {
   if(m_ColorStreamType == enColor)
      ProcessKinectBuffer(m_ColorStreamHandle, HookHandlerPtr, UserDataPtr, 3, 8+M_UNSIGNED, M_BGR32 + M_PACKED);
   else
      ProcessKinectBuffer(m_ColorStreamHandle, HookHandlerPtr, UserDataPtr, 1, 16+M_UNSIGNED, M_NULL);   
      
   }

void CKinectCamera::ProcessKinectBufferDepth(MIL_DIG_HOOK_FUNCTION_PTR HookHandlerPtr, void* UserDataPtr)
   {
   ProcessKinectBuffer(m_DepthStreamHandle, HookHandlerPtr, UserDataPtr, 1, 16+M_UNSIGNED, M_NULL);
   }

void CKinectCamera::ProcessKinectBuffer(HANDLE NuiImageStream,
                                        MIL_DIG_HOOK_FUNCTION_PTR HookHandlerPtr,
                                        void* UserDataPtr,
                                        MIL_INT NbBands,
                                        MIL_INT Type,
                                        MIL_INT Attribute)
   {
   NUI_IMAGE_FRAME NuiImageFrame;
   NUI_LOCKED_RECT NuiLockedRect;
   NUI_SURFACE_DESC NuiSurfaceDesc;

   // Get the image frame.
   if(FAILED(m_pNuiSensor->NuiImageStreamGetNextFrame(NuiImageStream, 0,  &NuiImageFrame)))
      return;

   // Get the surface descriptor.
   NuiImageFrame.pFrameTexture->GetLevelDesc(0, &NuiSurfaceDesc);

   // Get the locked rect of the image frame.
   NuiImageFrame.pFrameTexture->LockRect(0, &NuiLockedRect, NULL, 0);   

   // Create a mil image on the Kinect data. 
   MIL_ID MilKinectImage = MbufCreateColor(GetMilSystem(),
                                           NbBands,
                                           NuiSurfaceDesc.Width,
                                           NuiSurfaceDesc.Height,
                                           Type,
                                           M_IMAGE + M_PROC + Attribute,
                                           M_HOST_ADDRESS + M_PITCH_BYTE,
                                           NuiLockedRect.Pitch,
                                           (void**)(&NuiLockedRect.pBits),
                                           M_NULL);

   // Call the Mil hook function. 
   (*HookHandlerPtr)(0, MilKinectImage, UserDataPtr);

   // Free the created buffer.
   MbufFree(MilKinectImage);

   // Release the locked rect of the image frame.
   NuiImageFrame.pFrameTexture->UnlockRect(0);

   // Release the image frame.
   m_pNuiSensor->NuiImageStreamReleaseFrame(NuiImageStream, &NuiImageFrame);
   }

#elif (KINECT_CAMERA_VERSION == 2)

// Visual studio 2012 and above is required to run the Kinect 2.0
#if _MSC_VER < 1700
   #error Visual Studio 2012 or higher is required to use the Kinect 2.0 SDK 
#endif

// Include the header for the Kinect.
#pragma comment(lib, "Kinect20.lib")
#include "Kinect.h"

// Constants.
static const MIL_INT KINECT_COLOR_NB_FRAMES = 2;
static const MIL_INT KINECT_DEPTH_NB_FRAMES = 2;
static const MIL_INT KINECT_IR_NB_FRAMES = 2;

//*****************************************************************************
// Constructor.
//*****************************************************************************
CKinectOneCamera::CKinectOneCamera(MIL_ID MilSystem)
   : CKinectCameraBase(MilSystem)
   {
   // Create the sensor.
   if (FAILED(GetDefaultKinectSensor(&m_pKinectSensor)))
      m_SensorStatus = KINECT_CAMERA_NOT_OK;
   }

//*****************************************************************************
// Destructor.
//*****************************************************************************
CKinectOneCamera::~CKinectOneCamera()
   {
   // Close down the camera.
   CloseDown();
   }

void CKinectOneCamera::ShutdownCamera()
   {
   // Shut down the sensor.
   if (m_pKinectSensor)
      m_pKinectSensor->Close();

   // Release the sensor. This should also close the image streams handles.
   if (m_pKinectSensor)
      {
      m_pKinectSensor->Release();
      m_pKinectSensor = NULL;
      }
   }

//*****************************************************************************
// Initialization function.
//*****************************************************************************
MIL_INT CKinectOneCamera::InitCamera(ColorStreamTypeEnum ColorStreamType, bool UseDepth)
   {
   if (m_SensorStatus == KINECT_CAMERA_NOT_OK)
      return m_SensorStatus;

   // Set whether to use the IR camera.
   m_ColorStreamType = ColorStreamType;
      
   // Open the color/infrared streams of the Kinect.
   if (ColorStreamType == enColor)
      {
      IColorFrameSource* pColorFrameSource = NULL;
      if (FAILED(m_pKinectSensor->get_ColorFrameSource(&pColorFrameSource)))
         return KINECT_CAMERA_NOT_OK;

      if (FAILED(pColorFrameSource->OpenReader(&m_ColorFrameReader)))
         return KINECT_CAMERA_NOT_OK;

      WAITABLE_HANDLE NextColorFrameWaitableEvent;
      if (FAILED(m_ColorFrameReader->SubscribeFrameArrived(&NextColorFrameWaitableEvent)))
         return KINECT_CAMERA_NOT_OK;
      m_NextColorFrameEvent = (HANDLE)NextColorFrameWaitableEvent;
      }
   else if (ColorStreamType == enIR)
      {
      IInfraredFrameSource* pInfraredFrameSource = NULL;
      if (FAILED(m_pKinectSensor->get_InfraredFrameSource(&pInfraredFrameSource)))
         return KINECT_CAMERA_NOT_OK;

      if (FAILED(pInfraredFrameSource->OpenReader(&m_IRFrameReader)))
         return KINECT_CAMERA_NOT_OK;

      WAITABLE_HANDLE NextColorFrameWaitableEvent;
      if (FAILED(m_IRFrameReader->SubscribeFrameArrived(&NextColorFrameWaitableEvent)))
         return KINECT_CAMERA_NOT_OK;
      m_NextColorFrameEvent = (HANDLE)NextColorFrameWaitableEvent;
      }

   // Open the depth stream
   if (UseDepth)
      {
      IDepthFrameSource* pDepthFrameSource = NULL;
      if (FAILED(m_pKinectSensor->get_DepthFrameSource(&pDepthFrameSource)))
         return KINECT_CAMERA_NOT_OK;

      if (FAILED(pDepthFrameSource->OpenReader(&m_DepthFrameReader)))
         return KINECT_CAMERA_NOT_OK;

      WAITABLE_HANDLE NextDepthFrameWaitableEvent;
      if (FAILED(m_DepthFrameReader->SubscribeFrameArrived(&NextDepthFrameWaitableEvent)))
         return KINECT_CAMERA_NOT_OK;
      m_NextDepthFrameEvent = (HANDLE)NextDepthFrameWaitableEvent;

      // Put the depth steam in near mode.
      }

   // Open the kinect sensor
   if (FAILED(m_pKinectSensor->Open()))
      return KINECT_CAMERA_NOT_OK;

   return KINECT_CAMERA_OK;
   }

//*****************************************************************************
// Function that processes the Kinect buffer.
//*****************************************************************************

// Explicit template instanciation
template void CKinectOneCamera::ProcessKinectBuffer<IColorFrameReader, IColorFrameArrivedEventArgs, IColorFrameReference, IColorFrame>             
                          (IColorFrameReader* FrameReader, 
                          HANDLE NextFrameEvent, 
                          MIL_DIG_HOOK_FUNCTION_PTR HookHandlerPtr, 
                          void* UserDataPtr, MIL_INT NbBands, 
                          MIL_INT Type, MIL_INT Attribute);

template void CKinectOneCamera::ProcessKinectBuffer<IInfraredFrameReader, IInfraredFrameArrivedEventArgs, IInfraredFrameReference, IInfraredFrame> 
                          (IInfraredFrameReader* FrameReader, 
                          HANDLE NextFrameEvent, 
                          MIL_DIG_HOOK_FUNCTION_PTR HookHandlerPtr, 
                          void* UserDataPtr, 
                          MIL_INT NbBands, 
                          MIL_INT Type, 
                          MIL_INT Attribute);

template void CKinectOneCamera::ProcessKinectBuffer<IDepthFrameReader, IDepthFrameArrivedEventArgs, IDepthFrameReference, IDepthFrame>             
                         (IDepthFrameReader* FrameReader, 
                         HANDLE NextFrameEvent, 
                         MIL_DIG_HOOK_FUNCTION_PTR HookHandlerPtr, 
                         void* UserDataPtr, 
                         MIL_INT NbBands, 
                         MIL_INT Type, 
                         MIL_INT Attribute);

template <>
void CKinectOneCamera::GetFrameData<IColorFrame>(IColorFrame* pFrame, void** ppFrameData)
   {
   UINT nBufferSize;   
   pFrame->AccessRawUnderlyingBuffer(&nBufferSize, reinterpret_cast<BYTE**>(ppFrameData));
   }

template <class IFrame>
void CKinectOneCamera::GetFrameData(IFrame* pFrame, void** ppFrameData)
   {
   UINT nBufferSize;
   pFrame->AccessUnderlyingBuffer(&nBufferSize, reinterpret_cast<UINT16**>(ppFrameData));
   }


void CKinectOneCamera::ProcessKinectBufferColor(MIL_DIG_HOOK_FUNCTION_PTR HookHandlerPtr, void* UserDataPtr)
   {
   if (m_ColorStreamType == enColor)
      ProcessKinectBuffer<IColorFrameReader, IColorFrameArrivedEventArgs, IColorFrameReference, IColorFrame>(m_ColorFrameReader, (HANDLE)m_NextColorFrameEvent, HookHandlerPtr, UserDataPtr, 3, 8 + M_UNSIGNED, M_YUV16 + M_PACKED);
   else
      ProcessKinectBuffer<IInfraredFrameReader, IInfraredFrameArrivedEventArgs, IInfraredFrameReference, IInfraredFrame>(m_IRFrameReader, (HANDLE)m_NextColorFrameEvent, HookHandlerPtr, UserDataPtr, 1, 16 + M_UNSIGNED, M_NULL);

   }

void CKinectOneCamera::ProcessKinectBufferDepth(MIL_DIG_HOOK_FUNCTION_PTR HookHandlerPtr, void* UserDataPtr)
   {
   ProcessKinectBuffer<IDepthFrameReader, IDepthFrameArrivedEventArgs, IDepthFrameReference, IDepthFrame>(m_DepthFrameReader, (HANDLE)m_NextDepthFrameEvent, HookHandlerPtr, UserDataPtr, 1, 16 + M_UNSIGNED, M_NULL);
   }

template <class IFrameReader, class IFrameArrivedEventArgs, class IFrameReference, class IFrame>
void CKinectOneCamera::ProcessKinectBuffer(IFrameReader* pFrameReader,
                                           HANDLE NextFrameEvent,
                                           MIL_DIG_HOOK_FUNCTION_PTR HookHandlerPtr,
                                           void* UserDataPtr,
                                           MIL_INT NbBands,
                                           MIL_INT Type,
                                           MIL_INT Attribute)
   {
   // Get the frame arguments.
   IFrameArrivedEventArgs* pFrameEventArgs = NULL;
   pFrameReader->GetFrameArrivedEventData((WAITABLE_HANDLE)NextFrameEvent, &pFrameEventArgs);

   // Get the frame reference.
   IFrameReference* pFrameRef = NULL;
   pFrameEventArgs->get_FrameReference(&pFrameRef);

   // Acquire the frame.
   IFrame* pFrame = NULL;
   HRESULT hr = pFrameRef->AcquireFrame(&pFrame);

   if (SUCCEEDED(hr))
      {
      // Get the frame description.
      IFrameDescription* pFrameDescription = NULL;
      pFrame->get_FrameDescription(&pFrameDescription);

      void* pFrameData = NULL;
      GetFrameData(pFrame, &pFrameData);   

      // Create a mil image on the Kinect data. 
      int Width;
      int Height;
      unsigned int BytesPerPixels;
      pFrameDescription->get_Width(&Width);
      pFrameDescription->get_Height(&Height);
      pFrameDescription->get_BytesPerPixel(&BytesPerPixels);
      MIL_ID MilKinectImage = MbufCreateColor(GetMilSystem(),
                                              NbBands,
                                              Width,
                                              Height,
                                              Type,
                                              M_IMAGE + M_PROC + Attribute,
                                              M_HOST_ADDRESS + M_PITCH_BYTE,
                                              BytesPerPixels * Width,
                                              (void**)(&pFrameData),
                                              M_NULL);

      // Call the Mil hook function. 
      (*HookHandlerPtr)(0, MilKinectImage, UserDataPtr);

      // Free the created buffer.
      MbufFree(MilKinectImage);

      // Release the frame
      pFrame->Release();
      }

   // Release the frame reference
   pFrameRef->Release();
   }

#endif // KINECT_CAMERA_VERSION
#endif // USE_REAL_CAMERA