/* 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 <cuda.h>
#include <cuda_runtime.h>
#include <fstream>
#include <iostream>
#include <stdlib.h>
#include <sys/stat.h>
#include <sys/types.h>

#include "resize_convert.h"
#include "utils.h"

__global__ void floatToChar(float *src, unsigned char *dst, int height, int width, int batchSize)
{
    int x = threadIdx.x + blockIdx.x * blockDim.x;

    if (x >= height * width)
        return;

    int offset = height * width * 3;

    for (int j = 0; j < batchSize; j++) {
        // b
        *(dst + j * offset + x * 3 + 0) = (unsigned char)*(src + j * offset + height * width * 0 + x);
        // g
        *(dst + j * offset + x * 3 + 1) = (unsigned char)*(src + j * offset + height * width * 1 + x);
        // r
        *(dst + j * offset + x * 3 + 2) = (unsigned char)*(src + j * offset + height * width * 2 + x);
    }
}

void floatPlanarToChar(float *src, unsigned char *dst, int height, int width, int batchSize)
{
    floatToChar<<<(height * width - 1) / 1024 + 1, 1024, 0, NULL>>>(src, dst, height, width, batchSize);
}

void dumpRawBGR(float *d_srcBGR, int pitch, int width, int height, int batchSize, char *folder, char *tag)
{
    float *bgr, *d_bgr;
    int    frameSize;
    char   directory[120];
    char   mkdir_cmd[256];
#if !defined(_WIN32)
    sprintf(directory, "output/%s", folder);
    sprintf(mkdir_cmd, "mkdir -p %s 2> /dev/null", directory);
#else
    sprintf(directory, "output\\%s", folder);
    sprintf(mkdir_cmd, "mkdir %s 2> nul", directory);
#endif

    int ret = system(mkdir_cmd);

    frameSize = width * height * 3 * sizeof(float);
    bgr       = (float *)malloc(frameSize);
    if (bgr == NULL) {
        std::cerr << "Failed malloc for bgr\n";
        return;
    }

    d_bgr = d_srcBGR;
    for (int i = 0; i < batchSize; i++) {
        char           filename[256];
        std::ofstream *outputFile;

        checkCudaErrors(cudaMemcpy((void *)bgr, (void *)d_bgr, frameSize, cudaMemcpyDeviceToHost));
        snprintf(filename, sizeof(filename), "%s/%s_%d.raw", directory, tag, (i + 1));

        outputFile = new std::ofstream(filename);
        if (outputFile) {
            outputFile->write((char *)bgr, frameSize);
            delete outputFile;
        }

        d_bgr += pitch * height * 3;
    }

    free(bgr);
}

void dumpBGR(float *d_srcBGR, int pitch, int width, int height, int batchSize, char *folder, char *tag)
{
    dumpRawBGR(d_srcBGR, pitch, width, height, batchSize, folder, tag);
}

void dumpYUV(unsigned char *d_nv12, int size, char *folder, char *tag)
{
    unsigned char *nv12Data;
    std::ofstream *nv12File;
    char           filename[256];
    char           directory[60];
    char           mkdir_cmd[256];
#if !defined(_WIN32)
    sprintf(directory, "output/%s", folder);
    sprintf(mkdir_cmd, "mkdir -p %s 2> /dev/null", directory);
#else
    sprintf(directory, "output\\%s", folder);
    sprintf(mkdir_cmd, "mkdir %s 2> nul", directory);
#endif

    int ret = system(mkdir_cmd);

    snprintf(filename, sizeof(filename), "%s/%s.nv12", directory, tag);

    nv12File = new std::ofstream(filename);
    if (nv12File == NULL) {
        std::cerr << "Failed to new " << filename;
        return;
    }

    nv12Data = (unsigned char *)malloc(size * (sizeof(char)));
    if (nv12Data == NULL) {
        std::cerr << "Failed to allcoate memory\n";
        return;
    }

    cudaMemcpy((void *)nv12Data, (void *)d_nv12, size, cudaMemcpyDeviceToHost);

    nv12File->write((const char *)nv12Data, size);

    free(nv12Data);
    delete nv12File;
}
