/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkDataTransferHelper.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
#include "vtkDataTransferHelper.h"

#include "vtkDataArray.h"
#include "vtkImageData.h"
#include "vtkObjectFactory.h"
#include "vtkOpenGLRenderWindow.h"
#include "vtkPixelBufferObject.h"
#include "vtkSmartPointer.h"
#include "vtkStructuredData.h"
#include "vtkTextureObject.h"
#include <cassert>

static void vtkGetDimensions(int extents[6], int dims[3])
{
  dims[0] = extents[1] - extents[0] + 1;
  dims[1] = extents[3] - extents[2] + 1;
  dims[2] = extents[5] - extents[4] + 1;
}

vtkStandardNewMacro(vtkDataTransferHelper);
vtkCxxSetObjectMacro(vtkDataTransferHelper, Texture, vtkTextureObject);
vtkCxxSetObjectMacro(vtkDataTransferHelper, Array, vtkDataArray);
//------------------------------------------------------------------------------
vtkDataTransferHelper::vtkDataTransferHelper()
{
  this->Texture = nullptr;
  this->Context = nullptr;
  this->Array = nullptr;
  this->ShaderSupportsTextureInt = false;
  this->GPUExtent[0] = this->GPUExtent[1] = this->GPUExtent[2] = this->GPUExtent[3] =
    this->GPUExtent[4] = this->GPUExtent[5] = 0;
  this->CPUExtent[0] = this->CPUExtent[1] = this->CPUExtent[2] = this->CPUExtent[3] =
    this->CPUExtent[4] = this->CPUExtent[5] = 0;

  // invalid extent.
  this->TextureExtent[0] = 0;
  this->TextureExtent[1] = -1;
  this->TextureExtent[2] = 0;
  this->TextureExtent[3] = -1;
  this->TextureExtent[4] = 0;
  this->TextureExtent[5] = -1;

  this->MinTextureDimension = 1;
}

//------------------------------------------------------------------------------
vtkDataTransferHelper::~vtkDataTransferHelper()
{
  this->SetTexture(nullptr);
  this->SetArray(nullptr);
  this->SetContext(nullptr);
}

//------------------------------------------------------------------------------
// Description:
// Tells if the given extent (6 int) is valid. True if min
// extent<=max extent.
// \pre extent_exists: extent!=0
bool vtkDataTransferHelper::GetExtentIsValid(int* extent)
{
  assert("pre extent_exists:" && extent != nullptr);
  return extent[0] <= extent[1] && extent[2] <= extent[3] && extent[4] <= extent[5];
}

//------------------------------------------------------------------------------
// Description:
// Tells if CPUExtent is valid. True if min extent<=max extent.
bool vtkDataTransferHelper::GetCPUExtentIsValid()
{
  return this->GetExtentIsValid(this->CPUExtent);
}

//------------------------------------------------------------------------------
// Description:
// Tells if GPUExtent is valid. True if min extent<=max extent.
bool vtkDataTransferHelper::GetGPUExtentIsValid()
{
  return this->GetExtentIsValid(this->GPUExtent);
}

//------------------------------------------------------------------------------
// Description:
// Tells if TextureExtent is valid. True if min extent<=max extent.
bool vtkDataTransferHelper::GetTextureExtentIsValid()
{
  return this->GetExtentIsValid(this->TextureExtent);
}

//------------------------------------------------------------------------------
// Description:
// Returns if the context supports the required extensions.
bool vtkDataTransferHelper::IsSupported(vtkRenderWindow* renWin)
{
  vtkOpenGLRenderWindow* rw = vtkOpenGLRenderWindow::SafeDownCast(renWin);
  return (vtkPixelBufferObject::IsSupported(renWin) && vtkTextureObject::IsSupported(rw));
}

//------------------------------------------------------------------------------
vtkRenderWindow* vtkDataTransferHelper::GetContext()
{
  return this->Context;
}

//------------------------------------------------------------------------------
void vtkDataTransferHelper::SetContext(vtkRenderWindow* renWin)
{
  if (renWin == this->Context)
  {
    return;
  }

  if (this->Texture && this->Texture->GetContext() != renWin)
  {
    this->SetTexture(nullptr);
  }

  vtkOpenGLRenderWindow* openGLRenWin = vtkOpenGLRenderWindow::SafeDownCast(renWin);
  this->Context = openGLRenWin;
  // release the old PBO.
  this->PBO = nullptr;
  this->Modified();
}

//------------------------------------------------------------------------------
// Description:
// Upload GPUExtent from CPU vtkDataArray to GPU texture.
// It is possible to send a subset of the components or to specify and
// order of components or both. If components=0, componentList is ignored
// and all components are passed, a texture cannot have more than 4
// components.
// \pre array_exists: array!=0
// \pre array_not_empty: array->GetNumberOfTuples()>0
// \pre valid_cpu_extent: this->GetCPUExtentIsValid()
// \pre valid_cpu_extent_size:
// (CPUExtent[1]-CPUExtent[0]+1)*(CPUExtent[3]-CPUExtent[2]+1)*(CPUExtent[5]-CPUExtent[4]+1)==array->GetNumberOfTuples()
// \pre valid_gpu_extent: this->GetGPUExtentIsValid()
// \pre gpu_extent_in_cpu_extent: CPUExtent[0]<=GPUExtent[0] && GPUExtent[1]<=CPUExtent[1] &&
// CPUExtent[2]<=GPUExtent[2] && GPUExtent[3]<=CPUExtent[3] && CPUExtent[4]<=GPUExtent[4] &&
// GPUExtent[5]<=CPUExtent[5] \pre gpu_texture_size: !this->GetTextureExtentIsValid() ||
// (GPUExtent[1]-GPUExtent[0]+1)*(GPUExtent[3]-GPUExtent[2]+1)*(GPUExtent[5]-GPUExtent[4]+1)==(TextureExtent[1]-TextureExtent[0]+1)*(TextureExtent[3]-TextureExtent[2]+1)*(TextureExtent[5]-TextureExtent[4]+1)
// \pre valid_components: (components==0 && componentList==0 && array->GetNumberOfComponents()<=4)
// || (components>=1 && components<=array->GetNumberOfComponents() && components<=4 &&
// componentList!=0)
bool vtkDataTransferHelper::Upload(int components, int* componentList)
{
  assert("pre: array_exists" && this->Array != nullptr);
  assert("pre: array_not_empty" && this->Array->GetNumberOfTuples() > 0);
  assert("pre: valid_cpu_extent" && this->GetCPUExtentIsValid());
  assert("pre: valid_cpu_extent_size" &&
    ((this->CPUExtent[1] - this->CPUExtent[0] + 1) * (this->CPUExtent[3] - this->CPUExtent[2] + 1) *
        (this->CPUExtent[5] - this->CPUExtent[4] + 1) ==
      this->Array->GetNumberOfTuples()));
  assert("pre: valid_gpu_extent" && this->GetGPUExtentIsValid());
  assert("pre: gpu_extent_in_cpu_extent" &&
    (this->CPUExtent[0] <= this->GPUExtent[0] && this->GPUExtent[1] <= this->CPUExtent[1] &&
      this->CPUExtent[2] <= this->GPUExtent[2] && this->GPUExtent[3] <= this->CPUExtent[3] &&
      this->CPUExtent[4] <= this->GPUExtent[4] && this->GPUExtent[5] <= this->CPUExtent[5]));
  assert("pre: gpu_texture_size" &&
    (!this->GetTextureExtentIsValid() ||
      ((this->GPUExtent[1] - this->GPUExtent[0] + 1) *
          (this->GPUExtent[3] - this->GPUExtent[2] + 1) *
          (this->GPUExtent[5] - this->GPUExtent[4] + 1) ==
        (this->TextureExtent[1] - this->TextureExtent[0] + 1) *
          (this->TextureExtent[3] - this->TextureExtent[2] + 1) *
          (this->TextureExtent[5] - this->TextureExtent[4] + 1))));
  assert("pre: valid_components" &&
    ((components == 0 && componentList == nullptr && this->Array->GetNumberOfComponents() <= 4) ||
      (components >= 1 && components <= this->Array->GetNumberOfComponents() && components <= 4 &&
        componentList != nullptr)));

  if (!this->Context)
  {
    vtkErrorMacro("Cannot upload to GPU without context.");
    return false;
  }

  int cpudims[3];
  int gpudims[3];
  int texturedims[3];
  vtkGetDimensions(this->CPUExtent, cpudims);
  vtkGetDimensions(this->GPUExtent, gpudims);
  vtkDebugMacro(<< "CPUDims: " << cpudims[0] << ", " << cpudims[1] << ", " << cpudims[2] << endl);
  vtkDebugMacro(<< "GPUDims: " << gpudims[0] << ", " << gpudims[1] << ", " << gpudims[2] << endl);
  if (!this->GetTextureExtentIsValid())
  {
    // use GPU extent.
    texturedims[0] = gpudims[0];
    texturedims[1] = gpudims[1];
    texturedims[2] = gpudims[2];
  }
  else
  {
    vtkGetDimensions(this->TextureExtent, texturedims);
  }

  int numComps = this->Array->GetNumberOfComponents();

  // * Upload data to vtkPixelBufferObject.
  vtkSmartPointer<vtkPixelBufferObject> pbo = this->GetPBO();

  // We need to get the ContinuousIncrements as computed by the vtkImageData.
  // For that we create a dummy image data object,
  vtkSmartPointer<vtkImageData> tempImg = vtkSmartPointer<vtkImageData>::New();
  tempImg->SetDimensions(1, 1, 1);
  // scalars are needed for ComputeIncrements().
  tempImg->AllocateScalars(VTK_DOUBLE, 1);
  tempImg->SetExtent(this->CPUExtent);

  vtkIdType continuousInc[3];
  tempImg->GetContinuousIncrements(
    this->GPUExtent, continuousInc[0], continuousInc[1], continuousInc[2]);
  tempImg = nullptr;

  int pt[3] = { this->GPUExtent[0] - this->CPUExtent[0], this->GPUExtent[2] - this->CPUExtent[2],
    this->GPUExtent[4] - this->CPUExtent[4] };

  vtkIdType ptId = vtkStructuredData::ComputePointId(cpudims, pt);
  if (!pbo->Upload3D(this->Array->GetDataType(), this->Array->GetVoidPointer(ptId * numComps),
        reinterpret_cast<unsigned int*>(gpudims), numComps, continuousInc, components,
        componentList))
  {
    vtkErrorMacro("Failed to load data to pixel buffer.");
    return false;
  }

  // * Now, we need to create a Texture for the uploaded data.
  if (!this->Texture)
  {
    vtkTextureObject* tex = vtkTextureObject::New();
    tex->SetContext(vtkOpenGLRenderWindow::SafeDownCast(this->Context));
    this->SetTexture(tex);
    tex->Delete();
  }

  int tempdims[3] = { 0, 0, 0 };
  int dataDescription = vtkStructuredData::SetDimensions(texturedims, tempdims);
  int dimension = vtkStructuredData::GetDataDimension(dataDescription);

  bool uploaded = false;
  switch (dimension)
  {
    case 0: // 1 pixel image
    case 1:
    {
      unsigned int length = 0;
      switch (dataDescription)
      {
        case VTK_SINGLE_POINT:
          length = 1;
          break;
        case VTK_X_LINE:
          length = static_cast<unsigned int>(texturedims[0]);
          break;
        case VTK_Y_LINE:
          length = static_cast<unsigned int>(texturedims[1]);
          break;
        case VTK_Z_LINE:
          length = static_cast<unsigned int>(texturedims[2]);
          break;
      }
      switch (this->MinTextureDimension)
      {
        case 1:
          uploaded = this->Texture->Create1D(numComps, pbo, this->ShaderSupportsTextureInt);
          break;
        case 2:
          uploaded =
            this->Texture->Create2D(length, 1, numComps, pbo, this->ShaderSupportsTextureInt);
          break;
        case 3:
          uploaded =
            this->Texture->Create3D(length, 1, 1, numComps, pbo, this->ShaderSupportsTextureInt);
          break;
        default:
          assert("check: impossible case" && 0);
          break;
      }
    }
    break;

    case 2:
    {
      unsigned int width = 0;
      unsigned int height = 0;
#if 0
        switch (dataDescription)
        {
          case VTK_XY_PLANE:
            width = gpudims[0];
            height = gpudims[1];
            break;

          case VTK_YZ_PLANE:
            width = gpudims[1];
            height = gpudims[2];
            break;

          case VTK_XZ_PLANE:
            width = gpudims[0];
            height = gpudims[2];
            break;
        }
#else
#if 1
      switch (dataDescription)
      {
        case VTK_XY_PLANE:
          width = static_cast<unsigned int>(texturedims[0]);
          height = static_cast<unsigned int>(texturedims[1]);
          break;

        case VTK_YZ_PLANE:
          width = static_cast<unsigned int>(texturedims[1]);
          height = static_cast<unsigned int>(texturedims[2]);
          break;

        case VTK_XZ_PLANE:
          width = static_cast<unsigned int>(texturedims[0]);
          height = static_cast<unsigned int>(texturedims[2]);
          break;
      }
#else
      width = texturedims[0];
      height = texturedims[1];
#endif
#endif
      switch (this->MinTextureDimension)
      {
        case 1:
        case 2:
          uploaded =
            this->Texture->Create2D(width, height, numComps, pbo, this->ShaderSupportsTextureInt);
          break;
        case 3:
          uploaded = this->Texture->Create3D(
            width, height, 1, numComps, pbo, this->ShaderSupportsTextureInt);
          break;
        default:
          assert("check: impossible case" && 0);
          break;
      }
    }
    break;

    case 3:
      uploaded = this->Texture->Create3D(static_cast<unsigned int>(texturedims[0]),
        static_cast<unsigned int>(texturedims[1]), static_cast<unsigned int>(texturedims[2]),
        numComps, pbo, this->ShaderSupportsTextureInt);
      break;
  }

  pbo->ReleaseMemory();
  if (!uploaded)
  {
    vtkErrorMacro("Failed to upload data to texture.");
    return false;
  }
  return true;
}

//------------------------------------------------------------------------------
// Description:
// new comment:
// Download GPUExtent from GPU texture to CPU vtkDataArray.
// If Array is not provided, it will be created with the size of CPUExtent.
// But only the tuples covered by GPUExtent will be download. In this case,
// if GPUExtent does not cover all GPUExtent, some of the vtkDataArray will
// be uninitialized.
// Reminder: A=>B <=> !A||B
// \pre texture_exists: texture!=0
// \pre array_not_empty: array==0 || array->GetNumberOfTuples()>0
// \pre valid_cpu_extent: this->GetCPUExtentIsValid()
// \pre valid_cpu_extent_size: array==0 ||
// (CPUExtent[1]-CPUExtent[0]+1)*(CPUExtent[3]-CPUExtent[2]+1)*(CPUExtent[5]-CPUExtent[4]+1)==array->GetNumberOfTuples()
// \pre valid_gpu_extent: this->GetGPUExtentIsValid()
// \pre gpu_extent_in_cpu_extent: CPUExtent[0]<=GPUExtent[0] && GPUExtent[1]<=CPUExtent[1] &&
// CPUExtent[2]<=GPUExtent[2] && GPUExtent[3]<=CPUExtent[3] && CPUExtent[4]<=GPUExtent[4] &&
// GPUExtent[5]<=CPUExtent[5] \pre gpu_texture_size: !this->GetTextureExtentIsValid() ||
// (GPUExtent[1]-GPUExtent[0]+1)*(GPUExtent[3]-GPUExtent[2]+1)*(GPUExtent[5]-GPUExtent[4]+1)==(TextureExtent[1]-TextureExtent[0]+1)*(TextureExtent[3]-TextureExtent[2]+1)*(TextureExtent[5]-TextureExtent[4]+1)
// \pre valid_components: array==0 || array->GetNumberOfComponents()<=4
// \pre components_match: array==0 || (texture->GetComponents()==array->GetNumberOfComponents())
bool vtkDataTransferHelper::Download()
{
  return (this->DownloadAsync1() && this->DownloadAsync2());
}

//------------------------------------------------------------------------------
bool vtkDataTransferHelper::DownloadAsync1()
{
  if (!this->Context)
  {
    vtkErrorMacro("Cannot download from GPU without context.");
    return false;
  }

  assert("pre: texture_exists" && this->Texture != nullptr);
  assert(
    "pre: array_not_empty" && (this->Array == nullptr || this->Array->GetNumberOfTuples() > 0));
  assert("pre: valid_cpu_extent" && this->GetCPUExtentIsValid());
  assert("pre: valid_cpu_extent_size" &&
    (this->Array == nullptr ||
      ((this->CPUExtent[1] - this->CPUExtent[0] + 1) *
          (this->CPUExtent[3] - this->CPUExtent[2] + 1) *
          (this->CPUExtent[5] - this->CPUExtent[4] + 1) ==
        this->Array->GetNumberOfTuples())));
  assert("pre: valid_gpu_extent" && this->GetGPUExtentIsValid());
  assert("pre: gpu_extent_in_cpu_extent" &&
    (this->CPUExtent[0] <= this->GPUExtent[0] && this->GPUExtent[1] <= this->CPUExtent[1] &&
      this->CPUExtent[2] <= this->GPUExtent[2] && this->GPUExtent[3] <= this->CPUExtent[3] &&
      this->CPUExtent[4] <= this->GPUExtent[4] && this->GPUExtent[5] <= this->CPUExtent[5]));
  assert("pre: gpu_texture_size" &&
    (!this->GetTextureExtentIsValid() ||
      ((this->GPUExtent[1] - this->GPUExtent[0] + 1) *
          (this->GPUExtent[3] - this->GPUExtent[2] + 1) *
          (this->GPUExtent[5] - this->GPUExtent[4] + 1) ==
        (this->TextureExtent[1] - this->TextureExtent[0] + 1) *
          (this->TextureExtent[3] - this->TextureExtent[2] + 1) *
          (this->TextureExtent[5] - this->TextureExtent[4] + 1))));
  assert("pre: valid_components" &&
    (this->Array == nullptr || this->Array->GetNumberOfComponents() <= 4));
  assert("pre: components_match" &&
    (this->Array == nullptr ||
      (this->Texture->GetComponents() == this->Array->GetNumberOfComponents())));

  int numComps = this->Texture->GetComponents();
  int cpudims[3];
  int gpudims[3];
  vtkGetDimensions(this->CPUExtent, cpudims);
  vtkGetDimensions(this->GPUExtent, gpudims);

  // * Download data to vtkPixelBufferObject.
  vtkSmartPointer<vtkPixelBufferObject> pbo;
  pbo.TakeReference(this->Texture->Download());

  if (!pbo)
  {
    vtkErrorMacro("Failed to download texture to a Pixel Buffer object.");
    return false;
  }

  if (pbo->GetSize() < static_cast<unsigned int>(gpudims[0] * gpudims[1] * gpudims[2] * numComps))
  {
    vtkErrorMacro("GPU data size is smaller than GPUExtent.");
    return false;
  }

  this->AsyncDownloadPBO = pbo;
  return true;
}

//------------------------------------------------------------------------------
bool vtkDataTransferHelper::DownloadAsync2()
{
  if (!this->AsyncDownloadPBO)
  {
    vtkErrorMacro("DownloadAsync1() must be called successfully "
                  "before calling DownloadAsync2().");
    return false;
  }

  int numComps = this->Texture->GetComponents();
  int cpudims[3];
  int gpudims[3];
  vtkGetDimensions(this->CPUExtent, cpudims);
  vtkGetDimensions(this->GPUExtent, gpudims);

  if (!this->Array)
  {
    vtkDataArray* array = vtkDataArray::CreateDataArray(this->Texture->GetVTKDataType());
    this->SetArray(array);
    array->Delete();
    this->Array->SetNumberOfComponents(numComps);
    this->Array->SetNumberOfTuples(cpudims[0] * cpudims[1] * cpudims[2]);
  }

  // We need to get the ContinuousIncrements as computed by the vtkImageData.
  // For that we create a dummy image data object,
  vtkSmartPointer<vtkImageData> tempImg = vtkSmartPointer<vtkImageData>::New();
  tempImg->SetDimensions(1, 1, 1);
  // scalars are needed for ComputeIncrements().
  tempImg->AllocateScalars(VTK_DOUBLE, 1);
  tempImg->SetExtent(this->CPUExtent);

  vtkIdType continuousInc[3];
  tempImg->GetContinuousIncrements(
    this->GPUExtent, continuousInc[0], continuousInc[1], continuousInc[2]);
  tempImg = nullptr;

  int pt[3] = { this->GPUExtent[0] - this->CPUExtent[0], this->GPUExtent[2] - this->CPUExtent[2],
    this->GPUExtent[4] - this->CPUExtent[4] };

  vtkIdType ptId = vtkStructuredData::ComputePointId(cpudims, pt);
  bool reply = this->AsyncDownloadPBO->Download3D(this->Array->GetDataType(),
    this->Array->GetVoidPointer(ptId * numComps), reinterpret_cast<unsigned int*>(gpudims),
    numComps, continuousInc);
  this->AsyncDownloadPBO = nullptr;
  return reply;
}

//------------------------------------------------------------------------------
bool vtkDataTransferHelper::GetShaderSupportsTextureInt()
{
  return this->ShaderSupportsTextureInt;
}

//------------------------------------------------------------------------------
void vtkDataTransferHelper::SetShaderSupportsTextureInt(bool value)
{
  this->ShaderSupportsTextureInt = value;
}

//------------------------------------------------------------------------------
vtkPixelBufferObject* vtkDataTransferHelper::GetPBO()
{
  if (!this->PBO)
  {
    this->PBO = vtkSmartPointer<vtkPixelBufferObject>::New();
    this->PBO->SetContext(this->Context);
  }

  return this->PBO;
}

//------------------------------------------------------------------------------
void vtkDataTransferHelper::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);

  os << indent << "Array: " << this->Array << endl;
  os << indent << "Texture: " << this->Texture << endl;
  os << indent << "MinTextureDimension: " << this->MinTextureDimension << endl;
  os << indent << "CPUExtent: (" << this->CPUExtent[0] << ", " << this->CPUExtent[1] << ", "
     << this->CPUExtent[2] << ", " << this->CPUExtent[3] << ", " << this->CPUExtent[4] << ", "
     << this->CPUExtent[5] << ")" << endl;
  os << indent << "GPUExtent: (" << this->GPUExtent[0] << ", " << this->GPUExtent[1] << ", "
     << this->GPUExtent[2] << ", " << this->GPUExtent[3] << ", " << this->GPUExtent[4] << ", "
     << this->GPUExtent[5] << ")" << endl;
  os << indent << "TextureExtent: (" << this->TextureExtent[0] << ", " << this->TextureExtent[1]
     << ", " << this->TextureExtent[2] << ", " << this->TextureExtent[3] << ", "
     << this->TextureExtent[4] << ", " << this->TextureExtent[5] << ")" << endl;
}