text: /*
   * Copyright (c) Contributors to the Open 3D Engine Project.
   * For complete copyright and license terms please see the LICENSE at the root of this distribution.
   *
   * SPDX-License-Identifier: Apache-2.0 OR MIT
   *
   */
  
  #include <Atom/RHI/CommandList.h>
  #include <Atom/RHI/Factory.h>
  #include <Atom/RHI/FrameScheduler.h>
  #include <Atom/RHI/Image.h>
  #include <Atom/RHI/ImagePool.h>
  #include <Atom/RHI/ScopeProducerFunction.h>
  #include <Atom/RHI.Reflect/InputStreamLayoutBuilder.h>
  #include <Atom/RHI.Reflect/RenderAttachmentLayoutBuilder.h>
  
  #include <Atom/RPI.Public/Image/AttachmentImagePool.h>
  #include <Atom/RPI.Public/Image/ImageSystemInterface.h>
  #include <Atom/RPI.Public/Image/StreamingImage.h>
  #include <Atom/RPI.Public/Image/StreamingImagePool.h>
  #include <Atom/RPI.Public/Shader/Shader.h>
  
  #include <Atom/RPI.Reflect/Image/AttachmentImageAssetCreator.h>
  #include <Atom/RPI.Reflect/Shader/ShaderAsset.h>
  
  #include <AzCore/Math/Vector2.h>
  #include <AzCore/Math/Vector4.h>
  #include <AzCore/Serialization/SerializeContext.h>
  
  #include <RHI/MRTExampleComponent.h>
  #include <SampleComponentConfig.h>
  #include <SampleComponentManager.h>
  #include <Utils/Utils.h>
  
  namespace AtomSampleViewer
  {
      MRTExampleComponent::MRTExampleComponent()
      {
          m_time = 0.0f;
          m_attachmentID[0] = AZ::RHI::AttachmentId("redAttachemntID");
          m_attachmentID[1] = AZ::RHI::AttachmentId("greenAttachemntID");
          m_attachmentID[2] = AZ::RHI::AttachmentId("blueAttachemntID");
          m_clearValue = AZ::RHI::ClearValue::CreateVector4Float(0.0f, 0.0f, 0.0f, 0.0f);
  
          m_supportRHISamplePipeline = true;
      }
  
      void MRTExampleComponent::Reflect(AZ::ReflectContext* context)
      {
          if (AZ::SerializeContext* serializeContext = azrtti_cast<AZ::SerializeContext*>(context))
          {
              serializeContext->Class<MRTExampleComponent, AZ::Component>()
                  ->Version(0);
          }
      }
  
      void MRTExampleComponent::OnFramePrepare(AZ::RHI::FrameGraphBuilder& frameGraphBuilder)
      {
          m_time += 0.002f;
          float r = cosf(m_time);
          float g = sinf(m_time);
          float b = cosf(m_time + AZ::Constants::Pi);
          m_shaderResourceGroups[0]->SetConstant(m_shaderInputConstantIndices[0], AZ::Vector4(r, 0 ,0, 1));
          m_shaderResourceGroups[0]->SetConstant(m_shaderInputConstantIndices[1], AZ::Vector2(g, 0));
          m_shaderResourceGroups[0]->SetConstant(m_shaderInputConstantIndices[2], b);
          m_shaderResourceGroups[0]->Compile();
  
          BasicRHIComponent::OnFramePrepare(frameGraphBuilder);
      }
  
      void MRTExampleComponent::Activate()
      {
          // Init buffers & views
          CreateInputAssemblyBuffersAndViews();
          // Init render targets
          InitRenderTargets();
          // init pipeline state
          ReadRenderTargetShader();
          ReadScreenShader();
  
          CreateRenderTargetScope();
          CreateScreenScope();
  
          AZ::RHI::RHISystemNotificationBus::Handler::BusConnect();
      }
  
      void MRTExampleComponent::Deactivate()
      {
          AZ::RHI::RHISystemNotificationBus::Handler::BusDisconnect();
          m_windowContext = nullptr;
          m_bufferPool = nullptr;
          m_inputAssemblyBuffer = nullptr;
          m_renderTargetImageDescriptors.fill({});
          m_shaderResourceGroups.fill(nullptr);
          m_pipelineStates.fill(nullptr);
          m_scopeProducers.clear();
      }
  
      void MRTExampleComponent::ReadRenderTargetShader()
      {
          using namespace AZ;
          const char* MRTTargetShaderFilePath = "Shaders/RHI/MRTTarget.azshader";
          const char* sampleName = "MRTExample";
  
          auto shader = LoadShader(MRTTargetShaderFilePath, sampleName);
          if (shader == nullptr)
              return;
  
          RHI::PipelineStateDescriptorForDraw pipelineDesc;
          shader->GetVariant(RPI::ShaderAsset::RootShaderVariantStableId).ConfigurePipelineState(pipelineDesc);
          pipelineDesc.m_inputStreamLayout = m_inputStreamLayout;
  
          RHI::RenderAttachmentLayoutBuilder attachmentsBuilder;
          attachmentsBuilder.AddSubpass()
              ->RenderTargetAttachment(RHI::Format::R8G8B8A8_UNORM_SRGB)
              ->RenderTargetAttachment(RHI::Format::R16G16_FLOAT)
              ->RenderTargetAttachment(RHI::Format::R32_FLOAT);
          [[maybe_unused]] RHI::ResultCode result = attachmentsBuilder.End(pipelineDesc.m_renderAttachmentConfiguration.m_renderAttachmentLayout);
          AZ_Assert(result == RHI::ResultCode::Success, "Failed to create render attachment layout");
  
          m_pipelineStates[0] = shader->AcquirePipelineState(pipelineDesc);
          if (!m_pipelineStates[0])
          {
              AZ_Error(sampleName, false, "Failed to acquire default pipeline state for shader '%s'", MRTTargetShaderFilePath);
              return;
          }
  
          m_shaderResourceGroups[0] = CreateShaderResourceGroup(shader, "MrtTargetSrg", sampleName);
  
          FindShaderInputIndex(&m_shaderInputConstantIndices[0], m_shaderResourceGroups[0], AZ::Name{"rValue"}, sampleName);
          FindShaderInputIndex(&m_shaderInputConstantIndices[1], m_shaderResourceGroups[0], AZ::Name{"gValue"}, sampleName);
          FindShaderInputIndex(&m_shaderInputConstantIndices[2], m_shaderResourceGroups[0], AZ::Name{"bValue"}, sampleName);
      }
  
      void MRTExampleComponent::ReadScreenShader()
      {
          using namespace AZ;
          const char* MRTScreenShaderFilePath = "Shaders/RHI/MRTScreen.azshader";
          const char* sampleName = "MRTExample";
  
          auto shader = LoadShader(MRTScreenShaderFilePath, sampleName);
          if (shader == nullptr)
              return;
  
          RHI::PipelineStateDescriptorForDraw pipelineDesc;
          shader->GetVariant(RPI::ShaderAsset::RootShaderVariantStableId).ConfigurePipelineState(pipelineDesc);
          pipelineDesc.m_inputStreamLayout = m_inputStreamLayout;
  
          RHI::RenderAttachmentLayoutBuilder attachmentsBuilder;
          attachmentsBuilder.AddSubpass()
              ->RenderTargetAttachment(m_outputFormat);
          [[maybe_unused]] RHI::ResultCode result = attachmentsBuilder.End(pipelineDesc.m_renderAttachmentConfiguration.m_renderAttachmentLayout);
          AZ_Assert(result == RHI::ResultCode::Success, "Failed to create render attachment layout");
  
          m_pipelineStates[1] = shader->AcquirePipelineState(pipelineDesc);
          if (!m_pipelineStates[1])
          {
              AZ_Error(sampleName, false, "Failed to acquire default pipeline state for shader '%s'", MRTScreenShaderFilePath);
              return;
          }
  
          m_shaderResourceGroups[1] = CreateShaderResourceGroup(shader, "MrtScreenInstanceSrg", sampleName);
  
          FindShaderInputIndex(&m_shaderInputImageIndices[0], m_shaderResourceGroups[1], AZ::Name{"rMap"}, sampleName);
          FindShaderInputIndex(&m_shaderInputImageIndices[1], m_shaderResourceGroups[1], AZ::Name{"gMap"}, sampleName);
          FindShaderInputIndex(&m_shaderInputImageIndices[2], m_shaderResourceGroups[1], AZ::Name{"bMap"}, sampleName);
      }
  
      void MRTExampleComponent::CreateInputAssemblyBuffersAndViews()
      {
          using namespace AZ;
          const RHI::Ptr<RHI::Device> device = Utils::GetRHIDevice();
  
          m_bufferPool = RHI::Factory::Get().CreateBufferPool();
          RHI::BufferPoolDescriptor bufferPoolDesc;
          bufferPoolDesc.m_bindFlags = RHI::BufferBindFlags::InputAssembly;
          bufferPoolDesc.m_heapMemoryLevel = RHI::HeapMemoryLevel::Device;
          m_bufferPool->Init(*device, bufferPoolDesc);
  
          BufferData bufferData;
          SetFullScreenRect(bufferData.m_positions.data(), bufferData.m_uvs.data(), bufferData.m_indices.data());
  
          m_inputAssemblyBuffer = RHI::Factory::Get().CreateBuffer();
          RHI::ResultCode result = RHI::ResultCode::Success;
          RHI::BufferInitRequest request;
  
          request.m_buffer = m_inputAssemblyBuffer.get();
          request.m_descriptor = RHI::BufferDescriptor{ RHI::BufferBindFlags::InputAssembly, sizeof(bufferData) };
          request.m_initialData = &bufferData;
          result = m_bufferPool->InitBuffer(request);
          if (result != RHI::ResultCode::Success)
          {
              AZ_Error("MRTExample", false, "Failed to initialize buffer with error code %d", result);
              return;
          }
  
          m_streamBufferViews[0] =
          {
              *m_inputAssemblyBuffer,
              offsetof(BufferData, m_positions),
              sizeof(BufferData::m_positions),
              sizeof(VertexPosition)
          };
  
          m_streamBufferViews[1] =
          {
              *m_inputAssemblyBuffer,
              offsetof(BufferData, m_uvs),
              sizeof(BufferData::m_uvs),
              sizeof(VertexUV)
          };
  
          m_indexBufferView =
          {
              *m_inputAssemblyBuffer,
              offsetof(BufferData, m_indices),
              sizeof(BufferData::m_indices),
              RHI::IndexFormat::Uint16
          };
  
          RHI::InputStreamLayoutBuilder layoutBuilder;
          layoutBuilder.AddBuffer()->Channel("POSITION", RHI::Format::R32G32B32_FLOAT);
          layoutBuilder.AddBuffer()->Channel("UV", RHI::Format::R32G32_FLOAT);
          m_inputStreamLayout = layoutBuilder.End();
  
          RHI::ValidateStreamBufferViews(m_inputStreamLayout, m_streamBufferViews);
      }
  
      void MRTExampleComponent::InitRenderTargets()
      {
          using namespace AZ;
  
          RHI::ImageDescriptor imageDesc;
  
          imageDesc = RHI::ImageDescriptor::Create2D(RHI::ImageBindFlags::Color | RHI::ImageBindFlags::ShaderReadWrite, 8, 8, RHI::Format::R8G8B8A8_UNORM_SRGB);
          m_renderTargetImageDescriptors[0] = RHI::TransientImageDescriptor(m_attachmentID[0], imageDesc);
  
          imageDesc = RHI::ImageDescriptor::Create2D(RHI::ImageBindFlags::Color | RHI::ImageBindFlags::ShaderReadWrite, 8, 8, RHI::Format::R16G16_FLOAT);
          m_renderTargetImageDescriptors[1] = RHI::TransientImageDescriptor(m_attachmentID[1], imageDesc);
  
          imageDesc = RHI::ImageDescriptor::Create2D(RHI::ImageBindFlags::Color | RHI::ImageBindFlags::ShaderReadWrite, 8, 8, RHI::Format::R32_FLOAT);
          m_renderTargetImageDescriptors[2] = RHI::TransientImageDescriptor(m_attachmentID[2], imageDesc);
      }
  
      void MRTExampleComponent::CreateRenderTargetScope()
      {
          using namespace AZ;
  
          const auto prepareFunctionRenderTarget = [this](RHI::FrameGraphInterface frameGraph, [[maybe_unused]] ScopeData& scopeData)
          {
              // Create & Binds RenderTarget images
              {
                  frameGraph.GetAttachmentDatabase().CreateTransientImage(m_renderTargetImageDescriptors[0]);
                  RHI::ImageScopeAttachmentDescriptor descR;
                  descR.m_attachmentId = m_attachmentID[0];
                  descR.m_loadStoreAction.m_clearValue = m_clearValue;
                  descR.m_loadStoreAction.m_loadAction = RHI::AttachmentLoadAction::Clear;
  
                  frameGraph.GetAttachmentDatabase().CreateTransientImage(m_renderTargetImageDescriptors[1]);
                  RHI::ImageScopeAttachmentDescriptor descG;
                  descG.m_attachmentId = m_attachmentID[1];
                  descG.m_loadStoreAction.m_clearValue = m_clearValue;
                  descG.m_loadStoreAction.m_loadAction = RHI::AttachmentLoadAction::Clear;
  
                  frameGraph.GetAttachmentDatabase().CreateTransientImage(m_renderTargetImageDescriptors[2]);
                  RHI::ImageScopeAttachmentDescriptor descB;
                  descB.m_attachmentId = m_attachmentID[2];
                  descB.m_loadStoreAction.m_clearValue = m_clearValue;
                  descB.m_loadStoreAction.m_loadAction = RHI::AttachmentLoadAction::Clear;
  
                  const AZStd::array<RHI::ImageScopeAttachmentDescriptor, 3> imageScopeArray{{descR, descG, descB}};
                  frameGraph.UseColorAttachments(imageScopeArray);
              }
  
              frameGraph.SetEstimatedItemCount(1);
          };
  
          RHI::EmptyCompileFunction<ScopeData> compileFunctionRenderTarget;
          const auto executeFunctionRenderTarget = [this](const RHI::FrameGraphExecuteContext& context, [[maybe_unused]] const ScopeData& scopeData)
          {
              RHI::CommandList* commandList = context.GetCommandList();
  
              // Set persistent viewport and scissor state.
              const auto& imageSize = m_renderTargetImageDescriptors[0].m_imageDescriptor.m_size;
              RHI::Viewport viewport(0, static_cast<float>(imageSize.m_width), 0, static_cast<float>(imageSize.m_height));
              RHI::Scissor scissor(0, 0, static_cast<int32_t>(imageSize.m_width), static_cast<int32_t>(imageSize.m_height));
              commandList->SetViewports(&viewport, 1);
              commandList->SetScissors(&scissor, 1);
  
              RHI::DrawIndexed drawIndexed;
              drawIndexed.m_indexCount = 6;
              drawIndexed.m_instanceCount = 1;
  
              const RHI::ShaderResourceGroup* shaderResourceGroups[] = { m_shaderResourceGroups[0]->GetRHIShaderResourceGroup() };
  
              RHI::DrawItem drawItem;
              drawItem.m_arguments = drawIndexed;
              drawItem.m_pipelineState = m_pipelineStates[0].get();
              drawItem.m_indexBufferView = &m_indexBufferView;
              drawItem.m_shaderResourceGroupCount = static_cast<uint8_t>(RHI::ArraySize(shaderResourceGroups));
              drawItem.m_shaderResourceGroups = shaderResourceGroups;
              drawItem.m_streamBufferViewCount = static_cast<uint8_t>(m_streamBufferViews.size());
              drawItem.m_streamBufferViews = m_streamBufferViews.data();
  
              commandList->Submit(drawItem);
          };
  
          m_scopeProducers.emplace_back(aznew RHI::ScopeProducerFunction<
              ScopeData,
              decltype(prepareFunctionRenderTarget),
              decltype(compileFunctionRenderTarget),
              decltype(executeFunctionRenderTarget)>(
                  RHI::ScopeId{"MRTTarget"},
                  ScopeData{},
                  prepareFunctionRenderTarget,
                  compileFunctionRenderTarget,
                  executeFunctionRenderTarget));
      }
  
      void MRTExampleComponent::CreateScreenScope()
      {
          using namespace AZ;
  
          const auto prepareFunctionScreen = [this](RHI::FrameGraphInterface frameGraph, [[maybe_unused]] ScopeData& scopeData)
          {
              // Binds the swap chain as a color attachment.
              {
                  RHI::ImageScopeAttachmentDescriptor descriptor;
                  descriptor.m_attachmentId = m_outputAttachmentId;
                  frameGraph.UseColorAttachment(descriptor);
              }
  
              // Bind shader attachments which are rendered by RenderTargetScope.
              {
                  frameGraph.UseShaderAttachment(RHI::ImageScopeAttachmentDescriptor(m_attachmentID[0]), RHI::ScopeAttachmentAccess::Read);
                  frameGraph.UseShaderAttachment(RHI::ImageScopeAttachmentDescriptor(m_attachmentID[1]), RHI::ScopeAttachmentAccess::Read);
                  frameGraph.UseShaderAttachment(RHI::ImageScopeAttachmentDescriptor(m_attachmentID[2]), RHI::ScopeAttachmentAccess::Read);
              }
  
              // We will submit a single draw item.
              frameGraph.SetEstimatedItemCount(1);
          };
  
          const auto compileFunctionScreen = [this](const AZ::RHI::FrameGraphCompileContext& context, [[maybe_unused]] const ScopeData& scopeData)
          {
              const AZ::RHI::ImageView* imageViewR = context.GetImageView(m_attachmentID[0]);
              const AZ::RHI::ImageView* imageViewG = context.GetImageView(m_attachmentID[1]);
              const AZ::RHI::ImageView* imageViewB = context.GetImageView(m_attachmentID[2]);
  
              m_shaderResourceGroups[1]->SetImageView(m_shaderInputImageIndices[0], imageViewR);
              m_shaderResourceGroups[1]->SetImageView(m_shaderInputImageIndices[1], imageViewG);
              m_shaderResourceGroups[1]->SetImageView(m_shaderInputImageIndices[2], imageViewB);
              m_shaderResourceGroups[1]->Compile();
          };
  
          const auto executeFunctionScreen = [this](const RHI::FrameGraphExecuteContext& context, [[maybe_unused]] const ScopeData& scopeData)
          {
              RHI::CommandList* commandList = context.GetCommandList();
  
              // Set persistent viewport and scissor state.
              commandList->SetViewports(&m_viewport, 1);
              commandList->SetScissors(&m_scissor, 1);
  
              RHI::DrawIndexed drawIndexed;
              drawIndexed.m_indexCount = 6;
              drawIndexed.m_instanceCount = 1;
  
              const RHI::ShaderResourceGroup* shaderResourceGroups[] = { m_shaderResourceGroups[1]->GetRHIShaderResourceGroup() };
  
              RHI::DrawItem drawItem;
              drawItem.m_arguments = drawIndexed;
              drawItem.m_pipelineState = m_pipelineStates[1].get();
              drawItem.m_indexBufferView = &m_indexBufferView;
              drawItem.m_shaderResourceGroupCount = static_cast<uint8_t>(RHI::ArraySize(shaderResourceGroups));
              drawItem.m_shaderResourceGroups = shaderResourceGroups;
              drawItem.m_streamBufferViewCount = static_cast<uint8_t>(m_streamBufferViews.size());
              drawItem.m_streamBufferViews = m_streamBufferViews.data();
  
              commandList->Submit(drawItem);
          };
  
          m_scopeProducers.emplace_back(
              aznew RHI::ScopeProducerFunction<
              ScopeData,
              decltype(prepareFunctionScreen),
              decltype(compileFunctionScreen),
              decltype(executeFunctionScreen)>(
                  RHI::ScopeId{"MRTScreen"},
                  ScopeData{},
                  prepareFunctionScreen,
                  compileFunctionScreen,
                  executeFunctionScreen));
      }
  }
  