/*
 * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
 *
 * NVIDIA CORPORATION and its licensors retain all intellectual property
 * and proprietary rights in and to this software, related documentation
 * and any modifications thereto.  Any use, reproduction, disclosure or
 * distribution of this software and related documentation without an express
 * license agreement from NVIDIA CORPORATION is strictly prohibited.
*/

/* standard headers */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <stdbool.h>
#include <signal.h>

/* Nvidia headers */
#include "cmdline.h"
#include "log_utils.h"
#include "nvmedia_image.h"
#include "nvmedia_ldc.h"
#include "nvmedia_surface.h"
#include "image_utils.h"

static NvMediaBool quit_flag = NVMEDIA_FALSE;
static void
SigHandler (int signum) {
    signal(SIGINT, SIG_IGN);
    signal(SIGTERM, SIG_IGN);
    signal(SIGQUIT, SIG_IGN);
    signal(SIGHUP, SIG_IGN);
    signal(SIGKILL, SIG_IGN);
    signal(SIGSTOP, SIG_IGN);

    quit_flag = NVMEDIA_TRUE;

    signal(SIGINT, SIG_DFL);
    signal(SIGTERM, SIG_DFL);
    signal(SIGQUIT, SIG_DFL);
    signal(SIGHUP, SIG_DFL);
    signal(SIGKILL, SIG_DFL);
    signal(SIGSTOP, SIG_DFL);
}

static void
SigSetup (void) {
    struct sigaction action;
    memset(&action, 0, sizeof(action));
    action.sa_handler = SigHandler;
    sigaction(SIGINT, &action, NULL);
    sigaction(SIGTERM, &action, NULL);
    sigaction(SIGQUIT, &action, NULL);
    sigaction(SIGHUP, &action, NULL);
    sigaction(SIGKILL, &action, NULL);
    sigaction(SIGSTOP, &action, NULL);
}

/* Timeout for completion of LDC operations */
/* Set to infinite for now.
 * TODO: Set to appropriate value */
#define IMAGE_COMPLETION_WAIT_TIME_MS NVMEDIA_IMAGE_TIMEOUT_INFINITE

/* Checks NvMediaStatus and jumps to fail */
#define CHK_STATUS_AND_RETURN(function, status)                     \
    if (status != NVMEDIA_STATUS_OK) {                              \
        LOG_ERR("%s failed! status = %x\n", function, status);      \
        ret = -1;                                                   \
        goto ldc_cleanup;                                           \
    }

/* Checks ret and jumps to fail */
#define CHK_AND_RETURN(function)                                    \
    if (ret != 0) {                                                 \
        LOG_ERR("%s failed!\n", function);                          \
        goto ldc_cleanup;                                           \
    }

/* Global variables */

/* gRefTnr3Params contains the typical/recommended settings for different lighting conditions */
NvMediaTNR3Params gRefTnr3Params[] = {
    {
    /* Bright */
/* spatialSigmaLuma */   2,
/* spatialSigmaChroma */ 2,
/* rangeSigmaLuma */     15,
/* rangeSigmaChroma */   45,
/* sadMultiplier */      1.0f,
/* sadWeightLuma */      0.5f,
/* alphaSmoothEnable */  0,
/* alphaIncreaseCap */   0.125f,
/* alphaScaleIIR */      0.5f,
/* alphaMaxLuma */       0.9277f,
/* alphaMinLuma */       0.0,
/* alphaMaxChroma */     0.9277f,
/* alphaMinChroma */     0.0,
/* betaX1 */             0.3906f,
/* betaX2 */             1.0f,
/* minBeta */            0.3906f,
/* maxBeta */            1.0f
    },

    {
    /* Medium */
/* spatialSigmaLuma */   3,
/* spatialSigmaChroma */ 3,
/* rangeSigmaLuma */     20,
/* rangeSigmaChroma */   60,
/* sadMultiplier */      0.7143f,
/* sadWeightLuma */      0.5f,
/* alphaSmoothEnable */  1,
/* alphaIncreaseCap */   0.125f,
/* alphaScaleIIR */      0.5,
/* alphaMaxLuma */       0.8789f,
/* alphaMinLuma */       0.0,
/* alphaMaxChroma */     0.8789f,
/* alphaMinChroma */     0.1953,
/* betaX1 */             0.3906f,
/* betaX2 */             1.0f,
/* minBeta */            0.3906f,
/* maxBeta */            1.0f
    },
    {
    /* Low */
/* spatialSigmaLuma */   4,
/* spatialSigmaChroma */ 4,
/* rangeSigmaLuma */     40,
/* rangeSigmaChroma */   120,
/* sadMultiplier */      0.5f,
/* sadWeightLuma */      0.5f,
/* alphaSmoothEnable */  1,
/* alphaIncreaseCap */   0.125f,
/* alphaScaleIIR */      0.5,
/* alphaMaxLuma */       0.7324f,
/* alphaMinLuma */       0.0,
/* alphaMaxChroma */     0.7324f,
/* alphaMinChroma */     0.1953f,
/* betaX1 */             0.3906f,
/* betaX2 */             1.0f,
/* minBeta */            0.3906f,
/* maxBeta */            1.0f
    },
    {
    /* Very low */
/* spatialSigmaLuma */   6,
/* spatialSigmaChroma */ 6,
/* rangeSigmaLuma */     200,
/* rangeSigmaChroma */   200,
/* sadMultiplier */      0.3571f,
/* sadWeightLuma */      0.5f,
/* alphaSmoothEnable */  1,
/* alphaIncreaseCap */   0.125f,
/* alphaScaleIIR */      0.5,
/* alphaMaxLuma */       0.7324f,
/* alphaMinLuma */       0.0,
/* alphaMaxChroma */     0.7324f,
/* alphaMinChroma */     0.1953f,
/* betaX1 */             0.3906f,
/* betaX2 */             1.0f,
/* minBeta */            0.3906f,
/* maxBeta */            1.0f
    }
};

/* Reads Sparse Warp map described via config file and warp map file.
 * Returns
 *  0 => Success
 * -1 => Failure
 */
static int
sReadSparseWarpMap(
    NvMediaLDCSparseWarpMap *pWarpMap, /* LDC Sparse Warp map */
    char *SparseWarpMapFile)           /* Sparse Warp map file */
{
    size_t size = 0;
    FILE *f = NULL;

    if (!pWarpMap) {
        LOG_ERR("Invalid Sparse Warp map\n");
        return -1;
    }

    /* Read the sparse warp map file */
    size = (size_t) pWarpMap->mapStride * (size_t) pWarpMap->numVerPoints * sizeof(float_t) * 2;
    if (!size) {
        LOG_ERR("Invalid mapStride/numVerPoints\n");
        return -1;
    }

    pWarpMap->mapPtr = calloc(1, size);
    if (!pWarpMap->mapPtr) {
        LOG_ERR("Failed to allocate WarpMap mapPtr\n");
        return -1;
    }

    f = fopen(SparseWarpMapFile, "rb");
    if (!f) {
        LOG_ERR("Failed to open SparseWarpMapFile:%s\n", SparseWarpMapFile);
        return -1;
    }

    if (fread(pWarpMap->mapPtr, 1, size, f) != size) {
        LOG_ERR("Failed to read SparseWarpMapFile:%s\n", SparseWarpMapFile);
        fclose(f);
        return -1;
    }

    fclose(f);

    return 0;
}

/* Reads bit mask map from bit mask file.
 * Returns,
 *  0 => Success
 * -1 => Failure
 */
static int
sReadBitMaskMap(
        NvMediaLDCBitMaskMap *pBitMaskMap, /* NvMedia LDC Bit Mask Map */
        char *bitMaskFile)                 /* Bit Mask map file */
{
    /* Read bit map file */
    size_t size = (size_t) (pBitMaskMap->width) * (size_t) (pBitMaskMap->height);
    FILE *f = NULL;

    if (size) {
        pBitMaskMap->mapPtr = calloc(1, size);
        if (!pBitMaskMap->mapPtr) {
            LOG_ERR("Failed to allocate BitMaskMap mapPtr\n");
            return -1;
        }
    }

    f = fopen(bitMaskFile, "rb");
    if (!f) {
        LOG_ERR("Failed to open bitMaskFile:%s\n", bitMaskFile);
        return -1;
    }

    if (fread(pBitMaskMap->mapPtr, 1, size, f) != size) {
        LOG_ERR("Failed to read bitMaskFile:%s\n", bitMaskFile);
        fclose(f);
        return -1;
    }

    fclose(f);

    return 0;
}

/* Gets NvMedia surface type from input parameters
 * order, subsampling, memType and bpc.
 * Returns NvMediaSurfaceType.
 */
static NvMediaSurfaceType
sGetNvMediaYUVSurfaceType(
        uint32_t order,        /* order */
        uint32_t subSampling,  /* Sub-sampling type */
        uint32_t memType,      /* Memory type */
        uint32_t bpc)          /* Bits per component */
{
    NvMediaSurfaceType type;

    /* define NvMedia Surface attributes */
    NVM_SURF_FMT_DEFINE_ATTR(surfFormatAttrs);

    /* Set surface attributes */
    surfFormatAttrs[0].type = NVM_SURF_ATTR_SURF_TYPE;
    surfFormatAttrs[0].value = NVM_SURF_ATTR_SURF_TYPE_YUV;

    surfFormatAttrs[1].type = NVM_SURF_ATTR_LAYOUT;
    surfFormatAttrs[1].value = NVM_SURF_ATTR_LAYOUT_PL;

    surfFormatAttrs[2].type = NVM_SURF_ATTR_DATA_TYPE;
    surfFormatAttrs[2].value = NVM_SURF_ATTR_DATA_TYPE_UINT;

    surfFormatAttrs[3].type = NVM_SURF_ATTR_MEMORY;
    surfFormatAttrs[3].value = memType;

    surfFormatAttrs[4].type = NVM_SURF_ATTR_SUB_SAMPLING_TYPE;
    surfFormatAttrs[4].value = subSampling;

    surfFormatAttrs[5].type = NVM_SURF_ATTR_BITS_PER_COMPONENT;
    surfFormatAttrs[5].value = bpc;

    surfFormatAttrs[6].type = NVM_SURF_ATTR_COMPONENT_ORDER;
    surfFormatAttrs[6].value = order;

    type = NvMediaSurfaceFormatGetType(surfFormatAttrs, NVM_SURF_FMT_ATTR_MAX);
    return type;
}


/* Create an NvMediaImage based on the input attributes.
 * Returns,
 * NULL => Failed to create Image
 * non-NULL ptr => Pointer to NvMediaImage
 */
static int
sCreateSurface(
        NvMediaImage **pImage, /* Pointer to NvMediaImage handle */
        NvMediaDevice *device, /* NvMediaDevice */
        uint32_t width,        /* Width of the image */
        uint32_t height,       /* Height of the image */
        uint32_t order,        /* order */
        uint32_t subSampling,  /* Sub-sampling type */
        uint32_t memType,      /* Memory type */
        uint32_t bpc)          /* Bits per component */

{
    NvMediaSurfaceType type;
    NvMediaSurfAllocAttr surfAllocAttrs[NVM_SURF_ALLOC_ATTR_MAX];

    *pImage = NULL;

    type = sGetNvMediaYUVSurfaceType(order, subSampling, memType, bpc);
    if (type == NvMediaSurfaceType_Unsupported) {
        LOG_ERR("Unsupported surface format attributes\n");
        return -1;
    }

    surfAllocAttrs[0].type = NVM_SURF_ATTR_WIDTH;
    surfAllocAttrs[0].value = width;
    surfAllocAttrs[1].type = NVM_SURF_ATTR_HEIGHT;
    surfAllocAttrs[1].value = height;

    *pImage = NvMediaImageCreateNew(device,         /* device */
                                    type,           /* surface type */
                                    surfAllocAttrs, /* surf allocation attrs */
                                    2,              /* num attrs */
                                    0);             /* flags */
    if(!(*pImage)) {
        LOG_ERR("Unable to create NvMediaImage\n");
        return -1;
    }

    return 0;
}


/* Main LDC test function.
 * Creates NvMediaLDC, configures NvMediaLDC based on config from config file,
 * creates NvMediaImage(s), Reads images from yuv files in to NvMediaImage(s) via util functions and
 * calls NvMediaLDCProcess() and writes the output NvMediaImage(s) to file(s) using utils functions.
 * There are no sanity checks done on the arguments passed to NvMediaLDC* functions as they are
 * expected to be done at library and failed gracefully.
 */
int main(int argc, char *argv[])
{
    TestArgs args;
    NvMediaDevice *device = NULL;
    NvMediaImage *curr = NULL, *prev = NULL, *temp = NULL;
    NvMediaImage *out = NULL, *xSob = NULL, *dSample = NULL;

    NvMediaLDC *pLDC = NULL; /* NvMediaLDC handle */

    /* NvMediaLDC config parameters. To be read via config file */
    NvMediaLDCInitParams *pInitParams = NULL;
    NvMediaGeoTransParams *pGeoTrans = NULL;
    NvMediaLDCMode ldcMode = NVMEDIA_LDC_MODE_GEOTRANS;
    uint32_t srcWidth, srcHeight;
    uint32_t dstWidth, dstHeight;

    uint32_t ret = 0, frameIdx = 0;

    NvMediaStatus status = NVMEDIA_STATUS_OK;
    NvMediaTaskStatus taskStatus;

    uint64_t totalOutTimeTaken = 0;

    NvMediaVersion version;

    SigSetup();
    /* Read configuration from command line and config file */
    memset(&args, 0, sizeof(TestArgs));

    /* ParseArgs parses the command line and the LDC configuration file and populates all initParams
     * and run time configuration in to appropriate structures within args
     */
    if (ParseArgs(argc, argv, &args)) {
        PrintUsage();
        return -1;
    }

    /* Get copies of or reference to config populated by ParseArgs */
    pInitParams = &args.initParams;
    ldcMode = pInitParams->ldcMode;
    srcWidth = args.srcWidth;
    srcHeight = args.srcHeight;
    dstWidth = args.dstWidth;
    dstHeight = args.dstHeight;
    pGeoTrans = &(pInitParams->geoTransParams);

    /* Read bit mask map file, if present, when GEO_TRANS is enabled */
    if (ldcMode == NVMEDIA_LDC_MODE_GEOTRANS_TNR3 ||
        ldcMode == NVMEDIA_LDC_MODE_GEOTRANS) {
        if (pGeoTrans->bitMaskEnable) {
            ret = sReadBitMaskMap(&pGeoTrans->bitMaskMap,
                                  args.bitMaskFile);
            CHK_AND_RETURN("sReadBitMaskMap");
        }
    }

    /* Check version */
    status = NvMediaLDCGetVersion(&version);
    if (status == NVMEDIA_STATUS_OK) {
        LOG_INFO("Library version: %u.%u\n", version.major, version.minor);
        LOG_INFO("Header version:  %u.%u\n", NVMEDIA_LDC_VERSION_MAJOR, NVMEDIA_LDC_VERSION_MINOR);
        if ((version.major != NVMEDIA_LDC_VERSION_MAJOR) ||
            (version.minor != NVMEDIA_LDC_VERSION_MINOR)) {
            LOG_ERR("Library and Header mismatch!\n");
        }
    }

    pInitParams->srcSurfaceType = sGetNvMediaYUVSurfaceType(NVM_SURF_ATTR_COMPONENT_ORDER_YUV,
                                                        args.subSampling,
                                                        NVM_SURF_ATTR_MEMORY_SEMI_PLANAR,
                                                        args.bitsPerComponent);

    /* Create NvMedia device.
     * This is required to create the NvMediaImage surfaces
     * and to create NvMedia LDC handle
     */
    device = NvMediaDeviceCreate();
    if(!device) {
        LOG_ERR("NvMediaDeviceCreate failed\n");
        ret = -1;
        goto ldc_cleanup;
    }

    /* Create LDC Handle */
    status = NvMediaLDCCreateNew(device,
                                 &pLDC,
                                 srcWidth,
                                 srcHeight,
                                 &args.srcRect,
                                 dstWidth,
                                 dstHeight,
                                 &args.dstRect,
                                 pInitParams);

    CHK_STATUS_AND_RETURN("NvMediaLDCCreate", status);
    LOG_INFO("NvMediaLDCCreate done. LDC Handle:%p\n", pLDC);

    /* Generate/Feed the Sparse warp mapping when GEO_TRANS is enabled and geoTransMode is GEN/FEED
     * mapping */
    if (ldcMode == NVMEDIA_LDC_MODE_GEOTRANS_TNR3 ||
        ldcMode == NVMEDIA_LDC_MODE_GEOTRANS) {
        if (pGeoTrans->geoTransMode == NVMEDIA_GEOTRANS_MODE_GEN_MAPPING) {
            /* NvMediaLDC will generate the mapping */
            status = NvMediaLDCMappingGen(pLDC);
            CHK_STATUS_AND_RETURN("NvMediaLDCMappingGen", status);
            LOG_INFO("NvMediaLDCMappingGen(%p) done.\n", pLDC);
        }
        else if (pGeoTrans->geoTransMode == NVMEDIA_GEOTRANS_MODE_FEED_MAPPING) {
            /* Read the Sparse Warp Map from file */
            ret = sReadSparseWarpMap(&args.SparseWarpMap,
                                     args.SparseWarpMapFile);
            CHK_AND_RETURN("sReadSparseWarpMap");

            /* Feed the Sparse Warp Map to NvMediaLDC */
            status = NvMediaLDCFeedSparseWarpMap(pLDC,
                                                 &args.SparseWarpMap);
            CHK_STATUS_AND_RETURN("NvMediaLDCFeedSparseWarpMap", status);
            LOG_INFO("NvMediaLDCFeedSparseWarpMap(%s) done.\n", args.SparseWarpMapFile);
        }
    }

    /* Create NvMedia input and output surfaces */
    /* curr. Current input surface */
    ret = sCreateSurface(&curr,
                         device,
                         srcWidth,
                         srcHeight,
                         NVM_SURF_ATTR_COMPONENT_ORDER_YUV,
                         args.subSampling,
                         NVM_SURF_ATTR_MEMORY_SEMI_PLANAR,
                         args.bitsPerComponent);
    CHK_AND_RETURN("sCreateSurface(curr)");
    LOG_INFO("Current inptut surface %p created\n", curr);

    /* prev. Required when TNR3/TNR2 is enabled */
    if ((ldcMode == NVMEDIA_LDC_MODE_GEOTRANS_TNR3 ||
         ldcMode == NVMEDIA_LDC_MODE_TNR3 ||
         ldcMode == NVMEDIA_LDC_MODE_TNR2)         &&
        (args.numFrames > 1)) {
        ret = sCreateSurface(&prev,
                             device,
                             dstWidth,
                             dstHeight,
                             NVM_SURF_ATTR_COMPONENT_ORDER_YUV,
                             args.subSampling,
                             NVM_SURF_ATTR_MEMORY_SEMI_PLANAR,
                             args.bitsPerComponent);
        CHK_AND_RETURN("sCreateSurface(prev)");
        LOG_INFO("Previous output surface %p created\n", prev);
    }

    /* out. Output surface */
    ret = sCreateSurface(&out,
                         device,
                         dstWidth,
                         dstHeight,
                         NVM_SURF_ATTR_COMPONENT_ORDER_YUV,
                         args.subSampling,
                         NVM_SURF_ATTR_MEMORY_SEMI_PLANAR,
                         args.bitsPerComponent);
    CHK_AND_RETURN("sCreateSurface(out)");
    LOG_INFO("Output surface %p created\n", out);
    /* Start main processing loop */
    for (frameIdx = 0; ((frameIdx < args.numFrames) && (quit_flag != NVMEDIA_TRUE)); frameIdx++) {
        bool bNeedXob = false;
        bool bNeedDS = false;
        /* Get LDC control parameters read via config file.
         * These serve as initial config and can be modified within the processing loop
         */
        NvMediaLDCCtrlParams *ldcCtrlParams = &(args.ldcCtrlParams);

        LOG_DBG("Processing frame:%u\n", frameIdx);

        /* Read input file */
        status = ReadImage(args.infile,
                           frameIdx, /* frame number */
                           srcWidth,
                           srcHeight,
                           curr,
                           args.UVorder, /* inputUVOrderFlag */
                           0,  /* rawBytesPerPixel for RAW */
                           MSB_ALIGNED); /* pixelAlignment */
        CHK_STATUS_AND_RETURN("ReadImage", status);
        LOG_INFO("ReadImage(frame:%u from %s in to curr:%p) done\n", frameIdx, args.infile, curr);

        /* Check xSobelMode mode and create xSob and downSample surfaces if necessary */
        if (ldcMode == NVMEDIA_LDC_MODE_GEOTRANS) {
            switch (ldcCtrlParams->xSobelMode) {
                case NVMEDIA_GEOTRANS_DISABLE_XSOBEL_ENABLE_DS:
                case NVMEDIA_GEOTRANS_ENABLE_XSOBEL_ENABLE_DS:
                    bNeedXob = true;
                    bNeedDS = true;
                    break;
                case NVMEDIA_GEOTRANS_ENABLE_XSOBEL_DISABLE_DS:
                    bNeedXob = true;
                    bNeedDS = false;
                    break;
                case NVMEDIA_GEOTRANS_DISABLE_XSOBEL_DISABLE_DS:
                default:
                    break;
            }

            if (bNeedXob && !xSob) {
                /* Create Xob Surface */
                ret = sCreateSurface(&xSob,
                                     device,
                                     dstWidth,
                                     dstHeight,
                                     NVM_SURF_ATTR_COMPONENT_ORDER_LUMA,
                                     NVM_SURF_ATTR_SUB_SAMPLING_TYPE_NONE,
                                     NVM_SURF_ATTR_MEMORY_PACKED,
                                     args.bitsPerComponent);
                CHK_AND_RETURN("sCreateSurface(xSob)");
                LOG_INFO("xSobel output surface %p created\n", xSob);
            }

            if (bNeedDS && !dSample) {
                /* Create Down Sample surface */
                ret = sCreateSurface(&dSample,
                                     device,
                                     dstWidth >> 2, /* DS is 4x4 downsampled xSob */
                                     dstHeight >> 2, /* DS is 4x4 downsampled xSob */
                                     NVM_SURF_ATTR_COMPONENT_ORDER_LUMA,
                                     NVM_SURF_ATTR_SUB_SAMPLING_TYPE_NONE,
                                     NVM_SURF_ATTR_MEMORY_PACKED,
                                     args.bitsPerComponent);
                CHK_AND_RETURN("sCreateSurface(dSample)");
                LOG_INFO("downsampled xSobel output surface %p created\n", dSample);
            }
        }

        /* Update TNR3 Parameters if necessary */
        if (args.bUpdateTNR3Params) {
            if (ldcMode == NVMEDIA_LDC_MODE_TNR3 ||
                ldcMode == NVMEDIA_LDC_MODE_GEOTRANS_TNR3) {
                /* Update TNR3 Parameters
                 * The below code rotates through a pre-determined reference settings for different
                 * lighting conditions to demo the usage of API.
                 * */
                NvMediaTNR3Params *tnr3Params = NULL;
                uint32_t numRef = sizeof(gRefTnr3Params)/sizeof(NvMediaTNR3Params);
                tnr3Params = &(gRefTnr3Params[frameIdx%numRef]);
                status = NvMediaLDCUpdateTNR3Params(pLDC,
                                                    tnr3Params);
                CHK_STATUS_AND_RETURN("NvMediaLDCUpdateTNR3Params", status);
                LOG_INFO("NvMediaLDCUpdateTNR3Params(%u) done\n", frameIdx%numRef);
            }
        }

        /* NvMediaLDCProcess */
        status = NvMediaLDCProcess(pLDC,
                                   frameIdx ? prev : NULL, /* prevSurface */
                                   curr, /* curSurface */
                                   out, /* outputSurface */
                                   bNeedXob ? xSob : NULL, /* xSobel */
                                   bNeedDS ? dSample : NULL, /* downSample */
                                   ldcCtrlParams);
        CHK_STATUS_AND_RETURN("NvMediaLDCProcess", status);
        LOG_INFO("NvMediaLDCProcess(pLDC:%p) done\n", pLDC);

        /* Dump output(s) to file */
        if(out) {
            uint64_t timeEnd, timeStart = 0, timeTaken = 0;
            LOG_DBG("Waiting for out frame:%u\n", frameIdx);

            GetTimeMicroSec(&timeStart);
            /* Wait for HW to be done with the surface */
            status = NvMediaImageGetStatus(out, IMAGE_COMPLETION_WAIT_TIME_MS, &taskStatus);
            CHK_STATUS_AND_RETURN("NvMediaImageGetStatus(out)", status);
            GetTimeMicroSec(&timeEnd);
            timeTaken = (timeEnd - timeStart);
            totalOutTimeTaken += timeTaken;
            LOG_INFO("NvMediaImageGetStatus(out:%p) done in %llu us\n\n", out, timeTaken);

            status = WriteImage(args.outfile,
                                out,
                                args.UVorder,
                                frameIdx ? NVMEDIA_TRUE : NVMEDIA_FALSE, /* Append flag */
                                0, /* rawBytesPerPixel for RAW */
                                NULL);
            CHK_STATUS_AND_RETURN("WriteImage(out)", status);
            LOG_INFO("WriteImage(from out:%p in to %s) done\n", out, args.outfile);
        }
        if(xSob) {
            uint64_t timeEnd, timeStart = 0, timeTaken = 0;
            LOG_DBG("Waiting for xSob frame:%u\n", frameIdx);

            GetTimeMicroSec(&timeStart);
            /* Wait for HW to be done with the surface */
            status = NvMediaImageGetStatus(xSob, IMAGE_COMPLETION_WAIT_TIME_MS, &taskStatus);
            CHK_STATUS_AND_RETURN("NvMediaImageGetStatus(xSob)", status);
            GetTimeMicroSec(&timeEnd);
            timeTaken = (timeEnd - timeStart);
            LOG_INFO("NvMediaImageGetStatus(xSob:%p) done in %llu us\n", xSob, timeTaken);

            status = WriteImage(args.xsobelfile,
                                xSob,
                                args.UVorder,
                                frameIdx ? NVMEDIA_TRUE : NVMEDIA_FALSE, /* Append flag */
                                0, /* rawBytesPerPixel for RAW */
                                NULL);
            CHK_STATUS_AND_RETURN("WriteImage(out)", status);
            LOG_INFO("WriteImage(from xSob:%p in to %s) done\n", xSob, args.xsobelfile);
        }
        if(dSample) {
            uint64_t timeEnd, timeStart = 0, timeTaken = 0;
            LOG_DBG("Waiting for dSample frame:%u\n", frameIdx);

            GetTimeMicroSec(&timeStart);
            /* Wait for HW to be done with the surface */
            status = NvMediaImageGetStatus(dSample, IMAGE_COMPLETION_WAIT_TIME_MS, &taskStatus);
            CHK_STATUS_AND_RETURN("NvMediaImageGetStatus(dSample)", status);
            GetTimeMicroSec(&timeEnd);
            timeTaken = (timeEnd - timeStart);
            LOG_INFO("NvMediaImageGetStatus(dSample:%p) done in %llu us\n", dSample, timeTaken);

            status = WriteImage(args.downsamplefile,
                                dSample,
                                args.UVorder,
                                frameIdx ? NVMEDIA_TRUE : NVMEDIA_FALSE, /* Append flag */
                                0, /* rawBytesPerPixel for RAW */
                                NULL);
            CHK_STATUS_AND_RETURN("WriteImage(dSample)", status);
            LOG_INFO("WriteImage(from dSample:%p in to %s) done\n", dSample, args.downsamplefile);
        }

        /* Swap out and prev */
        if (ldcMode == NVMEDIA_LDC_MODE_GEOTRANS_TNR3 ||
            ldcMode == NVMEDIA_LDC_MODE_TNR3 ||
            ldcMode == NVMEDIA_LDC_MODE_TNR2) {

            LOG_DBG("swap out:%p and prev:%p \n", out, prev);
            temp = prev;
            prev = out;
            out = temp;
        }
    }

    if (frameIdx) {
        LOG_INFO("\nAverage time taken to process a frame: %llu us\n", totalOutTimeTaken/frameIdx);
    }

    /* Cleanup before exit */
ldc_cleanup:
    LOG_INFO("Processed %u frames, cleaning up\n", frameIdx);
    if(pLDC) {
        status = NvMediaLDCDestroy(pLDC);
        if (status != NVMEDIA_STATUS_OK) {
            LOG_ERR("NvMediaLDCDestroy failed!\n");
        }
        LOG_INFO("NvMediaLDCDestroy(pLDC:%p) done\n", pLDC);
    }

    /* Free any memory allocted */
    if (args.SparseWarpMap.mapPtr) {
        free(args.SparseWarpMap.mapPtr);
    }

    /* Destroy NvMediaImage surfaces */
    if (curr) {
        LOG_DBG("Destroying input surface\n");
        NvMediaImageDestroy(curr);
    }

    if (out) {
        LOG_DBG("Destroying output surface\n");
        NvMediaImageDestroy(out);
    }

    if (prev) {
        LOG_DBG("Destroying previous surface\n");
        NvMediaImageDestroy(prev);
    }

    if (xSob) {
        LOG_DBG("Destroying xSobel output surface\n");
        NvMediaImageDestroy(xSob);
    }

    if (dSample) {
        LOG_DBG("Destroying downsampled xSobel output surface\n");
        NvMediaImageDestroy(dSample);
    }

    /* Destroy NvMediaDevice */
    if(device) {
        LOG_DBG("Destroying device\n");
        NvMediaDeviceDestroy(device);
    }

    return ret;
}
