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

   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 "vtkSmartPointer.h"
#include "vtkObjectFactory.h"
#include "vtkOSPRayRendererNode.h"
#include "vtkOverlappingAMR.h"
#include "vtkUniformGridAMRDataIterator.h"
#include "vtkVolume.h"
#include "vtkVolumeMapper.h"
#include "vtkVolumeNode.h"
#include "vtkVolumeProperty.h"

#include "ospcommon/box.h"
#include "ospcommon/vec.h"
#include "ospray/ospray.h"
#include "ospray/version.h"

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;
  if (OSPRAY_VERSION_MAJOR < 1 || (OSPRAY_VERSION_MAJOR == 1 && OSPRAY_VERSION_MINOR < 4))
  {
    vtkErrorMacro("vtkOSPRayAMRVolumeMapperNode Requires at least OSPRay 1.4");
  }
}

//----------------------------------------------------------------------------
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;
    }

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

    vtkOSPRayRendererNode *orn =
                    static_cast<vtkOSPRayRendererNode *>(
                    this->GetFirstAncestorOfType("vtkOSPRayRendererNode"));

    osp::Model* 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)
    {
      this->OSPRayVolume = ospNewVolume("amr_volume");
      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();
        assert(level >= lastLevel);  //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]);

      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(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
}