/* Copyright (c) 2022, 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 "image_utils.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "misc_utils.h"
#include "nvmedia_surface.h"

#define MAXM_NUM_SURFACES 6

typedef struct
{
    float        heightFactor[6];
    float        widthFactor[6];
    unsigned int numSurfaces;
} ImgUtilSurfParams;

ImgUtilSurfParams ImgSurfParamsTable_RGBA = {
    .heightFactor = {1, 0, 0, 0, 0, 0},
    .widthFactor  = {1, 0, 0, 0, 0, 0},
    .numSurfaces  = 1,
};

ImgUtilSurfParams ImgSurfParamsTable_RAW = {
    .heightFactor = {1, 0, 0, 0, 0, 0},
    .widthFactor  = {1, 0, 0, 0, 0, 0},
    .numSurfaces  = 1,
};

ImgUtilSurfParams ImgSurfParamsTable_YUV[][4] = {
    {
        /* PLANAR */
        {
            /* 420 */
            .heightFactor = {1, 0.5, 0.5, 0, 0, 0},
            .widthFactor  = {1, 0.5, 0.5, 0, 0, 0},
            .numSurfaces  = 3,
        },
        {
            /* 422 */
            .heightFactor = {1, 1, 1, 0, 0, 0},
            .widthFactor  = {1, 0.5, 0.5, 0, 0, 0},
            .numSurfaces  = 3,
        },
        {
            /* 444 */
            .heightFactor = {1, 1, 1, 0, 0, 0},
            .widthFactor  = {1, 1, 1, 0, 0, 0},
            .numSurfaces  = 3,
        },
        {
            /* 422R */
            .heightFactor = {1, 0.5, 0.5, 0, 0, 0},
            .widthFactor  = {1, 1, 1, 0, 0, 0},
            .numSurfaces  = 3,
        },
    },
    {
        /* SEMI_PLANAR */
        {
            /* 420 */
            .heightFactor = {1, 0.5, 0, 0, 0, 0},
            .widthFactor  = {1, 0.5, 0, 0, 0, 0},
            .numSurfaces  = 2,
        },
        {
            /* 422 */
            .heightFactor = {1, 1, 0, 0, 0, 0},
            .widthFactor  = {1, 0.5, 0, 0, 0, 0},
            .numSurfaces  = 2,
        },
        {
            /* 444 */
            .heightFactor = {1, 1, 0.5, 0, 0, 0},
            .widthFactor  = {1, 1, 0.5, 0, 0, 0},
            .numSurfaces  = 2,
        },
        {
            /* 422R */
            .heightFactor = {1, 0.5, 0.5, 0, 0, 0},
            .widthFactor  = {1, 1, 0.5, 0, 0, 0},
            .numSurfaces  = 2,
        },
    },
    {
        /* PACKED */
        {
            /* 420 */
            .heightFactor = {1, 0.5, 0.5, 0, 0, 0},
            .widthFactor  = {1, 0.5, 0.5, 0, 0, 0},
            .numSurfaces  = 3,
        },
        {
            /* 422 */
            .heightFactor = {1, 1, 1, 0, 0, 0},
            .widthFactor  = {1, 0.5, 0.5, 0, 0, 0},
            .numSurfaces  = 3,
        },
        {
            /* 444 */
            .heightFactor = {1, 1, 1, 0, 0, 0},
            .widthFactor  = {1, 1, 1, 0, 0, 0},
            .numSurfaces  = 3,
        },
        {
            /* 422R */
            .heightFactor = {1, 0.5, 0.5, 0, 0, 0},
            .widthFactor  = {1, 1, 1, 0, 0, 0},
            .numSurfaces  = 3,
        },
    },
};

ImgUtilSurfParams ImgSurfParamsTable_Packed = {
    .heightFactor = {1, 0, 0, 0, 0, 0},
    .widthFactor  = {1, 0, 0, 0, 0, 0},
    .numSurfaces  = 1,
};


unsigned int ImgBytesPerPixelTable_RGBA[][6] = {
    {4, 0, 0, 0, 0, 0}, /* 8 */
};

unsigned int ImgBytesPerPixelTable_RGBA16[][6] = {
    {8, 0, 0, 0, 0, 0}, /* 16 */
};

unsigned int ImgBytesPerPixelTable_RG16[6] = {4, 0, 0, 0, 0, 0};

unsigned int ImgBytesPerPixelTable_Alpha[][6] = {
    {1, 0, 0, 0, 0, 0}, /* 8 */
    {2, 0, 0, 0, 0, 0}, /* 10 */
    {2, 0, 0, 0, 0, 0}, /* 12 */
    {2, 0, 0, 0, 0, 0}, /* 14 */
    {2, 0, 0, 0, 0, 0}, /* 16 */
    {4, 0, 0, 0, 0, 0}, /* 32 */
};

unsigned int ImgBytesPerPixelTable_RAW[][6] = {
    {1, 0, 0, 0, 0, 0}, /* 8 */
    {2, 0, 0, 0, 0, 0}, /* 10 */
    {2, 0, 0, 0, 0, 0}, /* 12 */
    {2, 0, 0, 0, 0, 0}, /* 14 */
    {2, 0, 0, 0, 0, 0}, /* 16 */
    {4, 0, 0, 0, 0, 0}, /* 32 */
    {4, 0, 0, 0, 0, 0}, /* 16_8_8 */
    {4, 0, 0, 0, 0, 0}, /* 10_8_8 */
    {4, 0, 0, 0, 0, 0}, /* 2_10_10_10 */
    {4, 0, 0, 0, 0, 0}, /* 20 */
};

unsigned int ImgBytesPerPixelTable_YUV[][9][6] = {{
                                                      /* PLANAR */
                                                      {1, 1, 1, 0, 0, 0}, /* 8 */
                                                      {2, 2, 2, 0, 0, 0}, /* 10 */
                                                      {2, 2, 2, 0, 0, 0}, /* 12 */
                                                      {2, 2, 2, 0, 0, 0}, /* 14 */
                                                      {2, 2, 2, 0, 0, 0}, /* 16 */
                                                      {4, 4, 4, 0, 0, 0}, /* 32 */
                                                      {2, 1, 1, 0, 0, 0}, /* 16_8_8 */
                                                      {2, 1, 1, 0, 0, 0}, /* 10_8_8 */
                                                      {4, 0, 0, 0, 0, 0}, /* 2_10_10_10 */
                                                  },
                                                  {
                                                      /* SEMI_PLANAR */
                                                      {1, 2, 0, 0, 0, 0}, /* 8 */
                                                      {2, 4, 0, 0, 0, 0}, /* 10 */
                                                      {2, 4, 0, 0, 0, 0}, /* 12 */
                                                      {2, 4, 0, 0, 0, 0}, /* 14 */
                                                      {2, 4, 0, 0, 0, 0}, /* 16 */
                                                      {4, 8, 0, 0, 0, 0}, /* 32 */
                                                      {2, 2, 0, 0, 0, 0}, /* 16_8_8 */
                                                      {2, 2, 0, 0, 0, 0}, /* 10_8_8 */
                                                      {4, 0, 0, 0, 0, 0}, /* 2_10_10_10 */
                                                  }};

static NvMediaStatus GetBytesPerCompForPackedYUV(unsigned int surfBPCidx, unsigned int *bytespercomp)
{
    switch (surfBPCidx) {
    case NVM_SURF_ATTR_BITS_PER_COMPONENT_8:
    case NVM_SURF_ATTR_BITS_PER_COMPONENT_LAYOUT_2_10_10_10:
        *bytespercomp = 1;
        break;
    case NVM_SURF_ATTR_BITS_PER_COMPONENT_10:
    case NVM_SURF_ATTR_BITS_PER_COMPONENT_12:
    case NVM_SURF_ATTR_BITS_PER_COMPONENT_14:
    case NVM_SURF_ATTR_BITS_PER_COMPONENT_16:
        *bytespercomp = 2;
        break;
    case NVM_SURF_ATTR_BITS_PER_COMPONENT_20:
        *bytespercomp = 3;
        break;
    case NVM_SURF_ATTR_BITS_PER_COMPONENT_32:
        *bytespercomp = 4;
        break;
    default:
        return NVMEDIA_STATUS_ERROR;
    }
    return NVMEDIA_STATUS_OK;
}

static NvMediaStatus GetSurfParams(unsigned int   surfaceType,
                                   float        **xScale,
                                   float        **yScale,
                                   unsigned int **bytePerPixel,
                                   uint32_t      *numSurfacesVal)
{
    NvMediaStatus status;
    unsigned int  surfType, surfMemoryType, surfSubSamplingType, surfBPC, surfCompOrder;
    float        *xScalePtr = NULL, *yScalePtr = NULL;
    unsigned int *bytePerPixelPtr = NULL;
    NVM_SURF_FMT_DEFINE_ATTR(srcAttr);
    uint32_t            numSurfaces     = 1;
    static unsigned int yuvpackedtbl[6] = {1, 0, 0, 0, 0, 0};
    unsigned int        numcomps        = 1;

    status = NvMediaSurfaceFormatGetAttrs(surfaceType, srcAttr, NVM_SURF_FMT_ATTR_MAX);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s:NvMediaSurfaceFormatGetAttrs failed\n", __func__);
        return NVMEDIA_STATUS_ERROR;
    }

    surfType            = srcAttr[NVM_SURF_ATTR_SURF_TYPE].value;
    surfMemoryType      = srcAttr[NVM_SURF_ATTR_MEMORY].value;
    surfSubSamplingType = srcAttr[NVM_SURF_ATTR_SUB_SAMPLING_TYPE].value;
    surfBPC             = srcAttr[NVM_SURF_ATTR_BITS_PER_COMPONENT].value;
    surfCompOrder       = srcAttr[NVM_SURF_ATTR_COMPONENT_ORDER].value;

    switch (surfType) {
    case NVM_SURF_ATTR_SURF_TYPE_YUV:
        if (surfSubSamplingType == NVM_SURF_ATTR_SUB_SAMPLING_TYPE_NONE
            && surfMemoryType == NVM_SURF_ATTR_MEMORY_PACKED) {

            xScalePtr   = &ImgSurfParamsTable_Packed.widthFactor[0];
            yScalePtr   = &ImgSurfParamsTable_Packed.heightFactor[0];
            numSurfaces = ImgSurfParamsTable_Packed.numSurfaces;

            if (NVMEDIA_STATUS_OK != GetBytesPerCompForPackedYUV(surfBPC, &yuvpackedtbl[0])) {
                printf("Invalid Bits per component and Packed YUV combination\n");
                return NVMEDIA_STATUS_ERROR;
            }

            switch (surfCompOrder) {
            case NVM_SURF_ATTR_COMPONENT_ORDER_VUYX:
            case NVM_SURF_ATTR_COMPONENT_ORDER_XYUV:
            case NVM_SURF_ATTR_COMPONENT_ORDER_XUYV:
                numcomps = 4;
                break;
            case NVM_SURF_ATTR_COMPONENT_ORDER_UYVY:
            case NVM_SURF_ATTR_COMPONENT_ORDER_VYUY:
            case NVM_SURF_ATTR_COMPONENT_ORDER_YVYU:
            case NVM_SURF_ATTR_COMPONENT_ORDER_YUYV:
                numcomps = 2;
                break;
            case NVM_SURF_ATTR_COMPONENT_ORDER_LUMA:
                numcomps = 1;
                break;
            default:
                printf("Invalid component Order  and Packed YUV combination\n");
                return NVMEDIA_STATUS_ERROR;
            }
            yuvpackedtbl[0] = yuvpackedtbl[0] * numcomps;
            bytePerPixelPtr = &yuvpackedtbl[0];
        }
        else {
            xScalePtr =
                &ImgSurfParamsTable_YUV[0][surfSubSamplingType - NVM_SURF_ATTR_SUB_SAMPLING_TYPE_420].widthFactor[0];
            yScalePtr =
                &ImgSurfParamsTable_YUV[0][surfSubSamplingType - NVM_SURF_ATTR_SUB_SAMPLING_TYPE_420].heightFactor[0];
            numSurfaces =
                ImgSurfParamsTable_YUV[0][surfSubSamplingType - NVM_SURF_ATTR_SUB_SAMPLING_TYPE_420].numSurfaces;
            bytePerPixelPtr = &ImgBytesPerPixelTable_YUV[0][surfBPC - NVM_SURF_ATTR_BITS_PER_COMPONENT_8][0];
        }

        break;
    case NVM_SURF_ATTR_SURF_TYPE_RGBA:
        if (surfCompOrder == NVM_SURF_ATTR_COMPONENT_ORDER_ALPHA) {
            bytePerPixelPtr = &ImgBytesPerPixelTable_Alpha[surfBPC - NVM_SURF_ATTR_BITS_PER_COMPONENT_8][0];
        }
        else if (surfCompOrder == NVM_SURF_ATTR_COMPONENT_ORDER_RG) {
            if (surfBPC == NVM_SURF_ATTR_BITS_PER_COMPONENT_16) {
                bytePerPixelPtr = &ImgBytesPerPixelTable_RG16[0];
            }
            else {
                printf("Invalid RGorder & Bitspercomp combination.Only RG16 is supported\n");
                return NVMEDIA_STATUS_ERROR;
            }
        }
        else { /* RGBA, ARGB, BGRA */
            if (surfBPC == NVM_SURF_ATTR_BITS_PER_COMPONENT_16) {
                bytePerPixelPtr = &ImgBytesPerPixelTable_RGBA16[0][0];
            }
            else if (surfBPC == NVM_SURF_ATTR_BITS_PER_COMPONENT_8) {
                bytePerPixelPtr = &ImgBytesPerPixelTable_RGBA[0][0];
            }
            else {
                printf("RGBA orders with 8 and 16bits only is supported \n");
                return NVMEDIA_STATUS_ERROR;
            }
        }
        xScalePtr   = &ImgSurfParamsTable_RGBA.widthFactor[0];
        yScalePtr   = &ImgSurfParamsTable_RGBA.heightFactor[0];
        numSurfaces = ImgSurfParamsTable_RGBA.numSurfaces;
        break;
    case NVM_SURF_ATTR_SURF_TYPE_RAW:
        bytePerPixelPtr = &ImgBytesPerPixelTable_RAW[surfBPC - NVM_SURF_ATTR_BITS_PER_COMPONENT_8][0];
        xScalePtr       = &ImgSurfParamsTable_RAW.widthFactor[0];
        yScalePtr       = &ImgSurfParamsTable_RAW.heightFactor[0];
        numSurfaces     = ImgSurfParamsTable_RAW.numSurfaces;
        break;
    default:
        printf("%s: Unsupported Pixel Format %d", __func__, surfType);
        return NVMEDIA_STATUS_ERROR;
    }

    if (xScale) {
        *xScale = xScalePtr;
    }
    if (yScale) {
        *yScale = yScalePtr;
    }
    if (bytePerPixel) {
        *bytePerPixel = bytePerPixelPtr;
    }
    if (numSurfacesVal) {
        *numSurfacesVal = numSurfaces;
    }

    return NVMEDIA_STATUS_OK;
}

NvMediaStatus
AllocateBufferToWriteImage(Blit2DTest *ctx, NvMediaImage *image, NvMediaBool uvOrderFlag, NvMediaBool appendFlag)
{
    uint32_t      imageSize = 0;
    unsigned int  size[3]   = {0};
    uint8_t      *buffer    = NULL;
    uint32_t      i, k, newk = 0;
    unsigned int *bytePerPixelPtr = NULL;
    ctx->numSurfaces              = 1;
    NvMediaImageSurfaceMap surfaceMap;
    NvMediaStatus          status = NVMEDIA_STATUS_ERROR;
    uint32_t               lineWidth, numRows, startOffset;

    if (!image) {
        printf("%s: Bad parameter\n", __func__);
        return NVMEDIA_STATUS_BAD_PARAMETER;
    }

    status = NvMediaImageLock(image, NVMEDIA_IMAGE_ACCESS_WRITE, &surfaceMap);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: NvMediaImageLock() failed\n", __func__);
        return status;
    }
    NvMediaImageUnlock(image);

    ctx->dstBuff = (uint8_t **)malloc(sizeof(uint8_t *) * MAXM_NUM_SURFACES);
    if (!ctx->dstBuff) {
        printf("%s: Out of memory\n", __func__);
        status = NVMEDIA_STATUS_OUT_OF_MEMORY;
        goto done;
    }

    ctx->dstBuffPitches = (uint32_t *)calloc(1, sizeof(uint32_t) * MAXM_NUM_SURFACES);
    if (!ctx->dstBuffPitches) {
        printf("%s: Out of memory\n", __func__);
        status = NVMEDIA_STATUS_OUT_OF_MEMORY;
        goto done;
    }

    ctx->heightSurface = surfaceMap.height;
    ctx->widthSurface  = surfaceMap.width;

    status = GetSurfParams(image->type, &ctx->xScalePtr, &ctx->yScalePtr, &bytePerPixelPtr, &ctx->numSurfaces);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: GetSurfParams() failed\n", __func__);
        goto done;
    }

    imageSize = 0;
    for (i = 0; i < ctx->numSurfaces; i++) {
        size[i] = (ctx->widthSurface * ctx->xScalePtr[i] * ctx->heightSurface * ctx->yScalePtr[i] * bytePerPixelPtr[i]);
        imageSize += size[i];
        ctx->dstBuffPitches[i] = (uint32_t)((float)ctx->widthSurface * ctx->xScalePtr[i]) * bytePerPixelPtr[i];
    }

    // Embedded data size needs to be included for RAW surftype
    size[0] += image->embeddedDataTopSize;
    size[0] += image->embeddedDataBottomSize;
    imageSize += image->embeddedDataTopSize;
    imageSize += image->embeddedDataBottomSize;

    buffer = (uint8_t *)calloc(1, imageSize);
    if (!buffer) {
        printf("%s: Out of memory\n", __func__);
        status = NVMEDIA_STATUS_OUT_OF_MEMORY;
        goto done;
    }

    ctx->dstBuffer = buffer;
    memset(buffer, 0xFF, imageSize);
    for (i = 0; i < ctx->numSurfaces; i++) {
        ctx->dstBuff[i] = buffer;
        buffer          = buffer + (uint32_t)(ctx->heightSurface * ctx->yScalePtr[i] * ctx->dstBuffPitches[i]);
    }

done:
    return status;
}

NvMediaStatus WriteImageToAllocatedBuffer(Blit2DTest   *ctx,
                                          NvMediaImage *image,
                                          NvMediaBool   uvOrderFlag,
                                          NvMediaBool   appendFlag,
                                          uint32_t      bytesPerPixel)
{
    NvMediaImageSurfaceMap surfaceMap;

    NvMediaStatus status = NvMediaImageLock(image, NVMEDIA_IMAGE_ACCESS_WRITE, &surfaceMap);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: NvMediaImageLock() failed\n", __func__);
        goto done;
    }
    status = NvMediaImageGetBits(image, NULL, (void **)ctx->dstBuff, ctx->dstBuffPitches);
    NvMediaImageUnlock(image);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: NvMediaVideoSurfaceGetBits() failed \n", __func__);
        goto done;
    }

done:

    return status;
}


static NvMediaStatus ReadImageNew(char         *fileName,
                                  uint32_t      frameNum,
                                  uint32_t      width,
                                  uint32_t      height,
                                  NvMediaImage *image,
                                  NvMediaBool   uvOrderFlag,
                                  uint32_t      bytesPerPixel,
                                  uint32_t      pixelAlignment)
{
    uint8_t              **pBuff        = NULL;
    uint32_t              *pBuffPitches = NULL;
    uint32_t               imageSize = 0, surfaceSize = 0;
    uint8_t               *buffer  = NULL;
    uint8_t               *pBuffer = NULL;
    uint32_t               i, j, k, newk = 0;
    float                 *xScalePtr = NULL, *yScalePtr = NULL;
    unsigned int          *bytePerPixelPtr = NULL;
    uint32_t               numSurfaces     = 1;
    unsigned int           uHeightSurface, uWidthSurface;
    NvMediaImageSurfaceMap surfaceMap;
    NvMediaStatus          status = NVMEDIA_STATUS_ERROR;
    FILE                  *file   = NULL;
    unsigned int           count, index;
    NVM_SURF_FMT_DEFINE_ATTR(srcAttr);
    unsigned int surfType, surfBPC;

    if (!image || !fileName) {
        printf("%s: Bad parameter\n", __func__);
        return NVMEDIA_STATUS_BAD_PARAMETER;
    }

    status = NvMediaImageLock(image, NVMEDIA_IMAGE_ACCESS_WRITE, &surfaceMap);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: NvMediaImageLock() failed\n", __func__);
        return status;
    }
    NvMediaImageUnlock(image);

    uHeightSurface = surfaceMap.height;
    uWidthSurface  = surfaceMap.width;

    if (width > uWidthSurface || height > uHeightSurface) {
        printf("%s: Bad parameter\n", __func__);
        return NVMEDIA_STATUS_BAD_PARAMETER;
    }

    pBuff = (uint8_t **)malloc(sizeof(uint8_t *) * MAXM_NUM_SURFACES);
    if (!pBuff) {
        printf("%s: Out of memory\n", __func__);
        status = NVMEDIA_STATUS_OUT_OF_MEMORY;
        goto done;
    }

    pBuffPitches = (uint32_t *)calloc(1, sizeof(uint32_t) * MAXM_NUM_SURFACES);
    if (!pBuffPitches) {
        printf("%s: Out of memory\n", __func__);
        status = NVMEDIA_STATUS_OUT_OF_MEMORY;
        goto done;
    }

    status = GetSurfParams(image->type, &xScalePtr, &yScalePtr, &bytePerPixelPtr, &numSurfaces);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: GetSurfParams() failed\n", __func__);
        goto done;
    }

    status = NvMediaSurfaceFormatGetAttrs(image->type, srcAttr, NVM_SURF_FMT_ATTR_MAX);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s:NvMediaSurfaceFormatGetAttrs failed\n", __func__);
        goto done;
    }
    surfType = srcAttr[NVM_SURF_ATTR_SURF_TYPE].value;
    surfBPC  = srcAttr[NVM_SURF_ATTR_BITS_PER_COMPONENT].value;

    surfaceSize = 0;
    imageSize   = 0;
    for (i = 0; i < numSurfaces; i++) {
        surfaceSize += (uWidthSurface * xScalePtr[i] * uHeightSurface * yScalePtr[i] * bytePerPixelPtr[i]);
        imageSize += (width * xScalePtr[i] * height * yScalePtr[i] * bytePerPixelPtr[i]);
        pBuffPitches[i] = (uint32_t)((float)uWidthSurface * xScalePtr[i]) * bytePerPixelPtr[i];
    }

    buffer = (uint8_t *)calloc(1, surfaceSize);
    if (!buffer) {
        printf("%s: Out of memory\n", __func__);
        status = NVMEDIA_STATUS_OUT_OF_MEMORY;
        goto done;
    }

    pBuffer = buffer;
    memset(buffer, 0x10, surfaceSize);
    for (i = 0; i < numSurfaces; i++) {
        pBuff[i] = buffer;
        if (i) {
            memset(pBuff[i], 0x80, (uHeightSurface * yScalePtr[i] * pBuffPitches[i]));
        }
        buffer = buffer + (uint32_t)(uHeightSurface * yScalePtr[i] * pBuffPitches[i]);
    }

    file = fopen(fileName, "rb");
    if (!file) {
        printf("%s: Error opening file: %s\n", __func__, fileName);
        status = NVMEDIA_STATUS_ERROR;
        goto done;
    }

    if (frameNum > 0) {
        if (fseeko(file, frameNum * (off_t)imageSize, SEEK_SET)) {
            printf("ReadImage: Error seeking file: %s\n", fileName);
            status = NVMEDIA_STATUS_ERROR;
            goto done;
        }
    }

    if ((surfType == NVM_SURF_ATTR_SURF_TYPE_RGBA) && strstr(fileName, ".png")) {
        printf("ReadImage: Does not support png format\n");
        status = NVMEDIA_STATUS_ERROR;
        goto done;
    }
    for (k = 0; k < numSurfaces; k++) {
        for (j = 0; j < height * yScalePtr[k]; j++) {
            newk  = (!uvOrderFlag && k) ? (numSurfaces - k) : k;
            index = j * pBuffPitches[newk];
            count = width * xScalePtr[newk] * bytePerPixelPtr[newk];
            if (fread(pBuff[newk] + index, count, 1, file) != 1) {
                status = NVMEDIA_STATUS_ERROR;
                printf("ReadImage: Error reading file: %s\n", fileName);
                goto done;
            }
            if ((surfType == NVM_SURF_ATTR_SURF_TYPE_YUV) && (pixelAlignment == LSB_ALIGNED)) {
                uint16_t *psrc = (uint16_t *)(pBuff[newk] + index);
                switch (surfBPC) {
                case NVM_SURF_ATTR_BITS_PER_COMPONENT_10:
                    for (i = 0; i < count / 2; i++) {
                        *(psrc + i) = (*(psrc + i)) << (16 - 10);
                    }
                    break;
                case NVM_SURF_ATTR_BITS_PER_COMPONENT_12:
                    for (i = 0; i < count / 2; i++) {
                        *(psrc + i) = (*(psrc + i)) << (16 - 12);
                    }
                    break;
                case NVM_SURF_ATTR_BITS_PER_COMPONENT_14:
                    for (i = 0; i < count / 2; i++) {
                        *(psrc + i) = (*(psrc + i)) << (16 - 14);
                    }
                    break;
                default:
                    break;
                }
            }
        }
    }

    status = NvMediaImageLock(image, NVMEDIA_IMAGE_ACCESS_WRITE, &surfaceMap);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: NvMediaImageLock() failed\n", __func__);
        goto done;
    }
    status = NvMediaImagePutBits(image, NULL, (void **)pBuff, pBuffPitches);
    NvMediaImageUnlock(image);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: Failed to put bits\n", __func__);
    }

done:
    if (pBuff) {
        free(pBuff);
    }

    if (pBuffPitches) {
        free(pBuffPitches);
    }

    if (pBuffer) {
        free(pBuffer);
    }

    if (file) {
        fclose(file);
    }

    return status;
}

NvMediaStatus ReadImage(char         *fileName,
                        uint32_t      frameNum,
                        uint32_t      width,
                        uint32_t      height,
                        NvMediaImage *image,
                        NvMediaBool   uvOrderFlag,
                        uint32_t      bytesPerPixel,
                        uint32_t      pixelAlignment)
{
    NvMediaStatus status;
    NVM_SURF_FMT_DEFINE_ATTR(srcAttr);

    status = NvMediaSurfaceFormatGetAttrs(image->type, srcAttr, NVM_SURF_FMT_ATTR_MAX);
    if (status == NVMEDIA_STATUS_OK) {
        return ReadImageNew(fileName, frameNum, width, height, image, uvOrderFlag, bytesPerPixel, pixelAlignment);
    }
    else {
        printf("%s:NvMediaSurfaceFormatGetAttrs failed\n", __func__);
        return status;
    }
}

NvMediaStatus InitImage(NvMediaImage *image, uint32_t width, uint32_t height)
{
    uint8_t              **pBuff        = NULL;
    uint32_t              *pBuffPitches = NULL;
    uint32_t               imageSize = 0, surfaceSize = 0;
    uint8_t               *buffer    = NULL;
    uint8_t               *pBuffer   = NULL;
    float                 *xScalePtr = NULL, *yScalePtr = NULL;
    unsigned int          *bytePerPixelPtr = NULL;
    uint32_t               numSurfaces     = 1;
    uint32_t               i;
    unsigned int           uHeightSurface, uWidthSurface;
    NvMediaImageSurfaceMap surfaceMap;
    NvMediaStatus          status = NVMEDIA_STATUS_ERROR;
    NVM_SURF_FMT_DEFINE_ATTR(srcAttr);

    if (!image) {
        printf("%s: Bad parameter\n", __func__);
        return NVMEDIA_STATUS_BAD_PARAMETER;
    }

    status = NvMediaImageLock(image, NVMEDIA_IMAGE_ACCESS_WRITE, &surfaceMap);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: NvMediaImageLock failed\n", __func__);
        return status;
    }
    NvMediaImageUnlock(image);


    uHeightSurface = surfaceMap.height;
    uWidthSurface  = surfaceMap.width;

    if (width > uWidthSurface || height > uHeightSurface) {
        printf("%s: Bad parameter\n", __func__);
        return NVMEDIA_STATUS_BAD_PARAMETER;
    }

    pBuff = (uint8_t **)calloc(1, sizeof(uint8_t *) * MAXM_NUM_SURFACES);
    if (!pBuff) {
        printf("%s: out of memory\n", __func__);
        status = NVMEDIA_STATUS_OUT_OF_MEMORY;
        goto done;
    }

    pBuffPitches = (uint32_t *)calloc(1, sizeof(uint32_t) * MAXM_NUM_SURFACES);
    if (!pBuffPitches) {
        printf("%s: out of memory\n", __func__);
        status = NVMEDIA_STATUS_OUT_OF_MEMORY;
        goto done;
    }

    status = GetSurfParams(image->type, &xScalePtr, &yScalePtr, &bytePerPixelPtr, &numSurfaces);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: GetSurfParams failed\n", __func__);
        goto done;
    }

    status = NvMediaSurfaceFormatGetAttrs(image->type, srcAttr, NVM_SURF_FMT_ATTR_MAX);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: NvMediaSurfaceFormatGetAttrs failed\n", __func__);
        goto done;
    }

    surfaceSize = 0;
    imageSize   = 0;
    for (i = 0; i < numSurfaces; i++) {
        surfaceSize += (uWidthSurface * xScalePtr[i] * uHeightSurface * yScalePtr[i] * bytePerPixelPtr[i]);
        imageSize += (width * xScalePtr[i] * height * yScalePtr[i] * bytePerPixelPtr[i]);
        pBuffPitches[i] = (uint32_t)((float)uWidthSurface * xScalePtr[i]) * bytePerPixelPtr[i];
    }

    buffer = (uint8_t *)calloc(1, surfaceSize);
    if (!buffer) {
        printf("%s: out of memory\n", __func__);
        status = NVMEDIA_STATUS_OUT_OF_MEMORY;
        goto done;
    }

    pBuffer = buffer;
    memset(buffer, 0x00, surfaceSize);
    for (i = 0; i < numSurfaces; i++) {
        pBuff[i] = buffer;
        buffer   = buffer + (uint32_t)(uHeightSurface * yScalePtr[i] * pBuffPitches[i]);
    }

    status = NvMediaImageLock(image, NVMEDIA_IMAGE_ACCESS_WRITE, &surfaceMap);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: NvMediaImageLock failed\n", __func__);
        goto done;
    }
    status = NvMediaImagePutBits(image, NULL, (void **)pBuff, pBuffPitches);
    NvMediaImageUnlock(image);
    if (status != NVMEDIA_STATUS_OK) {
        printf("%s: NvMediaImagePutBits failed\n", __func__);
    }

done:
    if (pBuff) {
        free(pBuff);
    }

    if (pBuffPitches) {
        free(pBuffPitches);
    }

    if (pBuffer) {
        free(pBuffer);
    }

    return status;
}
