//----------------------------------------------------------------------------------
// File:        CDX11TrueHDRDemo.cpp
// SDK Version: 1.0.2
//
// SPDX-FileCopyrightText: Copyright (c) 2023-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: LicenseRef-NvidiaProprietary
//
// NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
// property and proprietary rights in and to this material, related
// documentation and any modifications thereto. Any use, reproduction,
// disclosure or distribution of this material and related documentation
// without an express license agreement from NVIDIA CORPORATION or
// its affiliates is strictly prohibited.
//
//----------------------------------------------------------------------------------

///////////////////////////////////////////
// CDX11TrueHDRDemo.cpp
//
// This is an example of a DX11 class using the NGX wrapper class for TrueHDR.
// It initializes DX11 video devices, creates a video processor, and swap chains.
// The swap chains are created in the HDR format and in a standard ARGB format.
// It then reads uncompressed YUV (typically NV12) from a file.
// For TrueHDR it does a VpBlt from YUV to ARGB, 
//    then uses TrueHDR to go from ARGB to the HDR swap chain back buffer.
// For the comparison window, it does a VpBlt from YUV to the ARGB swap chain back buffer.
//

#include "CDX11TrueHDRDemo.h"

#include "CDx11NGXTrueHDR.h"


HRESULT CDX11TrueHDRDemo::Init(HWND hWndDisplay, HWND hWndDisplayAlt, Configuration_s* pConfiguration)
{
    m_CDx1xCommon.SetWindowOnNonNV(pConfiguration->bMonitorIsNotNV);

    m_pCDx11Api = new CDx11Api;
    if (!m_pCDx11Api) return E_OUTOFMEMORY;

    HRESULT hr = m_pCDx11Api->Init(&m_CDx1xCommon, hWndDisplay);
    if (FAILED(hr)) return hr;

    if (!m_CDx1xCommon.IsMonitorHDR())
    {
        fwprintf(stdout, L"Monitor is not in HDR mode. Exiting\n");
        return E_NOT_VALID_STATE;
    }

    if (!SetupStreamData(&m_StreamData, pConfiguration->pSzInputFile))
    {
        return E_INVALIDARG;
    }

    m_InWidth = m_StreamData.Width;
    m_InHeight = m_StreamData.Height;
    if (pConfiguration->bUseNativeSize)
    {
        m_OutWidth = m_InWidth;
        m_OutHeight = m_InHeight;
    }
    else
    {
        m_OutWidth = pConfiguration->DisplayWidth;
        m_OutHeight = pConfiguration->DisplayHeight;
    }
    m_StreamData.DstRect = { 0, 0, (LONG)m_OutWidth, (LONG)m_OutHeight };

    pConfiguration->WidthInUse = m_OutWidth;
    pConfiguration->HeightInUse = m_OutHeight;
    m_DstFormat = pConfiguration->bFP16 ? DXGI_FORMAT_R16G16B16A16_FLOAT : DXGI_FORMAT_R10G10B10A2_UNORM;
    m_DstFormatAlt = DXGI_FORMAT_R8G8B8A8_UNORM;

    m_pCNGXTrueHDR = new CDx11NGXTrueHDR;
    hr = m_pCNGXTrueHDR->CreateFeature(m_pCDx11Api->GetDevice());
    if (FAILED(hr))
    {
        SafeDelete(m_pCNGXTrueHDR);
        return hr;
    }
    // set input size for NGX
    m_rcNGXInput = { 0, 0, (LONG)m_OutWidth, (LONG)m_OutHeight };
    // set output size for NGX
    m_rcNGXOutput = { 0, 0, (LONG)m_OutWidth, (LONG)m_OutHeight };
    // set vp output format for ngx input format
    m_VPOutputFormat = DXGI_FORMAT_R8G8B8A8_UNORM;

    hr = CreateSurfaces();
    if (FAILED(hr)) return hr;

    hr = m_pCDx11Api->CreateVideoProcessor(m_VPOutputFormat, m_InWidth, m_InHeight, m_OutWidth, m_OutHeight);
    if (FAILED(hr)) return hr;

    hr = m_pCDx11Api->CreateSwapChain(hWndDisplay, m_DstFormat, m_OutWidth, m_OutHeight, &m_pSwapChain);
    if (FAILED(hr)) return hr;
    m_pSwapChain->UpdateSwapChainColorSpace(&m_CDx1xCommon, true);

    if (hWndDisplayAlt)
    {
        hr = m_pCDx11Api->CreateSwapChain(hWndDisplayAlt, m_DstFormatAlt, m_OutWidth, m_OutHeight, &m_pSwapChainAlt);
        if (FAILED(hr)) return hr;
        m_VPDefaultColorSpace = m_pSwapChainAlt->UpdateSwapChainColorSpace(&m_CDx1xCommon, false);
    }

    return hr;
}

void CDX11TrueHDRDemo::Shutdown()
{
    if (m_pCNGXTrueHDR)
    {
        m_pCNGXTrueHDR->ReleaseFeature();
        SafeDelete(m_pCNGXTrueHDR);
    }
    m_pCDx11Api->ShutDown();
    SafeDelete(m_pSwapChain);
    SafeDelete(m_pSwapChainAlt);
    SafeRelease(m_pNGXInputBuffer);
    CStreamData* pStreamData = &m_StreamData;
    SafeRelease(pStreamData->pSurface);
    SafeRelease(pStreamData->pStagingSurface);
    SafeDelete(pStreamData->pVideoFile);
}


HRESULT CDX11TrueHDRDemo::CreateSurfaces()
{
    HRESULT hr = S_OK;


    // Setup the input stream surface
    CStreamData* pStreamData = &m_StreamData;

    // create staging
    hr = m_pCDx11Api->CreateTexture2D(
                                pStreamData->Width,
                                pStreamData->Height,
                                pStreamData->StreamFormat,
                                1,
                                false,
                                true,
                                &pStreamData->pStagingSurface);
    if (FAILED(hr)) return hr;

    // create src surf
    hr = m_pCDx11Api->CreateTexture2D(
                                pStreamData->Width,
                                pStreamData->Height,
                                pStreamData->StreamFormat,
                                1,
                                true,
                                false,
                                &pStreamData->pSurface);
    if (FAILED(hr)) return hr;

    // create ngx input for output from vp
    hr = m_pCDx11Api->CreateTexture2D(
                                m_rcNGXInput.right,
                                m_rcNGXInput.bottom,
                                m_VPOutputFormat,
                                1,
                                true,
                                false,
                                &m_pNGXInputBuffer);
    if (FAILED(hr)) return hr;

    return S_OK;
}


bool CDX11TrueHDRDemo::SetupStreamData(CStreamData* pStreamData, char* szFile)
{
    strcpy_s(pStreamData->FileName, szFile);
    pStreamData->pVideoFile = new CVideoFile();
    char strFullPath[256];
    strcpy_s(strFullPath, pStreamData->FileName);
    HRESULT hr = pStreamData->pVideoFile->Open(strFullPath);
    if (FAILED(hr))
    {
        delete pStreamData->pVideoFile;
        pStreamData->pVideoFile = NULL;
        return false;
    }
    pStreamData->Enable = TRUE;
    pStreamData->ColorSpace = DXGI_COLOR_SPACE_YCBCR_STUDIO_G22_LEFT_P709;
    switch (pStreamData->pVideoFile->FourCC())
    {
    case FOURCC_NV12:
        pStreamData->StreamFormat = DXGI_FORMAT_NV12;
        break;
    case FOURCC_YUY2:
        pStreamData->StreamFormat = DXGI_FORMAT_YUY2;
        break;
    case FOURCC_ARGB:
        pStreamData->ColorSpace = DXGI_COLOR_SPACE_RGB_FULL_G22_NONE_P709;
        pStreamData->StreamFormat = DXGI_FORMAT_R8G8B8A8_UNORM;
        break;
    default:
        return false;
        break;
    }
    pStreamData->Width = pStreamData->pVideoFile->Width();
    pStreamData->Height = pStreamData->pVideoFile->Height();
    pStreamData->AlphaEnable = FALSE;
    pStreamData->Alpha = 1.0;
    pStreamData->AutoProcessingEnable = FALSE;
    pStreamData->DstRectEnable = TRUE;
    pStreamData->SrcRectEnable = FALSE;
    pStreamData->SrcRect = { 0, 0, pStreamData->Width, pStreamData->Height };
    pStreamData->DstRect = pStreamData->SrcRect;
    pStreamData->FrameFormat = D3D11_VIDEO_FRAME_FORMAT_PROGRESSIVE;
    pStreamData->LumaKeyEnable = FALSE;
    pStreamData->LumaKeyLower = 0.0f;
    pStreamData->LumaKeyUpper = 1.0f;
    pStreamData->OutputRate = D3D11_VIDEO_PROCESSOR_OUTPUT_RATE_NORMAL;
    pStreamData->OutputRateRepeatFrame = FALSE;
    pStreamData->OutputRateCustomRate.Denominator = 1;
    pStreamData->OutputRateCustomRate.Numerator = 1;
    pStreamData->PixelAspectRatioEnable = FALSE;
    pStreamData->PixelAspectRatioSource.Denominator = 1;
    pStreamData->PixelAspectRatioSource.Numerator = 1;
    pStreamData->PixelAspectRatioDestination.Denominator = 1;
    pStreamData->PixelAspectRatioDestination.Numerator = 1;
    pStreamData->RotationEnable = FALSE;
    pStreamData->Rotation = D3D11_VIDEO_PROCESSOR_ROTATION_IDENTITY;
    return true;
}


// Prepares the source surface for presentation by fill it with the video data.
// -----------------------------------------------------------------------------
HRESULT CDX11TrueHDRDemo::FillFrameFromFile(CStreamData* pStreamData, bool bBackward)
{
    // For interleaved content, each frame is two fields interleaved, so we should
    // only load 1 frame every 2 VPBlit() render calls (skip loading on every other frame).

    // Map the surface to fill with data from file
    D3D11_MAPPED_SUBRESOURCE Mapped = { 0 };
    HRESULT hr = m_pCDx11Api->MapStagingSurface(pStreamData->pStagingSurface, &Mapped, true);
    if (FAILED(hr)) return hr;

    if (m_uCurFrameCount)       // skip first frame
    {
        if (bBackward)
            pStreamData->CurFrameCount = pStreamData->pVideoFile->DecrementFrameCyclic();
        else
            pStreamData->CurFrameCount = pStreamData->pVideoFile->IncrementFrameCyclic();
    }
    // Fill the source buffer with the video data. This surface may be a different
    // format than the video file (ex. NV24 file to NV12 surface).
    hr = pStreamData->pVideoFile->FillBuffer((char*)Mapped.pData, pStreamData->pVideoFile->FourCC(), Mapped.RowPitch, false);
    if (FAILED(hr)) return hr;

    // unmap the surface
    hr = m_pCDx11Api->UnMapStagingSurface(pStreamData->pSurface, pStreamData->pStagingSurface, true, 0);

    return hr;
}

HRESULT CDX11TrueHDRDemo::NextFrame(bool bBackward)
{
    HRESULT hr = S_OK;
    CStreamData* pStreamData = &m_StreamData;

    hr = FillFrameFromFile(pStreamData, bBackward);
    if (FAILED(hr)) return hr;

    DXGI_COLOR_SPACE_TYPE outCSpace = m_pSwapChain->UpdateSwapChainColorSpace(&m_CDx1xCommon, true);

    ID3D11Texture2D*    pSrcSurf        = pStreamData->pSurface;
    RECT                srcRect         = pStreamData->SrcRect;

    ID3D11Texture2D* pSwapBackBuffer = nullptr;
    hr = m_pSwapChain->GetBuffer(&pSwapBackBuffer);
    if (FAILED(hr)) return hr;

    if (m_CDx1xCommon.IsMonitorHDR())
    {
        if (pStreamData->StreamFormat == DXGI_FORMAT_NV12 || pStreamData->StreamFormat == DXGI_FORMAT_YUY2)
        {
            // VpBlt to NGX input buffer
            hr = m_pCDx11Api->VpBltHd(&pSrcSurf, m_pNGXInputBuffer, m_rcNGXInput.right, m_rcNGXInput.bottom);
            if (FAILED(hr)) return hr;
            pSrcSurf = m_pNGXInputBuffer;
            srcRect  = m_rcNGXInput;
        }

        // Apply TrueHDR with EvaluateFeature
        UINT Contrast       = 100;
        UINT Saturation     = 100;
        UINT MiddleGray     = 50;
        UINT MaxLuminance   = m_CDx1xCommon.GetMonitorMaxLuminance();
        hr = m_pCNGXTrueHDR->EvaluateFeature(pSwapBackBuffer, m_rcNGXOutput,
                                                pSrcSurf, srcRect,
                                                Contrast, Saturation, MiddleGray, MaxLuminance);
        if (FAILED(hr)) return hr;
    }
    else
    {
        m_pCDx11Api->SetOutputColorSpace(outCSpace);
        hr = m_pCDx11Api->VpBltHd(&pSrcSurf, pSwapBackBuffer, m_rcNGXOutput.right, m_rcNGXOutput.bottom);
        if (FAILED(hr)) return hr;
        m_pCDx11Api->SetOutputColorSpace(m_VPDefaultColorSpace);
    }
    
    hr = m_pSwapChain->Present();
    if (FAILED(hr)) return hr;
    SafeRelease(pSwapBackBuffer);

    if (m_pSwapChainAlt)
    {
        m_pSwapChainAlt->GetBuffer(&pSwapBackBuffer);
        hr = m_pCDx11Api->VpBltHd(&m_StreamData.pSurface, pSwapBackBuffer, m_OutWidth, m_OutHeight);
        if (FAILED(hr)) return hr;

        hr = m_pSwapChainAlt->Present();
        if (FAILED(hr)) return hr;
        SafeRelease(pSwapBackBuffer);
    }

    m_uCurFrameCount++;
    return S_OK;
}

