/*=========================================================================
 *
 *  Copyright NumFOCUS
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         https://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/
#ifndef itkGPUDenseFiniteDifferenceImageFilter_hxx
#define itkGPUDenseFiniteDifferenceImageFilter_hxx


#include <list>
#include "itkImageRegionIterator.h"
#include "itkNumericTraits.h"
#include "itkNeighborhoodAlgorithm.h"
#include "itkOpenCLUtil.h"

namespace itk
{
template <typename TInputImage, typename TOutputImage, typename TParentImageFilter>
GPUDenseFiniteDifferenceImageFilter<TInputImage, TOutputImage, TParentImageFilter>::
  GPUDenseFiniteDifferenceImageFilter()
{
  /**
   * FiniteDifferenceImageFilter requires two GPU kernels
   * 1. CalculateChange
   * 2. ApplyUpdate
   * Kernel for 1 is defined in the specific DifferenceFunction.
   * Kernel for 2 can be used for entire subclass of GPUDense..,
   * it is defined and created here.
   */

  std::ostringstream defines;

  if (TInputImage::ImageDimension > 3 || ImageDimension < 1)
  {
    itkExceptionMacro("GPUDenseFiniteDifferenceImageFilter supports 1/2/3D image.");
  }

  defines << "#define DIM_" << TInputImage::ImageDimension << '\n';

  // PixelType is a Vector
  defines << "#define BUFPIXELTYPE ";
  GetTypenameInString(typeid(typename UpdateBufferType::PixelType), defines);

  defines << "#define OUTPIXELTYPE ";
  GetTypenameInString(typeid(typename TOutputImage::PixelType), defines);

  // assumes input and output pixel type is same
  defines << "#define PIXELDIM " << GetPixelDimension(typeid(typename TOutputImage::PixelType)) << '\n';

  const char * GPUSource = GPUDenseFiniteDifferenceImageFilter::GetOpenCLSource();

  // load and build program
  this->m_GPUKernelManager->LoadProgramFromString(GPUSource, defines.str().c_str());

  // create kernel
  m_ApplyUpdateGPUKernelHandle = this->m_GPUKernelManager->CreateKernel("ApplyUpdate");
}

template <typename TInputImage, typename TOutputImage, typename TParentImageFilter>
void
GPUDenseFiniteDifferenceImageFilter<TInputImage, TOutputImage, TParentImageFilter>::CopyInputToOutput()
{
  CPUSuperclass::CopyInputToOutput();

  // Marking GPU is dirty by setting CPU as modified
  this->GetOutput()->Modified();
}

template <typename TInputImage, typename TOutputImage, typename TParentImageFilter>
void
GPUDenseFiniteDifferenceImageFilter<TInputImage, TOutputImage, TParentImageFilter>::AllocateUpdateBuffer()
{
  // CPUSuperclass will call Image::Allocate() which will call GPUImage::Allocate() .
  CPUSuperclass::AllocateUpdateBuffer();
}

template <typename TInputImage, typename TOutputImage, typename TParentImageFilter>
void
GPUDenseFiniteDifferenceImageFilter<TInputImage, TOutputImage, TParentImageFilter>::ApplyUpdate(const TimeStepType & dt)
{
  this->GPUApplyUpdate(dt);
}

template <typename TInputImage, typename TOutputImage, typename TParentImageFilter>
void
GPUDenseFiniteDifferenceImageFilter<TInputImage, TOutputImage, TParentImageFilter>::GPUApplyUpdate(
  const TimeStepType & dt)
{
  // GPU version of ApplyUpdate
  // Single threaded version : Apply entire update buffer to output image
  using GPUBufferImage = typename itk::GPUTraits<UpdateBufferType>::Type;
  using GPUOutputImage = typename itk::GPUTraits<TOutputImage>::Type;

  typename GPUBufferImage::Pointer bfPtr = dynamic_cast<GPUBufferImage *>(GetUpdateBuffer());
  typename GPUOutputImage::Pointer otPtr =
    dynamic_cast<GPUOutputImage *>(this->GetOutput()); // this->ProcessObject::GetOutput(0)
                                                       // );
  typename GPUOutputImage::SizeType outSize = otPtr->GetLargestPossibleRegion().GetSize();

  int imgSize[3];
  imgSize[0] = imgSize[1] = imgSize[2] = 1;

  float timeStep = dt;

  int ImageDim = static_cast<int>(TInputImage::ImageDimension);

  for (int i = 0; i < ImageDim; ++i)
  {
    imgSize[i] = outSize[i];
  }

  size_t localSize[3], globalSize[3];
  localSize[0] = localSize[1] = localSize[2] = OpenCLGetLocalBlockSize(ImageDim);
  for (int i = 0; i < ImageDim; ++i)
  {
    // total # of threads
    globalSize[i] =
      localSize[i] * static_cast<unsigned int>(ceil(static_cast<float>(outSize[i]) / static_cast<float>(localSize[i])));
  }

  // arguments set up
  int argidx = 0;
  this->m_GPUKernelManager->SetKernelArgWithImage(m_ApplyUpdateGPUKernelHandle, argidx++, bfPtr->GetGPUDataManager());
  this->m_GPUKernelManager->SetKernelArgWithImage(m_ApplyUpdateGPUKernelHandle, argidx++, otPtr->GetGPUDataManager());
  this->m_GPUKernelManager->SetKernelArg(m_ApplyUpdateGPUKernelHandle, argidx++, sizeof(float), &(timeStep));
  for (int i = 0; i < ImageDim; ++i)
  {
    this->m_GPUKernelManager->SetKernelArg(m_ApplyUpdateGPUKernelHandle, argidx++, sizeof(int), &(imgSize[i]));
  }

  // launch kernel
  this->m_GPUKernelManager->LaunchKernel(
    m_ApplyUpdateGPUKernelHandle, static_cast<int>(TInputImage::ImageDimension), globalSize, localSize);

  // Explicitly call Modified on GetOutput here. Do we need this?
  // this->GetOutput()->Modified();
}

template <typename TInputImage, typename TOutputImage, typename TParentImageFilter>
auto
GPUDenseFiniteDifferenceImageFilter<TInputImage, TOutputImage, TParentImageFilter>::GPUCalculateChange() -> TimeStepType
{
  typename OutputImageType::Pointer output = this->GetOutput();

  TimeStepType timeStep;
  void *       globalData;

  auto * df = dynamic_cast<GPUFiniteDifferenceFunction<OutputImageType> *>(this->GetDifferenceFunction().GetPointer());

  globalData = df->GetGlobalDataPointer();
  this->m_ComputeUpdateTime.Start();
  df->GPUComputeUpdate(output, GetUpdateBuffer(), globalData);
  this->m_ComputeUpdateTime.Stop();

  //// Ask the finite difference function to compute the time step for
  //// this iteration.  We give it the global data pointer to use, then
  //// ask it to free the global data memory.
  timeStep = df->ComputeGlobalTimeStep(globalData);
  df->ReleaseGlobalDataPointer(globalData);

  return timeStep;
}

template <typename TInputImage, typename TOutputImage, typename TParentImageFilter>
void
GPUDenseFiniteDifferenceImageFilter<TInputImage, TOutputImage, TParentImageFilter>::PrintSelf(std::ostream & os,
                                                                                              Indent indent) const
{
  GPUSuperclass::PrintSelf(os, indent);
}

} // end namespace itk

#endif
