/* 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.
 */

// This sample needs at least CUDA 10.1. It demonstrates usages of the nvJPEG
// library nvJPEG encoder supports single and multiple image encode.

#include <cuda_runtime_api.h>
#include "helper_nvJPEG.hxx"


int dev_malloc(void **p, size_t s) { return (int)cudaMalloc(p, s); }
int dev_free(void *p) { return (int)cudaFree(p); }

bool is_interleaved(nvjpegOutputFormat_t format)
{
    if (format == NVJPEG_OUTPUT_RGBI || format == NVJPEG_OUTPUT_BGRI)
        return true;
    else
        return false;
}

struct encode_params_t {
  std::string input_dir;
  std::string output_dir;
  std::string format;
  std::string subsampling;
  int quality;
  int huf;
  int dev;
};

nvjpegEncoderParams_t encode_params;
nvjpegHandle_t nvjpeg_handle;
nvjpegJpegState_t jpeg_state;
nvjpegEncoderState_t encoder_state;

int decodeEncodeOneImage(std::string sImagePath, std::string sOutputPath, double &time, nvjpegOutputFormat_t output_format, nvjpegInputFormat_t input_format)
{
    time = 0.;
    cudaEvent_t startEvent = NULL, stopEvent = NULL;
    float loopTime = 0;
    checkCudaErrors(cudaEventCreate(&startEvent, cudaEventBlockingSync));
    checkCudaErrors(cudaEventCreate(&stopEvent, cudaEventBlockingSync));

    // Get the file name, without extension.
    // This will be used to rename the output file.    
    size_t position = sImagePath.rfind("/");
    std::string sFileName = (std::string::npos == position)? sImagePath : sImagePath.substr(position + 1, sImagePath.size());
    position = sFileName.rfind(".");
    sFileName = (std::string::npos == position)? sFileName : sFileName.substr(0, position);
    position = sFileName.rfind("/");
    sFileName = (std::string::npos == position) ? sFileName : sFileName.substr(position + 1, sFileName.length());
    position = sFileName.rfind("\\");
    sFileName = (std::string::npos == position) ? sFileName : sFileName.substr(position+1, sFileName.length());

    // Read an image from disk.
    std::ifstream oInputStream(sImagePath.c_str(), std::ios::in | std::ios::binary | std::ios::ate);
    if(!(oInputStream.is_open()))
    {
        std::cerr << "Cannot open image: " << sImagePath << std::endl;
        return 1;
    }
    
    // Get the size.
    std::streamsize nSize = oInputStream.tellg();
    oInputStream.seekg(0, std::ios::beg);

    // Image buffers. 
    unsigned char * pBuffer = NULL; 
    double encoder_time = 0.;
    
    std::vector<char> vBuffer(nSize);
    
    if (oInputStream.read(vBuffer.data(), nSize))
    {            
        unsigned char * dpImage = (unsigned char *)vBuffer.data();
        
        // Retrieve the componenet and size info.
        int nComponent = 0;
        nvjpegChromaSubsampling_t subsampling;
        int widths[NVJPEG_MAX_COMPONENT];
        int heights[NVJPEG_MAX_COMPONENT];
        if (NVJPEG_STATUS_SUCCESS != nvjpegGetImageInfo(nvjpeg_handle, dpImage, nSize, &nComponent, &subsampling, widths, heights))
        {
            std::cerr << "Error decoding JPEG header: " << sImagePath << std::endl;
            return 1;
        }

        // image information
        std::cout << "Image is " << nComponent << " channels." << std::endl;
        for (int i = 0; i < nComponent; i++)
        {
            std::cout << "Channel #" << i << " size: "  << widths[i]  << " x " << heights[i] << std::endl;    
        }
        
        switch (subsampling)
        {
            case NVJPEG_CSS_444:
                std::cout << "YUV 4:4:4 chroma subsampling" << std::endl;
                break;
            case NVJPEG_CSS_440:
                std::cout << "YUV 4:4:0 chroma subsampling" << std::endl;
                break;
            case NVJPEG_CSS_422:
                std::cout << "YUV 4:2:2 chroma subsampling" << std::endl;
                break;
            case NVJPEG_CSS_420:
                std::cout << "YUV 4:2:0 chroma subsampling" << std::endl;
                break;
            case NVJPEG_CSS_411:
                std::cout << "YUV 4:1:1 chroma subsampling" << std::endl;
                break;
            case NVJPEG_CSS_410:
                std::cout << "YUV 4:1:0 chroma subsampling" << std::endl;
                break;
            case NVJPEG_CSS_GRAY:
                std::cout << "Grayscale JPEG " << std::endl;
                break;
            case NVJPEG_CSS_UNKNOWN: 
                std::cout << "Unknown chroma subsampling" << std::endl;
                return 1;
        }

        {

            cudaError_t eCopy = cudaMalloc(&pBuffer, widths[0] * heights[0] * NVJPEG_MAX_COMPONENT);
            if(cudaSuccess != eCopy) 
            {
                std::cerr << "cudaMalloc failed for component Y: " << cudaGetErrorString(eCopy) << std::endl;
                return 1;
            }

            nvjpegImage_t imgdesc = 
            {
                {
                    pBuffer,
                    pBuffer + widths[0]*heights[0],
                    pBuffer + widths[0]*heights[0]*2,
                    pBuffer + widths[0]*heights[0]*3
                },
                {
                    (unsigned int)(is_interleaved(output_format) ? widths[0] * 3 : widths[0]),
                    (unsigned int)widths[0],
                    (unsigned int)widths[0],
                    (unsigned int)widths[0]
                }
            };
           
            int nReturnCode = 0;

            cudaDeviceSynchronize();

            nReturnCode = nvjpegDecode(nvjpeg_handle, jpeg_state, dpImage, nSize, output_format, &imgdesc, NULL);

            // alternatively decode by stages
            /*int nReturnCode = nvjpegDecodeCPU(nvjpeg_handle, dpImage, nSize, output_format, &imgdesc, NULL);
            nReturnCode = nvjpegDecodeMixed(nvjpeg_handle, NULL);
            nReturnCode = nvjpegDecodeGPU(nvjpeg_handle, NULL);*/
            cudaDeviceSynchronize();

            if(nReturnCode != 0)
            {
                std::cerr << "Error in nvjpegDecode." << std::endl;
                return 1;
            }

            checkCudaErrors(cudaEventRecord(startEvent, NULL));
            /////////////////////// encode ////////////////////
            if (NVJPEG_OUTPUT_YUV == output_format)
            {
                checkCudaErrors(nvjpegEncodeYUV(nvjpeg_handle,
                    encoder_state,
                    encode_params,
                    &imgdesc,
                    subsampling,
                    widths[0],
                    heights[0],
                    NULL));
            }
            else
            {
                checkCudaErrors(nvjpegEncodeImage(nvjpeg_handle,
                    encoder_state,
                    encode_params,
                    &imgdesc,
                    input_format,
                    widths[0],
                    heights[0],
                    NULL));
            }

            std::vector<unsigned char> obuffer;
            size_t length;
            checkCudaErrors(nvjpegEncodeRetrieveBitstream(
                nvjpeg_handle,
                encoder_state,
                NULL,
                &length,
                NULL));
            obuffer.resize(length);
            checkCudaErrors(nvjpegEncodeRetrieveBitstream(
                nvjpeg_handle,
                encoder_state,
                obuffer.data(),
                &length,
                NULL));

            checkCudaErrors(cudaEventRecord(stopEvent, NULL));
            checkCudaErrors(cudaEventSynchronize(stopEvent));
            checkCudaErrors(cudaEventElapsedTime(&loopTime, startEvent, stopEvent));
            encoder_time = static_cast<double>(loopTime);

            std::string output_filename = sOutputPath + "/" + sFileName + ".jpg";
            char directory[120];
            char mkdir_cmd[256];
            std::string folder = sOutputPath;
            output_filename = folder + "/"+ sFileName +".jpg";
#if !defined(_WIN32)
            sprintf(directory, "%s", folder.c_str());
            sprintf(mkdir_cmd, "mkdir -p %s 2> /dev/null", directory);
#else
            sprintf(directory, "%s", folder.c_str());
            sprintf(mkdir_cmd, "mkdir %s 2> nul", directory);
#endif

            int ret = system(mkdir_cmd);

            std::cout << "Writing JPEG file: " << output_filename << std::endl;
            std::ofstream outputFile(output_filename.c_str(), std::ios::out | std::ios::binary);
            outputFile.write(reinterpret_cast<const char *>(obuffer.data()), static_cast<int>(length));
            
            // Free memory
            checkCudaErrors(cudaFree(pBuffer));
        }
    }

    time = encoder_time;

    return 0;
}

int processArgs(encode_params_t param)
{
    std::string sInputPath(param.input_dir);
    std::string sOutputPath(param.output_dir);
    std::string sFormat(param.format);
    std::string sSubsampling(param.subsampling);
    nvjpegOutputFormat_t oformat = NVJPEG_OUTPUT_RGB;
    nvjpegInputFormat_t iformat = NVJPEG_INPUT_RGB;

    int error_code = 1;

    if (sFormat == "yuv")
    {
        oformat = NVJPEG_OUTPUT_YUV;
    } 
    else if (sFormat == "rgb")
    {
        oformat = NVJPEG_OUTPUT_RGB;
        iformat = NVJPEG_INPUT_RGB;
    }
    else if (sFormat == "bgr")
    {
        oformat = NVJPEG_OUTPUT_BGR;
        iformat = NVJPEG_INPUT_BGR;
    }
    else if (sFormat == "rgbi")
    {
        oformat = NVJPEG_OUTPUT_RGBI;
        iformat = NVJPEG_INPUT_RGBI;
    }
    else if (sFormat == "bgri")
    {
        oformat = NVJPEG_OUTPUT_BGRI;
        iformat = NVJPEG_INPUT_BGRI;
    }
    else 
    {
        std::cerr << "Unknown or unsupported output format: " << sFormat << std::endl;
        return error_code;
    }

    if (sSubsampling == "444")
    {
        checkCudaErrors(nvjpegEncoderParamsSetSamplingFactors(encode_params, NVJPEG_CSS_444, NULL));
    }
    else if (sSubsampling == "422")
    {
        checkCudaErrors(nvjpegEncoderParamsSetSamplingFactors(encode_params, NVJPEG_CSS_422, NULL));
    }
    else if (sSubsampling == "420")
    {
        checkCudaErrors(nvjpegEncoderParamsSetSamplingFactors(encode_params, NVJPEG_CSS_420, NULL));
    }
    else if (sSubsampling == "440")
    {
        checkCudaErrors(nvjpegEncoderParamsSetSamplingFactors(encode_params, NVJPEG_CSS_440, NULL));
    }
    else if (sSubsampling == "411")
    {
        checkCudaErrors(nvjpegEncoderParamsSetSamplingFactors(encode_params, NVJPEG_CSS_411, NULL));
    }
    else if (sSubsampling == "410")
    {
        checkCudaErrors(nvjpegEncoderParamsSetSamplingFactors(encode_params, NVJPEG_CSS_410, NULL));
    }
    else if (sSubsampling == "400")
    {
        checkCudaErrors(nvjpegEncoderParamsSetSamplingFactors(encode_params, NVJPEG_CSS_GRAY, NULL));
    }
    else 
    {
        std::cerr << "Unknown or unsupported subsampling: " << sSubsampling << std::endl;
        return error_code;
    }
    /*if( stat(sOutputPath.c_str(), &s) == 0 )
    {
        if( !(s.st_mode & S_IFDIR) )
        {
            std::cout << "Output path already exist as non-directory: " << sOutputPath << std::endl;
            return error_code;
        }
    }
    else
    {
        if (mkdir(sOutputPath.c_str(), 0775))
        {
            std::cout << "Cannot create output directory: " << sOutputPath << std::endl;
            return error_code;
        }
    }*/

    std::vector<std::string> inputFiles;
    if (readInput(sInputPath, inputFiles))
    {
        return error_code;
    }
    
    double total_time = 0., encoder_time = 0.;
    int total_images = 0;

    for (unsigned int i = 0; i < inputFiles.size(); i++)
    {
        std::string &sFileName = inputFiles[i];
        std::cout << "Processing file: " << sFileName << std::endl;
        int image_error_code = decodeEncodeOneImage(sFileName, sOutputPath, encoder_time, oformat, iformat);
        if (image_error_code)
        {
            std::cerr << "Error processing file: " << sFileName << std::endl;
            //return image_error_code;
        }
        else
        {
            total_images++;
            total_time += encoder_time;
        }                      
    }
    std::cout << "Total images processed: " << total_images << std::endl;
    std::cout << "Total time spent on encoding: " << total_time << std::endl;
    std::cout << "Avg time/image: " << total_time/total_images << std::endl;

    return 0;
}

// parse parameters
int findParamIndex(const char **argv, int argc, const char *parm) {
  int count = 0;
  int index = -1;

  for (int i = 0; i < argc; i++) {
    if (strncmp(argv[i], parm, 100) == 0) {
      index = i;
      count++;
    }
  }

  if (count == 0 || count == 1) {
    return index;
  } else {
    std::cout << "Error, parameter " << parm
              << " has been specified more than once, exiting\n"
              << std::endl;
    return -1;
  }

  return -1;
}


int main(int argc, const char *argv[]) 
{
  int pidx;

  if ((pidx = findParamIndex(argv, argc, "-h")) != -1 ||
      (pidx = findParamIndex(argv, argc, "--help")) != -1) {
    std::cout << "Usage: " << argv[0]
              << " -i images_dir  [-o output_dir] [-device=device_id]"                 
                 "[-q quality][-s 420/444] [-fmt output_format] [-huf 0]\n";
    std::cout << "Parameters: " << std::endl;
    std::cout << "\timages_dir\t:\tPath to single image or directory of images" << std::endl;
    std::cout << "\toutput_dir\t:\tWrite encoded images as jpeg to this directory" << std::endl;
    std::cout << "\tdevice_id\t:\tWhich device to use for encoding" << std::endl;
    std::cout << "\tQuality\t:\tUse image quality [default 70]" << std::endl;
    std::cout << "\tsubsampling\t:\tUse Subsampling [420, 444]" << std::endl;
    std::cout << "\toutput_format\t:\tnvJPEG output format for encoding. One "
                 "of [rgb, rgbi, bgr, bgri, yuv, y, unchanged]"
              << std::endl;
    std::cout << "\tHuffman Optimization\t:\tUse Huffman optimization [default 0]" << std::endl;
    return EXIT_SUCCESS;
  }

  encode_params_t params;

  params.input_dir = "./";
  if ((pidx = findParamIndex(argv, argc, "-i")) != -1) {
    params.input_dir = argv[pidx + 1];
  } else {
    // Search in default paths for input images.
    int found = getInputDir(params.input_dir, argv[0]);
    if (!found)
    {
      std::cout << "Please specify input directory with encoded images"<< std::endl;
      return EXIT_WAIVED;
    }
  }
  if ((pidx = findParamIndex(argv, argc, "-o")) != -1) {
    params.output_dir = argv[pidx + 1];
  } else {
      // by-default write the folder named "output" in cwd
      params.output_dir = "encode_output";
  }
  params.dev = 0;
  params.dev = findCudaDevice(argc, argv);

  params.quality = 70;
  if ((pidx = findParamIndex(argv, argc, "-q")) != -1) {
    params.quality = std::atoi(argv[pidx + 1]);
  }

  if ((pidx = findParamIndex(argv, argc, "-s")) != -1) {
    params.subsampling = argv[pidx + 1];
  } else {
      // by-default use subsampling as 420
      params.subsampling = "420";
  }
  if ((pidx = findParamIndex(argv, argc, "-fmt")) != -1) {
    params.format = argv[pidx + 1];
  } else {
   // by-default use output format yuv
    params.format = "yuv";
  }

  params.huf = 0;
  if ((pidx = findParamIndex(argv, argc, "-huf")) != -1) {
    params.huf = std::atoi(argv[pidx + 1]);
  }

    cudaDeviceProp props;
    checkCudaErrors(cudaGetDeviceProperties(&props, params.dev));

    printf("Using GPU %d (%s, %d SMs, %d th/SM max, CC %d.%d, ECC %s)\n",
         params.dev, props.name, props.multiProcessorCount,
         props.maxThreadsPerMultiProcessor, props.major, props.minor,
         props.ECCEnabled ? "on" : "off");

    nvjpegDevAllocator_t dev_allocator = {&dev_malloc, &dev_free};
    checkCudaErrors(nvjpegCreate(NVJPEG_BACKEND_DEFAULT, &dev_allocator, &nvjpeg_handle));
    checkCudaErrors(nvjpegJpegStateCreate(nvjpeg_handle, &jpeg_state));
    checkCudaErrors(nvjpegEncoderStateCreate(nvjpeg_handle, &encoder_state, NULL));
    checkCudaErrors(nvjpegEncoderParamsCreate(nvjpeg_handle, &encode_params, NULL));
    
    // sample input parameters
    checkCudaErrors(nvjpegEncoderParamsSetQuality(encode_params, params.quality, NULL));
    checkCudaErrors(nvjpegEncoderParamsSetOptimizedHuffman(encode_params, params.huf, NULL));

    pidx = processArgs(params);

    checkCudaErrors(nvjpegEncoderParamsDestroy(encode_params));
    checkCudaErrors(nvjpegEncoderStateDestroy(encoder_state));
    checkCudaErrors(nvjpegJpegStateDestroy(jpeg_state));
    checkCudaErrors(nvjpegDestroy(nvjpeg_handle));

    return pidx;
}