#include "FEAdaptor.h"
#include "FEDataStructures.h"
#include <iostream>

#include <vtkCPDataDescription.h>
#include <vtkCPInputDataDescription.h>
#include <vtkCPProcessor.h>
#include <vtkCPPythonScriptPipeline.h>
#include <vtkCellData.h>
#include <vtkCellType.h>
#include <vtkCommunicator.h>
#include <vtkDoubleArray.h>
#include <vtkFloatArray.h>
#include <vtkImageData.h>
#include <vtkMultiBlockDataSet.h>
#include <vtkMultiPieceDataSet.h>
#include <vtkMultiProcessController.h>
#include <vtkNew.h>
#include <vtkPointData.h>
#include <vtkPoints.h>

#include <mpi.h>

namespace
{
vtkCPProcessor* Processor = NULL;
vtkMultiBlockDataSet* VTKGrid;

void BuildVTKGrid(Grid& grid)
{
  vtkNew<vtkImageData> imageData;
  imageData->SetSpacing(grid.GetSpacing());
  imageData->SetOrigin(0, 0, 0); // Not necessary for (0,0,0)
  unsigned int* extents = grid.GetExtents();
  int extents2[6];
  for (int i = 0; i < 6; i++)
  {
    extents2[i] = static_cast<int>(extents[i]);
  }
  imageData->SetExtent(extents2);
  int mpiSize = 1;
  int mpiRank = 0;
  MPI_Comm_rank(MPI_COMM_WORLD, &mpiRank);
  MPI_Comm_size(MPI_COMM_WORLD, &mpiSize);
  vtkNew<vtkMultiPieceDataSet> multiPiece;
  multiPiece->SetNumberOfPieces(mpiSize);
  multiPiece->SetPiece(mpiRank, imageData.GetPointer());
  VTKGrid->SetNumberOfBlocks(1);
  VTKGrid->SetBlock(0, multiPiece.GetPointer());
}

void UpdateVTKAttributes(Grid& grid, Attributes& attributes)
{
  int mpiRank = 0;
  MPI_Comm_rank(MPI_COMM_WORLD, &mpiRank);
  vtkMultiPieceDataSet* multiPiece = vtkMultiPieceDataSet::SafeDownCast(VTKGrid->GetBlock(0));
  vtkDataSet* dataSet = vtkDataSet::SafeDownCast(multiPiece->GetPiece(mpiRank));
  if (dataSet->GetPointData()->GetNumberOfArrays() == 0)
  {
    // velocity array
    vtkNew<vtkDoubleArray> velocity;
    velocity->SetName("velocity");
    velocity->SetNumberOfComponents(3);
    velocity->SetNumberOfTuples(static_cast<vtkIdType>(grid.GetNumberOfLocalPoints()));
    dataSet->GetPointData()->AddArray(velocity.GetPointer());
  }
  if (dataSet->GetCellData()->GetNumberOfArrays() == 0)
  {
    // pressure array
    vtkNew<vtkFloatArray> pressure;
    pressure->SetName("pressure");
    pressure->SetNumberOfComponents(1);
    dataSet->GetCellData()->AddArray(pressure.GetPointer());
  }
  vtkDoubleArray* velocity =
    vtkDoubleArray::SafeDownCast(dataSet->GetPointData()->GetArray("velocity"));
  // The velocity array is ordered as vx0,vx1,vx2,..,vy0,vy1,vy2,..,vz0,vz1,vz2,..
  // so we need to create a full copy of it with VTK's ordering of
  // vx0,vy0,vz0,vx1,vy1,vz1,..
  double* velocityData = attributes.GetVelocityArray();
  vtkIdType numTuples = velocity->GetNumberOfTuples();
  for (vtkIdType i = 0; i < numTuples; i++)
  {
    double values[3] = { velocityData[i], velocityData[i + numTuples],
      velocityData[i + 2 * numTuples] };
    velocity->SetTypedTuple(i, values);
  }

  vtkFloatArray* pressure =
    vtkFloatArray::SafeDownCast(dataSet->GetCellData()->GetArray("pressure"));
  // The pressure array is a scalar array so we can reuse
  // memory as long as we ordered the points properly.
  float* pressureData = attributes.GetPressureArray();
  pressure->SetArray(pressureData, static_cast<vtkIdType>(grid.GetNumberOfLocalCells()), 1);
}

void BuildVTKDataStructures(Grid& grid, Attributes& attributes)
{
  if (VTKGrid == NULL)
  {
    // The grid structure isn't changing so we only build it
    // the first time it's needed. If we needed the memory
    // we could delete it and rebuild as necessary.
    VTKGrid = vtkMultiBlockDataSet::New();
    BuildVTKGrid(grid);
  }
  UpdateVTKAttributes(grid, attributes);
}
}

namespace FEAdaptor
{

void Initialize(int numScripts, char* scripts[])
{
  if (Processor == NULL)
  {
    Processor = vtkCPProcessor::New();
    Processor->Initialize();
  }
  else
  {
    Processor->RemoveAllPipelines();
  }
  for (int i = 1; i < numScripts; i++)
  {
    vtkNew<vtkCPPythonScriptPipeline> pipeline;
    pipeline->Initialize(scripts[i]);
    Processor->AddPipeline(pipeline.GetPointer());
  }
}

void Finalize()
{
  if (Processor)
  {
    Processor->Delete();
    Processor = NULL;
  }
  if (VTKGrid)
  {
    VTKGrid->Delete();
    VTKGrid = NULL;
  }
}

void CoProcess(
  Grid& grid, Attributes& attributes, double time, unsigned int timeStep, bool lastTimeStep)
{
  vtkNew<vtkCPDataDescription> dataDescription;
  dataDescription->AddInput("input");
  dataDescription->SetTimeData(time, timeStep);
  if (lastTimeStep == true)
  {
    // assume that we want to all the pipelines to execute if it
    // is the last time step.
    dataDescription->ForceOutputOn();
  }
  if (Processor->RequestDataDescription(dataDescription.GetPointer()) != 0)
  {
    BuildVTKDataStructures(grid, attributes);
    dataDescription->GetInputDescriptionByName("input")->SetGrid(VTKGrid);

    // figure out the whole extent of the grid
    unsigned int* extent = grid.GetExtents();
    int wholeExtent[6], tmp[6];
    for (int i = 0; i < 3; i++)
    {
      wholeExtent[i * 2] = -static_cast<int>(extent[i * 2]); // negate for parallel communication
      wholeExtent[i * 2 + 1] = static_cast<int>(extent[i * 2 + 1]);
    }
    vtkMultiProcessController::GetGlobalController()->AllReduce(
      wholeExtent, tmp, 6, vtkCommunicator::MAX_OP);
    for (int i = 0; i < 3; i++)
    {
      tmp[i * 2] = -tmp[i * 2];
    }
    dataDescription->GetInputDescriptionByName("input")->SetWholeExtent(tmp);
    Processor->CoProcess(dataDescription.GetPointer());
  }
}
} // end of Catalyst namespace
