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

   Program:   Visualization Toolkit
   Module:    vtkOSPRayAMRVolumeMapperNode.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 "vtkOSPRayAMRVolumeMapperNode.h"

#include "vtkAMRBox.h"
#include "vtkAMRInformation.h"
#include "vtkDataArray.h"
#include "vtkImageData.h"
#include "vtkOSPRayCache.h"
#include "vtkOSPRayRendererNode.h"
#include "vtkObjectFactory.h"
#include "vtkOverlappingAMR.h"
#include "vtkRenderer.h"
#include "vtkSmartPointer.h"
#include "vtkUniformGridAMRDataIterator.h"
#include "vtkVolume.h"
#include "vtkVolumeMapper.h"
#include "vtkVolumeNode.h"
#include "vtkVolumeProperty.h"

#include "RTWrapper/RTWrapper.h"

#include <cassert>

namespace ospray
{
namespace amr
{
struct BrickInfo
{
  /*! bounding box of integer coordinates of cells. note that
      this EXCLUDES the width of the rightmost cell: ie, a 4^3
      box at root level pos (0,0,0) would have a _box_ of
      [(0,0,0)-(3,3,3)] (because 3,3,3 is the highest valid
      coordinate in this box!), while its bounds would be
      [(0,0,0)-(4,4,4)]. Make sure to NOT use box.size() for the
      grid dimensions, since this will always be one lower than
      the dims of the grid.... */
  ospcommon::box3i box;
  //! level this brick is at
  int level;
  // width of each cell in this level
  float cellWidth;

  // inline ospcommon::box3f worldBounds() const {
  //  return ospcommon::box3f(ospcommon::vec3f(box.lower)*cellWidth,
  //        ospcommon::vec3f(box.upper+ospcommon::vec3i(1))*cellWidth);
  //}
};
}
}

using namespace ospray::amr;

vtkStandardNewMacro(vtkOSPRayAMRVolumeMapperNode);

//------------------------------------------------------------------------------
vtkOSPRayAMRVolumeMapperNode::vtkOSPRayAMRVolumeMapperNode()
{
  this->NumColors = 128;
  this->TransferFunction = nullptr;
  this->SamplingRate = 0.5f;
  this->OldSamplingRate = -1.f;
}

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

//------------------------------------------------------------------------------
void vtkOSPRayAMRVolumeMapperNode::Render(bool prepass)
{
  if (prepass)
  {
    vtkVolumeNode* volNode = vtkVolumeNode::SafeDownCast(this->Parent);
    if (!volNode)
    {
      vtkErrorMacro("invalid volumeNode");
      return;
    }
    vtkVolume* vol = vtkVolume::SafeDownCast(volNode->GetRenderable());
    if (vol->GetVisibility() == false)
    {
      return;
    }
    vtkVolumeMapper* mapper = vtkVolumeMapper::SafeDownCast(this->GetRenderable());
    if (!mapper)
    {
      vtkErrorMacro("invalid mapper");
      return;
    }
    if (!vol->GetProperty())
    {
      vtkErrorMacro("VolumeMapper had no vtkProperty");
      return;
    }

    vtkOSPRayRendererNode* orn =
      static_cast<vtkOSPRayRendererNode*>(this->GetFirstAncestorOfType("vtkOSPRayRendererNode"));
    vtkRenderer* ren = vtkRenderer::SafeDownCast(orn->GetRenderable());

    RTW::Backend* backend = orn->GetBackend();
    if (backend == nullptr)
      return;

    if (!this->TransferFunction)
    {
      this->TransferFunction = ospNewTransferFunction("piecewise_linear");
    }

    this->Cache->SetSize(vtkOSPRayRendererNode::GetTimeCacheSize(ren));

    OSPModel OSPRayModel = orn->GetOModel();
    if (!OSPRayModel)
    {
      return;
    }

    vtkOverlappingAMR* amr = vtkOverlappingAMR::SafeDownCast(mapper->GetInputDataObject(0, 0));
    if (!amr)
    {
      vtkErrorMacro("couldn't get amr data\n");
      return;
    }

    vtkVolumeProperty* volProperty = vol->GetProperty();

    bool volDirty = false;
    if (!this->OSPRayVolume || amr->GetMTime() > this->BuildTime)
    {
      double tstep = vtkOSPRayRendererNode::GetViewTime(ren);
      auto cached_Volume = this->Cache->Get(tstep);
      if (cached_Volume)
      {
        this->OSPRayVolume = static_cast<OSPVolume>(cached_Volume->object);
      }
      else
      {
        if (this->Cache->GetSize() == 0)
        {
          ospRelease(this->OSPRayVolume);
        }
        this->OSPRayVolume = ospNewVolume("amr_volume");
        if (this->Cache->HasRoom())
        {
          auto cacheEntry = std::make_shared<vtkOSPRayCacheItemObject>(backend, this->OSPRayVolume);
          this->Cache->Set(tstep, cacheEntry);
        }
        volDirty = true;

        unsigned int lastLevel = 0;
        std::vector<OSPData> brickDataArray;
        std::vector<BrickInfo> brickInfoArray;
        size_t totalDataSize = 0;

        vtkAMRInformation* amrInfo = amr->GetAMRInfo();
        vtkSmartPointer<vtkUniformGridAMRDataIterator> iter;
        iter.TakeReference(vtkUniformGridAMRDataIterator::SafeDownCast(amr->NewIterator()));
        for (iter->InitTraversal(); !iter->IsDoneWithTraversal(); iter->GoToNextItem())
        {
          unsigned int level = iter->GetCurrentLevel();
          if (!(level >= lastLevel))
          {
            vtkErrorMacro("ospray requires level info be ordered lowest to highest");
          };
          lastLevel = level;
          unsigned int index = iter->GetCurrentIndex();

          // note: this iteration "naturally" goes from datasets at lower levels to
          // those at higher levels.
          vtkImageData* data = vtkImageData::SafeDownCast(iter->GetCurrentDataObject());
          if (!data)
          {
            return;
          }
          float* dataPtr;
          int dim[3];

          const vtkAMRBox& box = amrInfo->GetAMRBox(level, index);
          const int* lo = box.GetLoCorner();
          const int* hi = box.GetHiCorner();
          ospcommon::vec3i lo_v = { lo[0], lo[1], lo[2] };
          ospcommon::vec3i hi_v = { hi[0], hi[1], hi[2] };
          dim[0] = hi[0] - lo[0] + 1;
          dim[1] = hi[1] - lo[1] + 1;
          dim[2] = hi[2] - lo[2] + 1;

          int fieldAssociation;
          mapper->SetScalarMode(VTK_SCALAR_MODE_USE_CELL_FIELD_DATA);
          vtkDataArray* cellArray =
            vtkDataArray::SafeDownCast(this->GetArrayToProcess(data, fieldAssociation));
          if (!cellArray)
          {
            std::cerr << "could not get data!\n";
            return;
          }

          if (cellArray->GetDataType() != VTK_FLOAT)
          {
            if (cellArray->GetDataType() == VTK_DOUBLE)
            {
              float* fdata = new float[dim[0] * dim[1] * size_t(dim[2])];
              double* dptr;
              dptr = (double*)cellArray->WriteVoidPointer(0, cellArray->GetSize());
              for (size_t i = 0; i < dim[0] * dim[1] * size_t(dim[2]); i++)
              {
                fdata[i] = dptr[i];
              }
              dataPtr = fdata;
            }
            else
            {
              std::cerr
                << "Only doubles and floats are supported in OSPRay AMR volume mapper currently";
              return;
            }
          }
          else
          {
            dataPtr = (float*)cellArray->WriteVoidPointer(0, cellArray->GetSize());
          }

          totalDataSize += sizeof(float) * size_t(dim[0] * dim[1]) * dim[2];
          OSPData odata =
            ospNewData(dim[0] * dim[1] * dim[2], OSP_FLOAT, dataPtr, OSP_DATA_SHARED_BUFFER);
          brickDataArray.push_back(odata);

          BrickInfo bi;
          ospcommon::box3i obox = { lo_v, hi_v };
          bi.box = obox;
          double spacing[3];
          amrInfo->GetSpacing(level, spacing);
          bi.cellWidth = spacing[0];
          // cell bounds:  origin + box.LoCorner*spacing,
          bi.level = level;
          totalDataSize += sizeof(BrickInfo);

          brickInfoArray.push_back(bi);
        }

        assert(brickDataArray.size() == brickInfoArray.size());
        ospSet1f(this->OSPRayVolume, "samplingRate", this->SamplingRate); // TODO: gui option

        double origin[3];
        vol->GetOrigin(origin);
        double* bds = mapper->GetBounds();
        origin[0] = bds[0];
        origin[1] = bds[2];
        origin[2] = bds[4];

        double spacing[3];
        amr->GetAMRInfo()->GetSpacing(0, spacing);
        ospSet3f(this->OSPRayVolume, "gridOrigin", origin[0], origin[1], origin[2]);
        ospSetString(this->OSPRayVolume, "voxelType", "float");

        OSPData brickDataData =
          ospNewData(brickDataArray.size(), OSP_OBJECT, &brickDataArray[0], 0);
        ospSetData(this->OSPRayVolume, "brickData", brickDataData);
        OSPData brickInfoData = ospNewData(
          brickInfoArray.size() * sizeof(brickInfoArray[0]), OSP_RAW, &brickInfoArray[0], 0);
        ospSetData(this->OSPRayVolume, "brickInfo", brickInfoData);
        ospSetObject(this->OSPRayVolume, "transferFunction", this->TransferFunction);
        this->BuildTime.Modified();
      }
    }
    if ((vol->GetProperty()->GetMTime() > this->PropertyTime) || volDirty)
    {
      //
      // transfer function
      //
      this->UpdateTransferFunction(backend, vol);
      ospSet1i(this->OSPRayVolume, "gradientShadingEnabled", volProperty->GetShade());
      this->PropertyTime.Modified();
    }

    if (this->OldSamplingRate != this->SamplingRate)
    {
      this->OldSamplingRate = this->SamplingRate;
      volDirty = true;
    }

    if (volDirty)
    {
      ospSet1f(this->OSPRayVolume, "samplingRate", this->SamplingRate);
      ospCommit(this->OSPRayVolume);
    }
    ospAddVolume(OSPRayModel, this->OSPRayVolume);
    ospCommit(OSPRayModel);
  } // prepass
}
