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

  Program:   ParaView
  Module:    vtkGmshWriter.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 "vtkGmshWriter.h"
#include "gmshCommon.h"

#include "vtkAlgorithm.h"
#include "vtkCellData.h"
#include "vtkCellType.h"
#include "vtkDataArray.h"
#include "vtkDataSetAttributes.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkNew.h"
#include "vtkObjectFactory.h"
#include "vtkPointData.h"
#include "vtkStreamingDemandDrivenPipeline.h"
#include "vtkType.h"
#include "vtkUnsignedCharArray.h"
#include "vtkUnstructuredGrid.h"

#include "vtkMultiBlockDataSet.h"
#include "vtkCompositeDataIterator.h"

#include "gmsh.h"

#include <numeric>
#include <set>
#include <unordered_map>
#include <vector>

//-----------------------------------------------------------------------------
struct MPhysicalGroup
{
  std::string Name = "";
  int Dimension = -1;
  int Entity = -1;
  vtkUnstructuredGrid* Grid = nullptr;
  vtkIdType CellOffset = 0;
  vtkIdType NodeOffset = 0;
};

struct GmshWriterInternal
{
  std::string ModelName = "PVModel";
  std::vector<MPhysicalGroup> PhysicalGroups;
  // Initialize sum to 1 because Gmsh tags begin at 1
  vtkIdType NodeOffsetSum = 1;
  vtkIdType CellOffsetSum = 1;

  double* TimeSteps = nullptr;
  unsigned int NumberOfTimeSteps = 0;
  unsigned int CurrentTimeStep = 0;
  double CurrentTime = 0.0;
};

//-----------------------------------------------------------------------------
namespace
{
static constexpr unsigned char MAX_TYPE = 15u;
typedef std::array<std::vector<std::size_t>, ::MAX_TYPE> ArrayVectorType;
// Associate a VTK type to a Gmsh type and its topological dimension.
// If dimension < 0 or GmshPrimitive >= MAX_TYPE, type is not supported.
// If dimension > 0 and GmshPrimitive == Unsupported, the cell first needs
// to be triangulated in order to be supported by Gmsh.
// clang-format off
static constexpr std::array<std::pair<GmshPrimitive, char>, MAX_TYPE> TRANSLATE_CELLS_TYPE = {
{ { GmshPrimitive::Unsupported, -1 } // VTK_EMPTY_CELL
, { GmshPrimitive::POINT,        0 } // VTK_VERTEX
, { GmshPrimitive::Unsupported, -1 } // VTK_POLY_VERTEX  
, { GmshPrimitive::LINE,         1 } // VTK_LINE
, { GmshPrimitive::Unsupported,  1 } // VTK_POLY_LINE
, { GmshPrimitive::TRIANGLE,     2 } // VTK_TRIANGLE
, { GmshPrimitive::Unsupported,  2 } // VTK_TRIANGLE_STRIP
, { GmshPrimitive::Unsupported,  2 } // VTK_POLYGON
, { GmshPrimitive::QUAD,         2 } // VTK_PIXEL
, { GmshPrimitive::QUAD,         2 } // VTK_QUAD
, { GmshPrimitive::TETRA,        3 } // VTK_TETRA
, { GmshPrimitive::HEXA,         3 } // VTK_VOXEL
, { GmshPrimitive::HEXA,         3 } // VTK_HEXAHEDRON
, { GmshPrimitive::PRISM,        3 } // VTK_WEDGE
, { GmshPrimitive::PYRAMID,      3 } // VTK_PYRAMID
}};
// clang-format on

// Reorder voxel vertices as hexahedron vertices to support it with Gmsh.
// `cellArray` contains every cell vertices, while `idx` is the index of the
// first vertex for the current voxel
void OrderVoxel(std::vector<std::size_t>& cellArray, std::size_t idx)
{
  std::swap(cellArray[idx + 2], cellArray[idx + 3]);
  std::swap(cellArray[idx + 6], cellArray[idx + 7]);
}

// Reorder pixel vertices as quad vertices to support it with Gmsh.
// `cellArray` contains every cell vertices, while `idx` is the index of the
// first vertex for the current pixel
void OrderPixel(std::vector<std::size_t>& cellArray, std::size_t idx)
{
  std::swap(cellArray[idx + 2], cellArray[idx + 3]);
}

// Transform VTK types not supported by Gmsh (like triange strips or polylines)
// as simpler primitives by triangulating them.
vtkIdType AddTriangulatedCell(const MPhysicalGroup& group, std::vector<std::size_t>& gmshNodeTags,
  const std::vector<std::size_t>& vtkCellToTranform, const int dim)
{
  vtkIdType createdCell = 0;
  vtkNew<vtkIdList> nodeIds;
  vtkNew<vtkPoints> dummy;

  for (std::size_t idx : vtkCellToTranform)
  {
    vtkCell* cell = group.Grid->GetCell(idx);
    cell->Triangulate(0, nodeIds, dummy);
    const vtkIdType numNodes = nodeIds->GetNumberOfIds();

    for (vtkIdType i = 0; i < numNodes; ++i)
    {
      gmshNodeTags.push_back(nodeIds->GetId(i) + group.NodeOffset);
    }

    createdCell += (numNodes / dim);
  }

  return createdCell;
}

// Add every supported VTK cells in the gmsh model by type.
vtkIdType FillCells(const MPhysicalGroup& group, const ArrayVectorType& vtkCellIdxsPerType)
{
  vtkDataArray* offsets = group.Grid->GetCells()->GetOffsetsArray();
  vtkDataArray* connectivity = group.Grid->GetCells()->GetConnectivityArray();

  vtkIdType numCreatedCell = 0;

  for (unsigned char currentType = 1; currentType < MAX_TYPE; ++currentType)
  {
    const std::vector<std::size_t>& vtkCellIdxs = vtkCellIdxsPerType[currentType];
    if (vtkCellIdxs.empty())
    {
      continue;
    }

    const char gmshType =
      static_cast<char>(TRANSLATE_CELLS_TYPE[currentType].first);
    // Type not natively supported, it will be triangulated later
    if (gmshType < 0)
    {
      continue;
    }

    // Add cells for this vtk type
    std::vector<std::size_t> gmshNodeTags;
    for (std::size_t index : vtkCellIdxs)
    {
      const int beginCell = gmshNodeTags.size();
      const long offsetBegin = static_cast<long>(*offsets->GetTuple(index));
      const long offsetEnd = static_cast<long>(*offsets->GetTuple(index + 1));
      for (unsigned int j = offsetBegin; j < offsetEnd; ++j)
      {
        const std::size_t vtkNodeIndex = static_cast<std::size_t>(*connectivity->GetTuple(j));
        gmshNodeTags.push_back(vtkNodeIndex + group.NodeOffset);
      }

      // Ordering if needed
      if (currentType == VTK_PIXEL)
      {
        ::OrderPixel(gmshNodeTags, beginCell);
      }
      else if (currentType == VTK_VOXEL)
      {
        ::OrderVoxel(gmshNodeTags, beginCell);
      }
    }

    // Translate non-supported type into supported one
    vtkIdType numCreatedCellFromTriangl = 0;
    const vtkIdType numOfCells = vtkCellIdxs.size();
    if (currentType == VTK_LINE)
    {
      numCreatedCellFromTriangl += ::AddTriangulatedCell(
        group, gmshNodeTags, vtkCellIdxsPerType[VTK_POLY_LINE], 2);
    }
    else if (currentType == VTK_TRIANGLE)
    {
      numCreatedCellFromTriangl += ::AddTriangulatedCell(
        group, gmshNodeTags, vtkCellIdxsPerType[VTK_TRIANGLE_STRIP], 3);
      numCreatedCellFromTriangl += ::AddTriangulatedCell(
        group, gmshNodeTags, vtkCellIdxsPerType[VTK_POLYGON], 3);
    }

    // Generate cell gmsh ids
    std::vector<std::size_t> gmshIds(numOfCells + numCreatedCellFromTriangl);
    std::iota(gmshIds.begin(), gmshIds.end(), numCreatedCell + group.CellOffset);
    numCreatedCell += gmshIds.size();

    gmsh::model::mesh::addElementsByType(group.Entity, gmshType, gmshIds, gmshNodeTags);
  }

  return numCreatedCell;
}
}

//-----------------------------------------------------------------------------
vtkStandardNewMacro(vtkGmshWriter);

//-----------------------------------------------------------------------------
vtkGmshWriter::vtkGmshWriter()
  : Internal(new GmshWriterInternal)
{
  this->SetNumberOfInputPorts(1);
}

//-----------------------------------------------------------------------------
vtkGmshWriter::~vtkGmshWriter()
{
  this->SetFileName(nullptr);
  delete this->Internal;
}

//-----------------------------------------------------------------------------
int vtkGmshWriter::LoadGridIntoPhysicalGroup(vtkUnstructuredGrid* input, std::string name, int desiredDimension)
{
  // std::vector<std::size_t> vtkCellIdxsPerType[::MAX_TYPE];
  ::ArrayVectorType vtkCellIdxsPerType;

  // =================== Load Cells =========================
  // Build list of gmsh indexes per type
  {
    vtkCellArray* cells = input->GetCells();
    const vtkIdType numberOfCells = cells->GetNumberOfCells();
    if (numberOfCells <= 0)
    {
      return 0;
    }

    for (vtkIdType i = 0; i < numberOfCells; ++i)
    {
      const int vtkCellType = input->GetCellType(i);

      if (vtkCellType >= 15)
      {
        continue;
      }
      const int vtkTopologicalDim = ::TRANSLATE_CELLS_TYPE[vtkCellType].second;
      if (vtkTopologicalDim < 0)
      {
        continue;
      }
      else if (vtkTopologicalDim != desiredDimension)
      {
        vtkGenericWarningMacro("Dropping " << name << ", cells of different dimension are found.");
        return -1;
      }

      vtkCellIdxsPerType[vtkCellType].push_back(i);
    }
  }
  // ========================================================

  // Define physical group and its entity
  const int currentEntity = gmsh::model::addDiscreteEntity(desiredDimension);
  gmsh::model::addPhysicalGroup(desiredDimension, {currentEntity});
  // TODO add its name

  MPhysicalGroup currentGroup;
  currentGroup.Dimension = desiredDimension;
  currentGroup.Name = name;
  currentGroup.Entity = currentEntity;
  currentGroup.Grid = input;
  currentGroup.CellOffset = this->Internal->CellOffsetSum;
  currentGroup.NodeOffset = this->Internal->NodeOffsetSum;
  this->Internal->PhysicalGroups.push_back(currentGroup);

  // ==================== Load Nodes ========================
  {
    vtkDataArray* points = input->GetPoints()->GetData();

    const vtkIdType numTuples = points->GetNumberOfTuples();
    const vtkIdType numCompnt = points->GetNumberOfComponents();
    std::vector<double> inlineCoords(numTuples * numCompnt);
    // Store point coordinates in a structure Gmsh can understand
    for (vtkIdType i = 0, counter = 0; i < numTuples; ++i)
    {
      for (vtkIdType j = 0; j < numCompnt; ++j)
      {
        inlineCoords[counter] = *(points->GetTuple(i) + j);
        ++counter;
      }
    }

    std::vector<std::size_t> gmshNodeTags(numTuples);
    std::iota(gmshNodeTags.begin(), gmshNodeTags.end(), currentGroup.NodeOffset);
    this->Internal->NodeOffsetSum += numTuples;

    gmsh::model::mesh::addNodes(
      desiredDimension, currentEntity, gmshNodeTags, inlineCoords);
  }
  // ========================================================

  // Actually load cells
  const vtkIdType numCells = ::FillCells(currentGroup, vtkCellIdxsPerType);
  this->Internal->CellOffsetSum += numCells;

  return 1;
}

//-----------------------------------------------------------------------------
int vtkGmshWriter::ProcessRequest(
  vtkInformation* request, vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  if (request->Has(vtkDemandDrivenPipeline::REQUEST_INFORMATION()))
  {
    return this->RequestInformation(request, inputVector, outputVector);
  }
  else if (request->Has(vtkStreamingDemandDrivenPipeline::REQUEST_UPDATE_EXTENT()))
  {
    return this->RequestUpdateExtent(request, inputVector, outputVector);
  }
  else if (request->Has(vtkDemandDrivenPipeline::REQUEST_DATA()))
  {
    return this->RequestData(request, inputVector, outputVector);
  }

  return this->Superclass::ProcessRequest(request, inputVector, outputVector);
}

//-----------------------------------------------------------------------------
int vtkGmshWriter::RequestInformation(
  vtkInformation*, vtkInformationVector** inputVector, vtkInformationVector*)
{
  vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);
  if (this->WriteAllTimeSteps && inInfo->Has(vtkStreamingDemandDrivenPipeline::TIME_STEPS()))
  {
    this->Internal->NumberOfTimeSteps =
      inInfo->Length(vtkStreamingDemandDrivenPipeline::TIME_STEPS());
  }
  else
  {
    this->Internal->NumberOfTimeSteps = 0;
  }
  this->Internal->CurrentTimeStep = 0;

  return 1;
}

//-----------------------------------------------------------------------------
int vtkGmshWriter::RequestUpdateExtent(
  vtkInformation*, vtkInformationVector** inputVector, vtkInformationVector*)
{
  vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);
  if (this->WriteAllTimeSteps && inInfo->Has(vtkStreamingDemandDrivenPipeline::TIME_STEPS()))
  {
    double* timeSteps = inInfo->Get(vtkStreamingDemandDrivenPipeline::TIME_STEPS());
    this->Internal->CurrentTime = timeSteps[this->Internal->CurrentTimeStep];
    inputVector[0]->GetInformationObject(0)->Set(
      vtkStreamingDemandDrivenPipeline::UPDATE_TIME_STEP(), this->Internal->CurrentTime);
  }
  return 1;
}


//-----------------------------------------------------------------------------
int vtkGmshWriter::RequestData(
  vtkInformation* request, vtkInformationVector** inputVector, vtkInformationVector*)
{
  // make sure the user specified a FileName
  if (!this->FileName)
  {
    vtkErrorMacro(<< "Please specify FileName to use");
    return 0;
  }

  vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);
  this->OriginalInput = vtkDataObject::SafeDownCast(inInfo->Get(vtkDataObject::DATA_OBJECT()));

  // If this is the first request
  // if (this->Internal->CurrentTimeStep == 0)
  {
    std::string file(this->FileName);
    gmsh::initialize();
    gmsh::option::setNumber("General.Verbosity", 1);
    gmsh::option::setNumber("PostProcessing.SaveMesh", 0);
    gmsh::model::add(this->Internal->ModelName.c_str());

    switch (this->OriginalInput->GetDataObjectType())
    {
      case VTK_UNSTRUCTURED_GRID:
      {
        vtkUnstructuredGrid* input = vtkUnstructuredGrid::SafeDownCast(this->OriginalInput);

        std::cout << "Writing unstructured grid" << std::endl;
        // LoadGridIntoPhysicalGroup(input, "");
        break;
      }
      case VTK_MULTIBLOCK_DATA_SET:
      {
        std::cout << "Multiblock dataset" << std::endl;
        vtkMultiBlockDataSet* compositeData = vtkMultiBlockDataSet::SafeDownCast(this->OriginalInput);
        vtkCompositeDataIterator* iter = compositeData->NewIterator();
        for (iter->InitTraversal(); !iter->IsDoneWithTraversal(); iter->GoToNextItem())
        {
          vtkUnstructuredGrid* leaf = vtkUnstructuredGrid::SafeDownCast(iter->GetCurrentDataObject());
          const char* blockName = compositeData->GetMetaData(iter)->Get(vtkCompositeDataSet::NAME());
          if (leaf == nullptr)
          {
            vtkWarningMacro("Leaf " << blockName << " is not an unstructured grid.");
            break;
          }

          const int vtkType = leaf->GetCellType(0);
          int dimension = -1;
          if (vtkType < ::MAX_TYPE)
          {
            dimension = ::TRANSLATE_CELLS_TYPE[vtkType].second;
          }
          if (dimension < 0)
          {
            vtkWarningMacro("Leaf " << blockName << " contains unsupported data type.");
            break;
          }

          std::cout << "writing physical group " << blockName << "of dimension " << dimension << std::endl;
          LoadGridIntoPhysicalGroup(leaf, blockName, dimension);          
        }

        break;
      }
      default:
      {
        vtkErrorMacro("Non supported type.");
        return 0;
      }
    }

    // Load and save topology
    gmsh::write(this->FileName);
  }

  gmsh::clear();
  gmsh::finalize();

  return 1;
}

//------------------------------------------------------------------------------
void vtkGmshWriter::WriteData()
{
  /*
  const int numViews = this->Internal->CellViews.size() + this->Internal->NodeViews.size();
  for (int tag = 0; tag < numViews; ++tag)
  {
    gmsh::view::write(tag, this->FileName, true);
  }
  */
}

//-----------------------------------------------------------------------------
int vtkGmshWriter::FillInputPortInformation(int, vtkInformation* info)
{
  info->Remove(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE());
  info->Append(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkUnstructuredGrid");
  info->Append(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkMultiBlockDataSet");
  return 1;
}

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

  os << indent << "FileName: " << (this->GetFileName() ? this->GetFileName() : "(none)") << indent
     << ", WriteAllTimeSteps: " << this->WriteAllTimeSteps << indent
     << ", WriteGmshSpecificArray: " << this->WriteGmshSpecificArray << std::endl;
}
