#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <vector>
#include <string>
#include "nvdspreprocess_interface.h"
#ifndef CONTEXTRA_RESOLUTION
#define CONTEXTRA_RESOLUTION 576
#endif

struct CustomCtx { cudaStream_t stream; bool fp16; };

template<typename T> __device__ T storeValue(float value);
template<> __device__ float storeValue<float>(float value) { return value; }
template<> __device__ __half storeValue<__half>(float value) { return __float2half_rn(value); }

// MODELCONTRACT-01: NVIDIA가 resize/RGB 변환한 device pixels를 NCHW로 정규화한다.
template<typename T>
__global__ void normalizeRgb(const unsigned char* rgb, T* out, unsigned pitch, unsigned channels)
{
    const unsigned i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= CONTEXTRA_RESOLUTION * CONTEXTRA_RESOLUTION) return;
    const float means[3] = {0.485f, 0.456f, 0.406f};
    const float stds[3] = {0.229f, 0.224f, 0.225f};
    for (unsigned c = 0; c < 3; ++c) {
        const float pixel = rgb[(i / CONTEXTRA_RESOLUTION) * pitch + (i % CONTEXTRA_RESOLUTION) * channels + c];
        // FP16 input storage does not change the FP32 normalization arithmetic.
        out[c * CONTEXTRA_RESOLUTION * CONTEXTRA_RESOLUTION + i] = storeValue<T>((pixel / 255.0f - means[c]) / stds[c]);
    }
}

extern "C" CustomCtx* initLib(CustomInitParams p)
{
    const auto& t = p.tensor_params;
    if (t.network_input_order != NvDsPreProcessNetworkInputOrder_kNCHW ||
        t.network_input_shape != std::vector<int>({1,3,CONTEXTRA_RESOLUTION,CONTEXTRA_RESOLUTION}) ||
        t.network_color_format != NvDsPreProcessFormat_RGB ||
        (t.data_type != NvDsDataType_FP32 && t.data_type != NvDsDataType_FP16) ||
        t.tensor_name != "input") return nullptr;
    auto* ctx = new CustomCtx{};
    ctx->fp16 = t.data_type == NvDsDataType_FP16;
    if (cudaStreamCreateWithFlags(&ctx->stream, cudaStreamNonBlocking) != cudaSuccess) {
        delete ctx;
        return nullptr;
    }
    return ctx;
}

extern "C" void deInitLib(CustomCtx* ctx)
{
    if (ctx) { cudaStreamDestroy(ctx->stream); delete ctx; }
}

extern "C" NvDsPreProcessStatus CustomTransformation(
    NvBufSurface* in, NvBufSurface* out, CustomTransformParams& p)
{
    if (NvBufSurfTransformSetSessionParams(&p.transform_config_params) != NvBufSurfTransformError_Success)
        return NVDSPREPROCESS_CUSTOM_TRANSFORMATION_FAILED;
    return NvBufSurfTransform(in, out, &p.transform_params) == NvBufSurfTransformError_Success
        ? NVDSPREPROCESS_SUCCESS : NVDSPREPROCESS_CUSTOM_TRANSFORMATION_FAILED;
}

extern "C" NvDsPreProcessStatus CustomTensorPreparation(
    CustomCtx* ctx, NvDsPreProcessBatch* batch, NvDsPreProcessCustomBuf*& buf,
    CustomTensorParams& params, NvDsPreProcessAcquirer* acquirer)
{
    buf = nullptr;
    if (!ctx || !batch || !acquirer || batch->units.size() != 1 ||
        !batch->units[0].converted_frame_ptr ||
        params.params.data_type != (ctx->fp16 ? NvDsDataType_FP16 : NvDsDataType_FP32) ||
        params.params.buffer_size < 3 * CONTEXTRA_RESOLUTION * CONTEXTRA_RESOLUTION * (ctx->fp16 ? sizeof(__half) : sizeof(float)))
        return NVDSPREPROCESS_INVALID_PARAMS;
    const unsigned channels = batch->scaling_pool_format == NvDsPreProcessFormat_RGB ? 3 :
        batch->scaling_pool_format == NvDsPreProcessFormat_RGBA ? 4 : 0;
    if (!channels || batch->pitch < CONTEXTRA_RESOLUTION * channels) return NVDSPREPROCESS_INVALID_PARAMS;
    buf = acquirer->acquire();
    if (!buf || !buf->memory_ptr) {
        if (buf) acquirer->release(buf);
        buf = nullptr;
        return NVDSPREPROCESS_RESOURCE_ERROR;
    }
    if (ctx->fp16) {
        normalizeRgb<<<(CONTEXTRA_RESOLUTION*CONTEXTRA_RESOLUTION+255)/256,256,0,ctx->stream>>>(
            static_cast<const unsigned char*>(batch->units[0].converted_frame_ptr),
            static_cast<__half*>(buf->memory_ptr), batch->pitch, channels);
    } else {
        normalizeRgb<<<(CONTEXTRA_RESOLUTION*CONTEXTRA_RESOLUTION+255)/256,256,0,ctx->stream>>>(
            static_cast<const unsigned char*>(batch->units[0].converted_frame_ptr),
            static_cast<float*>(buf->memory_ptr), batch->pitch, channels);
    }
    const auto launch = cudaGetLastError();
    const auto sync = cudaStreamSynchronize(ctx->stream);
    if (launch != cudaSuccess || sync != cudaSuccess) {
        acquirer->release(buf); buf = nullptr;
        return NVDSPREPROCESS_CUDA_ERROR;
    }
    params.params.network_input_shape = {1,3,CONTEXTRA_RESOLUTION,CONTEXTRA_RESOLUTION};
    return NVDSPREPROCESS_SUCCESS;
}