// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause
// .NAME TestUniformGridGhostDataGenerator.cxx -- Serial test for ghost data
//
// .SECTION Description
//  Serial tests for 2-D and 3-D ghost data generation of multi-block uniform
//  grid datasets. The tests apply an XYZ field to the nodes and cells of the
//  domain and ensure that the created ghost data have the correct fields.

// C++ includes
#include <cassert>
#include <iostream>
#include <set>
#include <sstream>
#include <string>
#include <vector>

// VTK includes
#include "vtkCell.h"
#include "vtkCellData.h"
#include "vtkDataObject.h"
#include "vtkDataSet.h"
#include "vtkDoubleArray.h"
#include "vtkGhostCellsGenerator.h"
#include "vtkInformation.h"
#include "vtkIntArray.h"
#include "vtkLogger.h"
#include "vtkMathUtilities.h"
#include "vtkMultiBlockDataSet.h"
#include "vtkMultiPieceDataSet.h"
#include "vtkPointData.h"
#include "vtkStreamingDemandDrivenPipeline.h"
#include "vtkUniformGrid.h"
#include "vtkUniformGridPartitioner.h"
#include "vtkUnsignedCharArray.h"
#include "vtkUnsignedIntArray.h"
#include "vtkXMLImageDataWriter.h"
#include "vtkXMLMultiBlockDataWriter.h"

//------------------------------------------------------------------------------
bool CheckNodeFieldsForGrid(vtkUniformGrid* grid)
{
  assert("pre: grid should not be nullptr" && (grid != nullptr));
  assert("pre: grid should have a NODE-XYZ array" && grid->GetPointData()->HasArray("NODE-XYZ"));

  double xyz[3];
  vtkDoubleArray* array =
    vtkArrayDownCast<vtkDoubleArray>(grid->GetPointData()->GetArray("NODE-XYZ"));
  assert("pre: num tuples must match number of nodes" &&
    (array->GetNumberOfTuples() == grid->GetNumberOfPoints()));
  assert("pre: num components must be 3" && (array->GetNumberOfComponents() == 3));

  for (vtkIdType idx = 0; idx < grid->GetNumberOfPoints(); ++idx)
  {
    grid->GetPoint(idx, xyz);

    for (int i = 0; i < 3; ++i)
    {
      if (!vtkMathUtilities::FuzzyCompare(xyz[i], array->GetComponent(idx, i)))
      {
        return false;
      } // END if fuzzy-compare
    }   // END for all components
  }     // END for all nodes
  return true;
}

//------------------------------------------------------------------------------
bool CheckCellFieldsForGrid(vtkUniformGrid* grid)
{
  assert("pre: grid should not be nullptr" && (grid != nullptr));
  assert("pre: grid should have a NODE-XYZ array" && grid->GetCellData()->HasArray("CELL-XYZ"));

  double centroid[3];
  double xyz[3];
  vtkDoubleArray* array =
    vtkArrayDownCast<vtkDoubleArray>(grid->GetCellData()->GetArray("CELL-XYZ"));
  assert("pre: num tuples must match number of nodes" &&
    (array->GetNumberOfTuples() == grid->GetNumberOfCells()));
  assert("pre: num components must be 3" && (array->GetNumberOfComponents() == 3));

  for (vtkIdType cellIdx = 0; cellIdx < grid->GetNumberOfCells(); ++cellIdx)
  {
    vtkCell* c = grid->GetCell(cellIdx);
    assert("pre: cell is not nullptr" && (c != nullptr));

    double xsum = 0.0;
    double ysum = 0.0;
    double zsum = 0.0;
    for (vtkIdType node = 0; node < c->GetNumberOfPoints(); ++node)
    {
      vtkIdType meshPntIdx = c->GetPointId(node);
      grid->GetPoint(meshPntIdx, xyz);
      xsum += xyz[0];
      ysum += xyz[1];
      zsum += xyz[2];
    } // END for all nodes

    centroid[0] = xsum / c->GetNumberOfPoints();
    centroid[1] = ysum / c->GetNumberOfPoints();
    centroid[2] = zsum / c->GetNumberOfPoints();

    for (int i = 0; i < 3; ++i)
    {
      if (!vtkMathUtilities::FuzzyCompare(centroid[i], array->GetComponent(cellIdx, i)))
      {
        return false;
      } // END if fuzz-compare
    }   // END for all components
  }     // END for all cells
  return true;
}

//------------------------------------------------------------------------------
int CheckFields(vtkMultiBlockDataSet* mbds, bool hasNodeData, bool hasCellData)
{
  assert("pre: input multi-block is nullptr" && (mbds != nullptr));

  if (!hasNodeData && !hasCellData)
  {
    return 0;
  }

  for (unsigned int block = 0; block < mbds->GetNumberOfBlocks(); ++block)
  {
    vtkUniformGrid* grid = vtkUniformGrid::SafeDownCast(mbds->GetBlock(block));
    assert("pre: grid is not nullptr" && (grid != nullptr));

    if (hasNodeData)
    {
      if (!CheckNodeFieldsForGrid(grid))
      {
        return 1;
      }
    }

    if (hasCellData)
    {
      if (!CheckCellFieldsForGrid(grid))
      {
        return 1;
      }
    }

  } // END for all blocks

  return 0;
}

//------------------------------------------------------------------------------
// Description:
// Write the uniform grid multi-block dataset into an XML file.
void WriteMultiBlock(vtkMultiBlockDataSet* mbds, const std::string& prefix)
{
  assert("pre: Multi-block is nullptr!" && (mbds != nullptr));

  vtkXMLMultiBlockDataWriter* writer = vtkXMLMultiBlockDataWriter::New();
  assert("pre: Cannot allocate writer" && (writer != nullptr));

  std::ostringstream oss;
  oss.str("");
  oss << prefix << mbds->GetNumberOfBlocks() << "." << writer->GetDefaultFileExtension();
  writer->SetFileName(oss.str().c_str());
  writer->SetInputData(mbds);
  writer->Write();

  writer->Delete();
}

//------------------------------------------------------------------------------
// Description:
// Adds and XYZ vector field in the nodes of the data-set
void AddNodeCenteredXYZField(vtkMultiBlockDataSet* mbds)
{
  assert("pre: Multi-block is nullptr!" && (mbds != nullptr));

  for (unsigned int block = 0; block < mbds->GetNumberOfBlocks(); ++block)
  {
    vtkUniformGrid* grid = vtkUniformGrid::SafeDownCast(mbds->GetBlock(block));
    assert("pre: grid is nullptr for the given block" && (grid != nullptr));

    vtkDoubleArray* nodeXYZArray = vtkDoubleArray::New();
    nodeXYZArray->SetName("NODE-XYZ");
    nodeXYZArray->SetNumberOfComponents(3);
    nodeXYZArray->SetNumberOfTuples(grid->GetNumberOfPoints());

    double xyz[3];
    for (vtkIdType pntIdx = 0; pntIdx < grid->GetNumberOfPoints(); ++pntIdx)
    {
      grid->GetPoint(pntIdx, xyz);
      nodeXYZArray->SetComponent(pntIdx, 0, xyz[0]);
      nodeXYZArray->SetComponent(pntIdx, 1, xyz[1]);
      nodeXYZArray->SetComponent(pntIdx, 2, xyz[2]);
    } // END for all points

    grid->GetPointData()->AddArray(nodeXYZArray);
    nodeXYZArray->Delete();
  } // END for all blocks
}

//------------------------------------------------------------------------------
// Description:
// Adds and XYZ vector field in the nodes of the dataset
void AddCellCenteredXYZField(vtkMultiBlockDataSet* mbds)
{
  assert("pre: Multi-block is nullptr!" && (mbds != nullptr));

  for (unsigned int block = 0; block < mbds->GetNumberOfBlocks(); ++block)
  {
    vtkUniformGrid* grid = vtkUniformGrid::SafeDownCast(mbds->GetBlock(block));
    assert("pre: grid is nullptr for the given block" && (grid != nullptr));

    vtkDoubleArray* cellXYZArray = vtkDoubleArray::New();
    cellXYZArray->SetName("CELL-XYZ");
    cellXYZArray->SetNumberOfComponents(3);
    cellXYZArray->SetNumberOfTuples(grid->GetNumberOfCells());

    double centroid[3];
    double xyz[3];
    for (vtkIdType cellIdx = 0; cellIdx < grid->GetNumberOfCells(); ++cellIdx)
    {
      vtkCell* c = grid->GetCell(cellIdx);
      assert("pre: cell is not nullptr" && (c != nullptr));

      double xsum = 0.0;
      double ysum = 0.0;
      double zsum = 0.0;
      for (vtkIdType node = 0; node < c->GetNumberOfPoints(); ++node)
      {
        vtkIdType meshPntIdx = c->GetPointId(node);
        grid->GetPoint(meshPntIdx, xyz);
        xsum += xyz[0];
        ysum += xyz[1];
        zsum += xyz[2];
      } // END for all nodes

      centroid[0] = xsum / c->GetNumberOfPoints();
      centroid[1] = ysum / c->GetNumberOfPoints();
      centroid[2] = zsum / c->GetNumberOfPoints();

      cellXYZArray->SetComponent(cellIdx, 0, centroid[0]);
      cellXYZArray->SetComponent(cellIdx, 1, centroid[1]);
      cellXYZArray->SetComponent(cellIdx, 2, centroid[2]);
    } // END for all cells

    grid->GetCellData()->AddArray(cellXYZArray);
    cellXYZArray->Delete();
  } // END for all blocks
}

//------------------------------------------------------------------------------
// Description:
// Creates a test data-set.
vtkMultiBlockDataSet* GetDataSet(double globalOrigin[3], int WholeExtent[6], double gridSpacing[3],
  const int numPartitions, const int numGhosts, const bool AddNodeData, const bool AddCellData)
{
  // STEP 0: Get the global grid dimensions
  int dims[3];
  vtkStructuredData::GetDimensionsFromExtent(WholeExtent, dims);

  // STEP 1: Get the whole grid
  vtkUniformGrid* wholeGrid = vtkUniformGrid::New();
  wholeGrid->SetOrigin(globalOrigin);
  wholeGrid->SetSpacing(gridSpacing);
  wholeGrid->SetDimensions(dims);

  // STEP 2: Partition the whole grid
  vtkUniformGridPartitioner* gridPartitioner = vtkUniformGridPartitioner::New();
  gridPartitioner->SetInputData(wholeGrid);
  gridPartitioner->SetNumberOfPartitions(numPartitions);
  gridPartitioner->SetNumberOfGhostLayers(numGhosts);
  gridPartitioner->Update();

  // STEP 3: Get the partitioned dataset
  vtkMultiBlockDataSet* mbds = vtkMultiBlockDataSet::SafeDownCast(gridPartitioner->GetOutput());
  mbds->Register(nullptr);

  // STEP 4: Delete temporary data
  wholeGrid->Delete();
  gridPartitioner->Delete();

  // STEP 5: Add node-centered and cell-centered fields
  if (AddNodeData)
  {
    AddNodeCenteredXYZField(mbds);
  }

  if (AddCellData)
  {
    AddCellCenteredXYZField(mbds);
  }

  return (mbds);
}

//------------------------------------------------------------------------------
// Description:
// Tests UniformGridGhostDataGenerator
int Test2D(
  const bool hasNodeData, const bool hasCellData, const int numPartitions, const int numGhosts)
{
  int rc = 0;

  int WholeExtent[6] = { 0, 49, 0, 49, 0, 0 };
  double h[3] = { 0.5, 0.5, 0.5 };
  double p[3] = { 0.0, 0.0, 0.0 };

  vtkMultiBlockDataSet* mbds =
    GetDataSet(p, WholeExtent, h, numPartitions, numGhosts, hasNodeData, hasCellData);
  // WriteMultiBlock(mbds, "INITIAL");

  vtkGhostCellsGenerator* ghostDataGenerator = vtkGhostCellsGenerator::New();

  ghostDataGenerator->SetInputData(mbds);
  ghostDataGenerator->SetNumberOfGhostLayers(1);
  ghostDataGenerator->BuildIfRequiredOff();
  ghostDataGenerator->Update();

  vtkMultiBlockDataSet* ghostedDataSet =
    vtkMultiBlockDataSet::SafeDownCast(ghostDataGenerator->GetOutputDataObject(0));
  //  WriteMultiBlock( ghostedDataSet, "GHOSTED" );

  rc = CheckFields(ghostedDataSet, hasNodeData, hasCellData);
  mbds->Delete();
  ghostDataGenerator->Delete();
  return rc;
}

//------------------------------------------------------------------------------
// Description:
// Tests UniformGridGhostDataGenerator
int Test3D(
  const bool hasNodeData, const bool hasCellData, const int numPartitions, const int numGhosts)
{
  int rc = 0;
  int WholeExtent[6] = { 0, 49, 0, 49, 0, 49 };
  double h[3] = { 0.5, 0.5, 0.5 };
  double p[3] = { 0.0, 0.0, 0.0 };

  vtkMultiBlockDataSet* mbds =
    GetDataSet(p, WholeExtent, h, numPartitions, numGhosts, hasNodeData, hasCellData);
  //  WriteMultiBlock( mbds, "INITIAL");

  vtkGhostCellsGenerator* ghostDataGenerator = vtkGhostCellsGenerator::New();

  ghostDataGenerator->SetInputData(mbds);
  ghostDataGenerator->SetNumberOfGhostLayers(1);
  ghostDataGenerator->BuildIfRequiredOff();
  ghostDataGenerator->Update();

  vtkMultiBlockDataSet* ghostedDataSet =
    vtkMultiBlockDataSet::SafeDownCast(ghostDataGenerator->GetOutput(0));
  //  WriteMultiBlock( ghostedDataSet, "GHOSTED" );

  rc = CheckFields(ghostedDataSet, hasNodeData, hasCellData);
  mbds->Delete();
  ghostDataGenerator->Delete();
  return rc;
}

//------------------------------------------------------------------------------
// Description:
// Tests UniformGridGhostDataGenerator
int TestUniformGridGhostDataGenerator(int, char*[])
{
  int rc = 0;

  rc += Test2D(true, false, 4, 1);
  vtkLog(INFO, "Running Test2D(true, true, 4, 1)");
  rc += Test2D(true, true, 16, 1);
  vtkLog(INFO, "Running Test3D(true, false, 4, 1)");
  rc += Test3D(false, true, 8, 1);
  vtkLog(INFO, "test end " << rc);
  return (rc);
}
