#include <mex.h>
#include <matrix.h>

#include <opencv2/opencv.hpp>

std::string readMATLABString(const mxArray *prhs)
{
    int buflen = (mxGetM(prhs) * mxGetN(prhs)) + 1;
    char *input_buf = (char *)mxCalloc(buflen, sizeof(char));
    mxGetString(prhs, input_buf, buflen);

    std::string toReturn(input_buf);
    mxFree(input_buf);

    return toReturn;
}

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{
    std::string vidfilename = readMATLABString(prhs[0]);

    double *frameIdxs = mxGetPr(prhs[1]);
    const mwSize *dims = mxGetDimensions(prhs[1]);

    cv::VideoCapture vc(vidfilename);
    int w = vc.get(cv::CAP_PROP_FRAME_WIDTH);
    int h = vc.get(cv::CAP_PROP_FRAME_HEIGHT);
    int nf = vc.get(cv::CAP_PROP_FRAME_COUNT);
    int tw = w, th = h;
    if(nrhs > 2)
    {
        double * tres = mxGetPr(prhs[2]);
        tw = tres[0], th = tres[1];
    }

    mwSize outdim[4] = {th, tw, 3, dims[1]};
    plhs[0] = mxCreateNumericArray(4, outdim, mxSINGLE_CLASS, mxREAL);
    float *data = (float *)mxGetPr(plhs[0]);
    cv::Mat m(h, w, CV_8UC3);
    cv::Mat timg(th, tw, CV_32FC3);

    for(int i=0; i<dims[1]; i++)
    {
        long frame = frameIdxs[i];

        vc.set(cv::CAP_PROP_POS_FRAMES, frame);
        if (frame >= nf)
            m.setTo(0);
        else
            vc >> m;
        if(m.size().area() > 0)
            cv::resize(m, timg, timg.size(), 0, 0, cv::INTER_AREA);
        float *ptr = &data[i*tw*th*3];
        unsigned char *srcptr = timg.data;
        for(int x=0, z=0; x<tw; x++)
        {
            for(int y=0; y<th; y++, z++)
            {
                ptr[z] = srcptr[(y*tw+x)*3+2] / 255.f;
                ptr[tw*th+z] = srcptr[(y*tw+x)*3+1] / 255.f;
                ptr[tw*th*2+z] = srcptr[(y*tw+x)*3] / 255.f;
            }
        }
    }
}