#include "XnCppWrapper.h"
#include "opencv/cv.h"
#include "opencv/highgui.h"
#include "stdio.h"

#define MAX_DEPTH 10000

using namespace xn;

//Global Variables
xn::Context context;
DepthGenerator Xn_depth;
ImageGenerator Xn_image;
XnStatus nRetVal = XN_STATUS_OK;


unsigned short depth[MAX_DEPTH];
char *depth_data;

void raw2depth(){
    int i;
    for ( i=0; i<MAX_DEPTH; i++) {
        //float dist =100/(-0.0030711016*i + 3.3309495161); //ros distance (cm)
        //float dist2 = 12.36 * tanf(i/ 2842.5 + 1.1863);//Stéphane Magnenat distance (cm)
        float v = (float)i/MAX_DEPTH;//for visualization purposes only
        v = powf(v, 2);
        v = v*36*256;
        depth[i] = v;
        //printf("%f    ",dist2);

        /*
         * XnUInt64 F_;
         * XnDouble pixel_size_;
        // get the focal length in mm (ZPS = zero plane distance)
            depth_.GetIntProperty ("ZPD", F_)
        // get the pixel size in mm ("ZPPS" = pixel size at zero plane)
            depth_.GetRealProperty ("ZPPS", pixel_size_)
            X = (u - 320) * depth_md_[k] * pixel_size_ * 0.001 / F_;
            Y = (v - 240) * depth_md_[k] * pixel_size_ * 0.001 / F_;
            Z = depth_md_[k] * 0.001; // from mm in meters!
         */
    }
}

void depth2rgb(const XnDepthPixel* Xn_disparity){
    int i;
    //const unsigned short *disparity = Xn_disparity;

    for (i=0; i<307200; i++) {
        int pval = depth[Xn_disparity[i]];
        //printf("%d: %u %d \n",i,*Xn_disparity,depth[*Xn_disparity]);
        //fflush(stdout);
        int lb = pval & 0xff;
        switch (pval>>8) {
            case 0:
                depth_data[3*i+0] = 255;
                depth_data[3*i+1] = 255-lb;
                depth_data[3*i+2] = 255-lb;
                break;
            case 1:
                depth_data[3*i+0] = 255;
                depth_data[3*i+1] = lb;
                depth_data[3*i+2] = 0;
                break;
            case 2:
                depth_data[3*i+0] = 255-lb;
                depth_data[3*i+1] = 255;
                depth_data[3*i+2] = 0;
                break;
            case 3:
                depth_data[3*i+0] = 0;
                depth_data[3*i+1] = 255;
                depth_data[3*i+2] = lb;
                break;
            case 4:
                depth_data[3*i+0] = 0;
                depth_data[3*i+1] = 255-lb;
                depth_data[3*i+2] = 255;
                break;
            case 5:
                depth_data[3*i+0] = 0;
                depth_data[3*i+1] = 0;
                depth_data[3*i+2] = 255-lb;
                break;
            default:
                depth_data[3*i+0] = 0;
                depth_data[3*i+1] = 0;
                depth_data[3*i+2] = 0;
                break;
        }
    }
}

void rgbd_init(){
/// Initialize context object
nRetVal = context.Init();
if (nRetVal != XN_STATUS_OK){
        printf("Failed to initialize context: %s\n", xnGetStatusString(nRetVal));
    }

/// Create a DepthGenerator node
nRetVal = Xn_depth.Create(context);
if (nRetVal != XN_STATUS_OK){
        printf("Failed to create depth generator: %s\n", xnGetStatusString(nRetVal));
    }

/// Create a ImageGenerator node
nRetVal = Xn_image.Create(context);
if (nRetVal != XN_STATUS_OK){
        printf("Failed to create image generator: %s\n", xnGetStatusString(nRetVal));
    }

/// Make it start generating data
nRetVal = context.StartGeneratingAll();
if (nRetVal != XN_STATUS_OK){
        printf("Failed generating data: %s\n", xnGetStatusString(nRetVal));
    }

///Sync the DepthGenerator with the ImageGenerator
//~ nRetVal = Xn_depth.GetFrameSyncCap().FrameSyncWith(Xn_image);
//~ if (nRetVal != XN_STATUS_OK){
        //~ printf("Failed to sync Depth and RGB streams: %s\n", xnGetStatusString(nRetVal));
    //~ }
/// Set the view point of the DepthGenerator to match the ImageGenerator
//~ nRetVal = Xn_depth.GetAlternativeViewPointCap().SetViewPoint(Xn_image);
//~ if (nRetVal != XN_STATUS_OK){
        //~ printf("Failed to match Depth and RGB points of view: %s\n", xnGetStatusString(nRetVal));
    //~ }
}

int main(int argc, char **argv) {

//RGB-D Sensor Initialization
rgbd_init();

//OpenCV stuff
cvNamedWindow("RGB", CV_WINDOW_AUTOSIZE);
cvNamedWindow("Depth", CV_WINDOW_AUTOSIZE);
IplImage *rgbimg = cvCreateImageHeader(cvSize(640,480), 8, 3);
IplImage *depthimg = cvCreateImageHeader(cvSize(640,480), 8, 3);

raw2depth();
depth_data = (char*)malloc(640*480*3);

while (true) {
    // Wait for new data to be available
    nRetVal = context.WaitAndUpdateAll();
    if (nRetVal != XN_STATUS_OK)
    {
        printf("Failed updating data: %s\n", xnGetStatusString(nRetVal));
        continue;
    }

    // Take current depth map
    const XnDepthPixel* pDepthMap = Xn_depth.GetDepthMap();
    // Take current image
    const XnRGB24Pixel* pImage = Xn_image.GetRGB24ImageMap();

    //process depth map
    depth2rgb(pDepthMap);
    cvSetData(depthimg,depth_data, 640*3);
    cvShowImage("Depth",depthimg);

    //process image data
    XnRGB24Pixel* ucpImage = const_cast<XnRGB24Pixel*> (pImage);
    cvSetData(rgbimg,ucpImage, 640*3);
    cvCvtColor(rgbimg,rgbimg,CV_RGB2BGR);
    cvShowImage("RGB",rgbimg);

    char key =cvWaitKey(30);
    if (key==13)
        continue;
    if (key==27)
        break;
}

// Clean-up
context.Shutdown();
}