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

  Program:   ParaView
  Module:    vtkImageVolumeRepresentation.cxx

  Copyright (c) Kitware, Inc.
  All rights reserved.
  See Copyright.txt or http://www.paraview.org/HTML/Copyright.html 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 "vtkImageVolumeRepresentation.h"

#include "vtkAlgorithmOutput.h"
#include "vtkCellData.h"
#include "vtkColorTransferFunction.h"
#include "vtkCommand.h"
#include "vtkContourValues.h"
#include "vtkDataSet.h"
#include "vtkImageData.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkMath.h"
#include "vtkMultiBlockVolumeMapper.h"
#include "vtkNew.h"
#include "vtkObjectFactory.h"
#include "vtkPVLODVolume.h"
#include "vtkPVRenderView.h"
#include "vtkPartitionedDataSet.h"
#include "vtkPolyDataMapper.h"
#include "vtkRectilinearGrid.h"
#include "vtkRenderer.h"
#include "vtkSMPTools.h"
#include "vtkSmartPointer.h"
#include "vtkSmartVolumeMapper.h"
#include "vtkStreamingDemandDrivenPipeline.h"
#include "vtkStructuredData.h"
#include "vtkUniformGrid.h"
#include "vtkUnsignedCharArray.h"
#include "vtkVolumeProperty.h"

#include <algorithm>
#include <map>
#include <string>

class vtkImageVolumeRepresentation::vtkRenderingInternals
{
public:
  vtkNew<vtkMultiBlockVolumeMapper> VolumeMapper;

  vtkSmartPointer<vtkDataObject> Cache;

  double CroppingOrigin[3] = { 0, 0, 0 };
  double CroppingScale[3] = { 1, 1, 1 };
};

namespace
{
//----------------------------------------------------------------------------
void vtkGetNonGhostExtent(int* resultExtent, vtkImageData* dataSet)
{
  // this is really only meant for topologically structured grids
  dataSet->GetExtent(resultExtent);

  if (vtkUnsignedCharArray* ghostArray = vtkUnsignedCharArray::SafeDownCast(
        dataSet->GetCellData()->GetArray(vtkDataSetAttributes::GhostArrayName())))
  {
    // We have a ghost array. We need to iterate over the array to prune ghost
    // extents.

    int pntExtent[6];
    std::copy(resultExtent, resultExtent + 6, pntExtent);

    int validCellExtent[6];
    vtkStructuredData::GetCellExtentFromPointExtent(pntExtent, validCellExtent);

    // The start extent is the location of the first cell with ghost value 0.
    for (vtkIdType cc = 0, numTuples = ghostArray->GetNumberOfTuples(); cc < numTuples; ++cc)
    {
      if (ghostArray->GetValue(cc) == 0)
      {
        int ijk[3];
        vtkStructuredData::ComputeCellStructuredCoordsForExtent(cc, pntExtent, ijk);
        validCellExtent[0] = ijk[0];
        validCellExtent[2] = ijk[1];
        validCellExtent[4] = ijk[2];
        break;
      }
    }

    // The end extent is the  location of the last cell with ghost value 0.
    for (vtkIdType cc = (ghostArray->GetNumberOfTuples() - 1); cc >= 0; --cc)
    {
      if (ghostArray->GetValue(cc) == 0)
      {
        int ijk[3];
        vtkStructuredData::ComputeCellStructuredCoordsForExtent(cc, pntExtent, ijk);
        validCellExtent[1] = ijk[0];
        validCellExtent[3] = ijk[1];
        validCellExtent[5] = ijk[2];
        break;
      }
    }

    // convert cell-extents to pt extents.
    resultExtent[0] = validCellExtent[0];
    resultExtent[2] = validCellExtent[2];
    resultExtent[4] = validCellExtent[4];

    resultExtent[1] = std::min(validCellExtent[1] + 1, resultExtent[1]);
    resultExtent[3] = std::min(validCellExtent[3] + 1, resultExtent[3]);
    resultExtent[5] = std::min(validCellExtent[5] + 1, resultExtent[5]);
  }
}
}

vtkStandardNewMacro(vtkImageVolumeRepresentation);
//----------------------------------------------------------------------------
vtkImageVolumeRepresentation::vtkImageVolumeRepresentation() = default;

//----------------------------------------------------------------------------
vtkImageVolumeRepresentation::~vtkImageVolumeRepresentation() = default;

//----------------------------------------------------------------------------
void vtkImageVolumeRepresentation::InitializeForDataProcessing()
{
  this->Superclass::InitializeForDataProcessing();
}

//----------------------------------------------------------------------------
void vtkImageVolumeRepresentation::InitializeForRendering()
{
  this->Superclass::InitializeForRendering();
  this->RInternals.reset(new vtkImageVolumeRepresentation::vtkRenderingInternals());

  auto& rinternals = *(this->RInternals);
  this->GetActor()->SetMapper(rinternals.VolumeMapper);
  this->SetVisibility(this->GetVisibility());
}

//----------------------------------------------------------------------------
bool vtkImageVolumeRepresentation::AddToView(vtkPVView* view)
{
  vtkPVRenderView* rview = vtkPVRenderView::SafeDownCast(view);
  if (rview && this->RInternals)
  {
    rview->GetRenderer()->AddActor(this->GetActor());
  }

  return this->Superclass::AddToView(view);
}

//----------------------------------------------------------------------------
bool vtkImageVolumeRepresentation::RemoveFromView(vtkPVView* view)
{
  vtkPVRenderView* rview = vtkPVRenderView::SafeDownCast(view);
  if (rview && this->RInternals)
  {
    rview->GetRenderer()->RemoveActor(this->GetActor());
  }

  return this->Superclass::RemoveFromView(view);
}

//----------------------------------------------------------------------------
int vtkImageVolumeRepresentation::FillInputPortInformation(int, vtkInformation* info)
{
  info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkImageData");
  info->Append(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkRectilinearGrid");
  info->Append(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkPartitionedDataSet");
  if (this->IsInitializedForRendering())
  {
    info->Set(vtkAlgorithm::INPUT_IS_OPTIONAL(), 1);
  }
  return 1;
}

//----------------------------------------------------------------------------
void vtkImageVolumeRepresentation::UpdateMapperParameters()
{
  if (!this->RInternals)
  {
    return;
  }

  auto& rinternals = *(this->RInternals);

  const char* colorArrayName = nullptr;
  int fieldAssociation = vtkDataObject::FIELD_ASSOCIATION_POINTS;

  vtkInformation* info = this->GetInputArrayInformation(0);
  if (info && info->Has(vtkDataObject::FIELD_ASSOCIATION()) &&
    info->Has(vtkDataObject::FIELD_NAME()))
  {
    colorArrayName = info->Get(vtkDataObject::FIELD_NAME());
    fieldAssociation = info->Get(vtkDataObject::FIELD_ASSOCIATION());
  }

  if (this->GetUseSeparateOpacityArray())
  {
    // See AppendOpacityComponent() for the construction of this array.
    std::string combinedName(colorArrayName);
    combinedName += "_and_opacity";
    rinternals.VolumeMapper->SelectScalarArray(combinedName.c_str());
  }
  else
  {
    rinternals.VolumeMapper->SelectScalarArray(colorArrayName);
  }
  switch (fieldAssociation)
  {
    case vtkDataObject::FIELD_ASSOCIATION_CELLS:
      rinternals.VolumeMapper->SetScalarMode(VTK_SCALAR_MODE_USE_CELL_FIELD_DATA);
      break;

    case vtkDataObject::FIELD_ASSOCIATION_NONE:
      rinternals.VolumeMapper->SetScalarMode(VTK_SCALAR_MODE_USE_FIELD_DATA);
      break;

    case vtkDataObject::FIELD_ASSOCIATION_POINTS:
    default:
      rinternals.VolumeMapper->SetScalarMode(VTK_SCALAR_MODE_USE_POINT_FIELD_DATA);
      break;
  }

  // this is necessary since volume mappers don't like empty arrays.
  this->SetVisibility(colorArrayName != nullptr && colorArrayName[0] != 0);

  if (rinternals.VolumeMapper->GetCropping())
  {
    double planes[6];
    for (int i = 0; i < 6; i++)
    {
      planes[i] = rinternals.CroppingOrigin[i / 2] +
        this->GetDataBounds()[i] * rinternals.CroppingScale[i / 2];
    }
    rinternals.VolumeMapper->SetCroppingRegionPlanes(planes);
  }

  if (this->GetActor()->GetMapper() != rinternals.VolumeMapper)
  {
    this->GetActor()->SetMapper(rinternals.VolumeMapper);
  }

  if (this->GetProperty())
  {
    if (this->GetMapScalars())
    {
      if (this->GetMultiComponentsMapping() || this->GetUseSeparateOpacityArray())
      {
        this->GetProperty()->SetIndependentComponents(false);
      }
      else
      {
        this->GetProperty()->SetIndependentComponents(true);
      }
    }
    else
    {
      this->GetProperty()->SetIndependentComponents(false);
    }

    // Update the mapper's vector mode
    vtkColorTransferFunction* ctf = this->GetProperty()->GetRGBTransferFunction(0);

    // When vtkScalarsToColors::MAGNITUDE mode is active, vtkSmartVolumeMapper
    // uses an internally generated (single-component) dataset.  However,
    // unchecking MapScalars (e.g. IndependentComponents == 0) requires 2C or 4C
    // data. In that case, vtkScalarsToColors::COMPONENT is forced in order to
    // make vtkSmartVolumeMapper use the original multiple-component dataset.
    int const indep = this->GetProperty()->GetIndependentComponents();
    int const mode = indep ? ctf->GetVectorMode() : vtkScalarsToColors::COMPONENT;
    int const comp = indep ? ctf->GetVectorComponent() : 0;

    if (auto smartVolumeMapper = vtkSmartVolumeMapper::SafeDownCast(rinternals.VolumeMapper))
    {
      smartVolumeMapper->SetVectorMode(mode);
      smartVolumeMapper->SetVectorComponent(comp);
    }
    else if (auto mbMapper = vtkMultiBlockVolumeMapper::SafeDownCast(rinternals.VolumeMapper))
    {
      mbMapper->SetVectorMode(mode);
      mbMapper->SetVectorComponent(comp);
    }
  }
}

//----------------------------------------------------------------------------
bool vtkImageVolumeRepresentation::RequestDataToRender(
  vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  auto* input = vtkDataObject::GetData(inputVector[0], 0);
  auto* outInfo = outputVector->GetInformationObject(0);
  outInfo->Set(vtkDataObject::DATA_OBJECT(), input);
  // if the input was aborted inform the view downstream so we can skip rendering
  auto* inInfo = inputVector[0]->GetInformationObject(0);
  outInfo->Set(vtkAlgorithm::ABORTED(),
    inInfo->Get(vtkAlgorithm::ABORTED()) | outInfo->Get(vtkAlgorithm::ABORTED()));
  return true;
}

//----------------------------------------------------------------------------
bool vtkImageVolumeRepresentation::RequestPrepareForRender(
  vtkInformation* request, vtkInformationVector* inputVector)
{
  assert(this->RInternals);
  auto& rinternals = (*this->RInternals);
  vtkMath::UninitializeBounds(this->GetDataBounds());
  this->SetDataSize(0);

  if (inputVector->GetNumberOfInformationObjects() == 1)
  {
    if (auto inputID = vtkImageData::GetData(inputVector, 0))
    {
      vtkSmartPointer<vtkImageData> cache = nullptr;
      if (auto inputUG = vtkUniformGrid::GetData(inputVector, 0))
      {
        cache = vtkSmartPointer<vtkUniformGrid>::New();
        cache->ShallowCopy(inputUG);

        if (this->GetUseSeparateOpacityArray())
        {
          this->AppendOpacityComponent(cache);
        }
      }
      else
      {
        cache = vtkSmartPointer<vtkImageData>::New();
        cache->ShallowCopy(inputID);

        if (this->GetUseSeparateOpacityArray())
        {
          this->AppendOpacityComponent(cache);
        }
      }
      if (inputID->HasAnyGhostCells())
      {
        int ext[6];
        vtkGetNonGhostExtent(ext, cache);
        // Yup, this will modify the "input", but that okay for now. Ultimately,
        // we will teach the volume mapper to handle ghost cells and this won't
        // be needed. Once that's done, we'll need to teach the KdTree
        // generation code to handle overlap in extents, however.
        cache->Crop(ext);
      }

      rinternals.VolumeMapper->SetInputData(cache);
      this->SetDataBounds(cache->GetBounds());
      this->UpdateMapperParameters();

      this->SetDataSize(cache->GetActualMemorySize());
      rinternals.Cache = cache.GetPointer();
    }
    else if (auto inputRD = vtkRectilinearGrid::GetData(inputVector, 0))
    {
      vtkNew<vtkRectilinearGrid> cache;
      cache->ShallowCopy(inputRD);

      if (this->GetUseSeparateOpacityArray())
      {
        this->AppendOpacityComponent(cache);
      }

      rinternals.VolumeMapper->SetInputData(cache);
      this->UpdateMapperParameters();

      this->SetDataSize(cache->GetActualMemorySize());
      rinternals.Cache = cache.GetPointer();
      this->SetDataBounds(cache->GetBounds());
    }
    else if (auto inputPD = vtkPartitionedDataSet::GetData(inputVector, 0))
    {
      if (!vtkMultiBlockVolumeMapper::SafeDownCast(rinternals.VolumeMapper))
      {
        vtkWarningMacro("Representation does not support rendering paritioned datasets yet.");
      }
      else
      {
        vtkNew<vtkPartitionedDataSet> cache;
        cache->CopyStructure(inputPD);
        for (unsigned int cc = 0; cc < inputPD->GetNumberOfPartitions(); ++cc)
        {
          auto partition = inputPD->GetPartition(cc);
          if (this->GetUseSeparateOpacityArray())
          {
            this->AppendOpacityComponent(partition);
          }
          auto partitionID = vtkImageData::SafeDownCast(partition);
          auto partitionRG = vtkRectilinearGrid::SafeDownCast(partition);
          if (partitionID)
          {
            cache->SetPartition(cc, partitionID);
          }
          else
          {
            cache->SetPartition(cc, partitionRG);
          }
        }
        rinternals.Cache = cache.GetPointer();
        double bounds[6];
        cache->GetBounds(bounds);
        this->SetDataBounds(bounds);
      }
    }
  }

  return true;
}

//----------------------------------------------------------------------------
bool vtkImageVolumeRepresentation::RequestRender(vtkInformation* vtkNotUsed(request))
{
  return true;
}

//----------------------------------------------------------------------------
void vtkImageVolumeRepresentation::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
  if (this->RInternals)
  {
    auto& rinternals = (*this->RInternals);
    os << indent << "Cropping Origin: " << rinternals.CroppingOrigin[0] << ", "
       << rinternals.CroppingOrigin[1] << ", " << rinternals.CroppingOrigin[2] << endl;
    os << indent << "Cropping Scale: " << rinternals.CroppingScale[0] << ", "
       << rinternals.CroppingScale[1] << ", " << rinternals.CroppingScale[2] << endl;
  }
}

//----------------------------------------------------------------------------
void vtkImageVolumeRepresentation::SetRequestedRenderMode(int mode)
{
  if (this->RInternals)
  {
    auto& rinternals = *(this->RInternals);
    if (auto smartVolumeMapper = vtkSmartVolumeMapper::SafeDownCast(rinternals.VolumeMapper))
    {
      smartVolumeMapper->SetRequestedRenderMode(mode);
    }
    else if (auto mbMapper = vtkMultiBlockVolumeMapper::SafeDownCast(rinternals.VolumeMapper))
    {
      mbMapper->SetRequestedRenderMode(mode);
    }
  }
}

//----------------------------------------------------------------------------
void vtkImageVolumeRepresentation::SetBlendMode(int blend)
{
  if (this->RInternals)
  {
    this->RInternals->VolumeMapper->SetBlendMode(static_cast<vtkVolumeMapper::BlendModes>(blend));
  }
}

//----------------------------------------------------------------------------
void vtkImageVolumeRepresentation::SetCropping(int crop)
{
  if (this->RInternals)
  {
    this->RInternals->VolumeMapper->SetCropping(crop != 0);
  }
}
//----------------------------------------------------------------------------
void vtkImageVolumeRepresentation::SetCroppingOrigin(double value0, double value1, double value2)
{
  if (this->RInternals)
  {
    this->RInternals->CroppingOrigin[0] = value0;
    this->RInternals->CroppingOrigin[1] = value1;
    this->RInternals->CroppingOrigin[2] = value2;
  }
}

//----------------------------------------------------------------------------
void vtkImageVolumeRepresentation::SetCroppingScale(double value0, double value1, double value2)
{
  if (this->RInternals)
  {
    this->RInternals->CroppingScale[0] = value0;
    this->RInternals->CroppingScale[1] = value1;
    this->RInternals->CroppingScale[2] = value2;
  }
}

//----------------------------------------------------------------------------
void vtkImageVolumeRepresentation::SetIsosurfaceValues(const std::vector<double>& values)
{
  if (this->RInternals)
  {
    this->GetProperty()->GetIsoSurfaceValues()->SetNumberOfContours(values.size());
    for (int i = 0; i < values.size(); i++)
    {
      this->GetProperty()->GetIsoSurfaceValues()->SetValue(i, values[i]);
    }
  }
}
