/*
* Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions
* are met:
*  * Redistributions of source code must retain the above copyright
*    notice, this list of conditions and the following disclaimer.
*  * Redistributions in binary form must reproduce the above copyright
*    notice, this list of conditions and the following disclaimer in the
*    documentation and/or other materials provided with the distribution.
*  * Neither the name of NVIDIA CORPORATION nor the names of its
*    contributors may be used to endorse or promote products derived
*    from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
* EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
* PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
* CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
* EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
* PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
* PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
* OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <opencv2/core/version.hpp>
#if CV_MAJOR_VERSION >= 3
#    include <opencv2/imgcodecs.hpp>
#    include <opencv2/videoio.hpp>
#else
#    include <opencv2/highgui/highgui.hpp>
#endif

#include <opencv2/imgproc/imgproc.hpp>

#include <vpi/Context.h>
#include <vpi/Event.h>
#include <vpi/Image.h>
#include <vpi/Stream.h>
#include <vpi/algo/ImageFormatConverter.h>
#include <vpi/algo/PerspectiveImageWarp.h>

#include <algorithm>
#include <cstring> // for memset
#include <fstream>
#include <iostream>
#include <random>

#define CHECK_STATUS(STMT)                                      \
    do                                                          \
    {                                                           \
        VPIStatus status = (STMT);                              \
        if (status != VPI_SUCCESS)                              \
        {                                                       \
            throw std::runtime_error(vpiStatusGetName(status)); \
        }                                                       \
    } while (0);

static void MatrixMultiply(VPIPerspectiveTransform &r, const VPIPerspectiveTransform &a,
                           const VPIPerspectiveTransform &b)
{
    for (int i = 0; i < 3; ++i)
    {
        for (int j = 0; j < 3; ++j)
        {
            r[i][j] = a[i][0] * b[0][j];
            for (int k = 1; k < 3; ++k)
            {
                r[i][j] += a[i][k] * b[k][j];
            }
        }
    }
}

// Utility function to wrap a cv::Mat into a VPIImage
static VPIImage ToVPIImage(const cv::Mat &frame)
{
    VPIImageData imgData;
    memset(&imgData, 0, sizeof(imgData));

    switch (frame.type())
    {
    case CV_8U:
        imgData.type = VPI_IMAGE_TYPE_U8;
        break;
    case CV_8UC3:
        imgData.type = VPI_IMAGE_TYPE_BGR8;
        break;
    case CV_8UC4:
        imgData.type = VPI_IMAGE_TYPE_BGRA8;
        break;
    default:
        throw std::runtime_error("Frame type not supported");
    }

    // First fill VPIImageData with the, well, image data...
    imgData.numPlanes           = 1;
    imgData.planes[0].width     = frame.cols;
    imgData.planes[0].height    = frame.rows;
    imgData.planes[0].rowStride = frame.step[0];
    imgData.planes[0].data      = frame.data;

    // Now create a VPIImage that wraps it.
    VPIImage img;
    CHECK_STATUS(vpiImageWrapHostMem(&imgData, 0, &img));
    return img;
};

// Utility function to wrap a VPIImageData into a cv::Mat
static cv::Mat ToCV(const VPIImageData &imgData)
{
    cv::Mat out;

    switch (imgData.type)
    {
    case VPI_IMAGE_TYPE_BGR8:
        out = cv::Mat(imgData.planes[0].height, imgData.planes[0].width, CV_8UC3, imgData.planes[0].data,
                      imgData.planes[0].rowStride);
        break;

    case VPI_IMAGE_TYPE_BGRA8:
        out = cv::Mat(imgData.planes[0].height, imgData.planes[0].width, CV_8UC4, imgData.planes[0].data,
                      imgData.planes[0].rowStride);
        break;

    case VPI_IMAGE_TYPE_U8:
        out = cv::Mat(imgData.planes[0].height, imgData.planes[0].width, CV_8UC1, imgData.planes[0].data,
                      imgData.planes[0].rowStride);
        break;

    default:
        throw std::runtime_error("Frame type not supported");
    }

    return out;
}

int main(int argc, char *argv[])
{
    // We'll create all our objects under this context, so that
    // we don't have to track what objects to destroy. Just destroying
    // the context will destroy all objects.
    VPIContext ctx = nullptr;

    int retval = 0;

    try
    {
        if (argc != 4)
        {
            throw std::runtime_error(std::string("Usage: ") + argv[0] + " <cpu|pva|cuda> <input_video> <output>");
        }

        std::string strDevType     = argv[1];
        std::string strInputVideo  = argv[2];
        std::string strOutputVideo = argv[3];

        // Load the input video
        cv::VideoCapture invid;
        if (!invid.open(strInputVideo))
        {
            throw std::runtime_error("Can't open '" + strInputVideo + "'");
        }

        // Create our context.
        CHECK_STATUS(vpiContextCreate(0, &ctx));

        // Activate it. From now on all created objects will be owned by it.
        CHECK_STATUS(vpiContextSetCurrent(ctx));

        // Now process the device type
        VPIDeviceType devType;

        if (strDevType == "cpu")
        {
            devType = VPI_DEVICE_TYPE_CPU;
        }
        else if (strDevType == "cuda")
        {
            devType = VPI_DEVICE_TYPE_CUDA;
        }
        else if (strDevType == "pva")
        {
            devType = VPI_DEVICE_TYPE_PVA;
        }
        else
        {
            throw std::runtime_error("Backend '" + strDevType +
                                     "' not recognized, it must be either cpu, cuda or pva.");
        }

        // Create the stream for the given backend.
        VPIStream streamWarp;
        CHECK_STATUS(vpiStreamCreate(devType, &streamWarp));

        VPIStream streamConv;
        if (devType == VPI_DEVICE_TYPE_PVA)
        {
            // PVA backend doesn't have currently Image Format Converter algorithm. We'll use CUDA
            // backend to do that.
            CHECK_STATUS(vpiStreamCreate(VPI_DEVICE_TYPE_CUDA, &streamConv));
        }
        else
        {
            // On other cases, we can use the same stream to do format conversion.
            streamConv = streamWarp;
        }

        int w = invid.get(cv::CAP_PROP_FRAME_WIDTH);
        int h = invid.get(cv::CAP_PROP_FRAME_HEIGHT);

        cv::VideoWriter outVideo(strOutputVideo, invid.get(cv::CAP_PROP_FOURCC), invid.get(cv::CAP_PROP_FPS),
                                 cv::Size(w, h));

        VPIImage imgInput, imgOutput;
        CHECK_STATUS(vpiImageCreate(w, h, VPI_IMAGE_TYPE_NV12, 0, &imgInput));
        CHECK_STATUS(vpiImageCreate(w, h, VPI_IMAGE_TYPE_NV12, 0, &imgOutput));

        // Create a Perspective Image Warp payload.
        VPIPayload warp;
        CHECK_STATUS(vpiCreatePerspectiveImageWarp(streamWarp, &warp));

        // Create events for stream synchronization (if needed)
        VPIEvent ev1, ev2;
        CHECK_STATUS(vpiEventCreate(0, &ev1));
        CHECK_STATUS(vpiEventCreate(0, &ev2));

        VPIPerspectiveTransform xform;
        memset(&xform, 0, sizeof(xform));

        int curFrame = 1;
        cv::Mat cvFrame;
        while (invid.read(cvFrame))
        {

            VPIImage frameBGR = ToVPIImage(cvFrame);

            // First convert it to NV12
            CHECK_STATUS(vpiSubmitImageFormatConverter(streamConv, frameBGR, imgInput, VPI_CONVERSION_CAST, 1, 0));

            // If format convertion is using a different stream than warping,
            if (streamConv != streamWarp)
            {
                // make sure warping starts only after format conversion is finished.
                CHECK_STATUS(vpiEventRecord(ev1, streamConv));
                CHECK_STATUS(vpiStreamWaitFor(streamWarp, ev1));
            }

            // move image's center to origin of coordinate system
            VPIPerspectiveTransform t1 = {{1, 0, -w / 2.0f}, {0, 1, -h / 2.0f}, {0, 0, 1}};

            // Apply some time-dependent perspective transform
            float v1                  = sin(curFrame / 30.0 * 2 * M_PI / 2) * 0.0005f;
            float v2                  = cos(curFrame / 30.0 * 2 * M_PI / 3) * 0.0005f;
            VPIPerspectiveTransform P = {{0.66, 0, 0}, {0, 0.66, 0}, {v1, v2, 1}};

            // move image's center back to where it was.
            VPIPerspectiveTransform t2 = {{1, 0, w / 2.0f}, {0, 1, h / 2.0f}, {0, 0, 1}};

            // Apply the transforms defined above.
            VPIPerspectiveTransform tmp;
            MatrixMultiply(tmp, P, t1);
            MatrixMultiply(xform, t2, tmp);

            // Do perspective warp
            CHECK_STATUS(vpiSubmitPerspectiveImageWarp(warp, imgInput, xform, imgOutput, VPI_INTERP_LINEAR,
                                                       VPI_BOUNDARY_COND_ZERO, 0));

            if (streamConv != streamWarp)
            {
                // make sure format conversion starts only after warping is finished.
                CHECK_STATUS(vpiEventRecord(ev2, streamWarp));
                CHECK_STATUS(vpiStreamWaitFor(streamConv, ev2));
            }
            // Convert output back to BGR
            CHECK_STATUS(vpiSubmitImageFormatConverter(streamConv, imgOutput, frameBGR, VPI_CONVERSION_CAST, 1, 0));
            CHECK_STATUS(vpiStreamSync(streamConv));

            // Now add it to the output video stream
            VPIImageData imgdata;
            CHECK_STATUS(vpiImageLock(frameBGR, VPI_LOCK_READ, &imgdata));
            outVideo << ToCV(imgdata);
            CHECK_STATUS(vpiImageUnlock(frameBGR));

        float elapsedConvert, elapsedDewarp, elapsedTotal;
        CHECK_STATUS(vpiEventElapsedTime(ev1, ev2, &elapsedDewarp));

            // Not needed anymore
            vpiImageDestroy(frameBGR);
            printf("Frame: %d, warp-time: %f ms\n", curFrame++, elapsedDewarp);
        }
    }
    catch (std::exception &e)
    {
        std::cerr << e.what() << std::endl;
        retval = 1;
    }

    // Clean up
    vpiContextDestroy(ctx);

    return retval;
}