// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause
#include "vtkCellGridWriter.h"

#include "vtkArrayDispatch.h"
#include "vtkCellAttribute.h"
#include "vtkCellGrid.h"
#include "vtkCellGridIOQuery.h"
#include "vtkCellMetadata.h"
#include "vtkDataArray.h"
#include "vtkDataArrayRange.h"
#include "vtkDataSetAttributes.h"
#include "vtkIOCellGrid.h"
#include "vtkInformation.h"
#include "vtkObjectFactory.h"
#include "vtkOptions.h"
#include "vtkStringFormatter.h"
#include "vtkVariant.h"

#include <fstream>
#include <sstream>

VTK_ABI_NAMESPACE_BEGIN

vtkStandardNewMacro(vtkCellGridWriter);

vtkCellGridWriter::vtkCellGridWriter()
{
  // Ensure vtkCellGridIOQuery is registered.
  vtkIOCellGrid::RegisterCellsAndResponders();
}

vtkCellGridWriter::~vtkCellGridWriter()
{
  this->SetFileName(nullptr);
}

void vtkCellGridWriter::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
  os << indent << "FileName: " << this->FileName << "\n";
  os << indent << "FileFormat: " << this->FileFormat << "\n";
}

vtkCellGrid* vtkCellGridWriter::GetInput()
{
  return vtkCellGrid::SafeDownCast(this->Superclass::GetInputDataObject(0, 0));
}

vtkCellGrid* vtkCellGridWriter::GetInput(int port)
{
  return vtkCellGrid::SafeDownCast(this->Superclass::GetInputDataObject(port, 0));
}

int vtkCellGridWriter::FillInputPortInformation(int vtkNotUsed(port), vtkInformation* info)
{
  info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkCellGrid");
  return 1;
}

std::string dataTypeToString(int dataType)
{
  switch (dataType)
  {
#if VTK_SIZEOF_CHAR == 1
    case VTK_CHAR:
      return "vtktypeint8";
    case VTK_SIGNED_CHAR:
      return "vtktypeint8";
    case VTK_UNSIGNED_CHAR:
      return "vtktypeuint8";
#else
#error "Unhandled char size " VTK_SIZEOF_CHAR
#endif

    case VTK_DOUBLE:
      return "double";
    case VTK_FLOAT:
      return "float";

#if VTK_SIZEOF_INT == 2
    case VTK_INT:
      return "vtktypeint16";
    case VTK_UNSIGNED_INT:
      return "vtktypeuint16";
#elif VTK_SIZEOF_INT == 4
    case VTK_INT:
      return "vtktypeint32";
    case VTK_UNSIGNED_INT:
      return "vtktypeuint32";
#elif VTK_SIZEOF_INT == 8
    case VTK_INT:
      return "vtktypeint64";
    case VTK_UNSIGNED_INT:
      return "vtktypeuint64";
#else
#error "Unhandled int size " VTK_SIZEOF_INT
#endif

#if VTK_SIZEOF_LONG == 4
    case VTK_LONG:
      return "vtktypeint32";
    case VTK_UNSIGNED_LONG:
      return "vtktypeuint32";
#elif VTK_SIZEOF_LONG == 8
    case VTK_LONG:
      return "vtktypeint32";
    case VTK_UNSIGNED_LONG:
      return "vtktypeuint32";
#else
#error "Unhandled int size " VTK_SIZEOF_INT
#endif

#if VTK_SIZEOF_LONG == 4
    case VTK_LONG_LONG:
      return "vtktypeint32";
    case VTK_UNSIGNED_LONG_LONG:
      return "vtktypeuint32";
#elif VTK_SIZEOF_LONG == 8
    case VTK_LONG_LONG:
      return "vtktypeint64";
    case VTK_UNSIGNED_LONG_LONG:
      return "vtktypeuint64";
#else
#error "Unhandled long size " VTK_SIZEOF_LONG
#endif

#if VTK_SIZEOF_SHORT == 1
    case VTK_SHORT:
      return "vtktypeint8";
    case VTK_UNSIGNED_SHORT:
      return "vtktypeuint8";
#elif VTK_SIZEOF_SHORT == 2
    case VTK_SHORT:
      return "vtktypeint16";
    case VTK_UNSIGNED_SHORT:
      return "vtktypeuint16";
#elif VTK_SIZEOF_SHORT == 4
    case VTK_SHORT:
      return "vtktypeint32";
    case VTK_UNSIGNED_SHORT:
      return "vtktypeuint32";
#else
#error "Unhandled short size " VTK_SIZEOF_SHORT
#endif

#ifdef VTK_USE_64BIT_IDS
    case VTK_ID_TYPE:
      return "vtktypeint64";
#else
    case VTK_ID_TYPE:
      return "vtktypeint32";
#endif

    default:
      break;
  }
  return "unhandled";
}

struct WriteDataArrayWorker
{
  WriteDataArrayWorker(nlohmann::json& result)
    : m_result(result)
  {
  }

  template <typename InArrayT>
  void operator()(InArrayT* inArray)
  {
    using T = vtk::GetAPIType<InArrayT>;
    const auto inRange = vtk::DataArrayValueRange(inArray);
    T val;
    for (const auto& value : inRange)
    {
      val = value;
      m_result.push_back(val);
    }
  }

  nlohmann::json& m_result;
};

nlohmann::json serializeArrayValues(vtkAbstractArray* arr)
{
  auto result = nlohmann::json::array();
  if (!arr)
  {
    return result;
  }

  if (auto* darr = vtkDataArray::SafeDownCast(arr))
  {
    using Dispatcher = vtkArrayDispatch::DispatchByValueType<vtkArrayDispatch::AllTypes>;
    WriteDataArrayWorker worker(result);
    if (!Dispatcher::Execute(darr, worker))
    {
      worker(darr);
    }
  }
  else
  {
    for (vtkIdType ii = 0; ii < arr->GetNumberOfValues(); ++ii)
    {
      result.push_back(static_cast<std::string>(arr->GetVariantValue(ii).ToString()));
    }
  }
  return result;
}

bool getCachedRange(vtkCellGridRangeQuery::CacheMap& rangeCache, vtkCellAttribute* attribute,
  nlohmann::json& rangeInfo)
{
  auto it = rangeCache.find(attribute);
  if (it == rangeCache.end())
  {
    return false;
  }
  rangeInfo = nlohmann::json::object();
  int compNum = 0;
  for (const auto& compRange : it->second)
  {
    nlohmann::json entry{ { "min", compRange.FiniteRange[0] },
      { "max", compRange.FiniteRange[1] } };
    if (compNum == 1)
    {
      rangeInfo["L₁"] = entry;
    }
    else if (compNum == 0)
    {
      rangeInfo["L₂"] = entry;
    }
    else
    {
      rangeInfo[vtk::to_string(compNum - 2)] = entry;
    }
    ++compNum;
  }
  return true;
}

bool vtkCellGridWriter::ToJSON(nlohmann::json& data, vtkCellGrid* grid)
{
  // Iterate all the vtkDataSetAttributes held by the grid.
  // As we go, store a map from each array in each vtkDataSetAttributes
  // to a "location" for the array so we can refer to the arrays later
  // by a persistent name instead of by pointer.
  std::unordered_map<vtkAbstractArray*, vtkStringToken> arrayLocations;
  auto arrayGroups = nlohmann::json::object();
  for (const auto& entry : grid->GetArrayGroups())
  {
    std::string groupName;
    vtkStringToken groupToken(entry.first);
    if (!groupToken.Data().empty())
    {
      groupName = groupToken.Data();
    }
    else
    {
      std::ostringstream groupId;
      groupId << entry.first;
      groupName = groupId.str();
    }
    nlohmann::json arraysInGroup;

    // Fetch arrays serving in specific roles.
    vtkDataArray* groupScalars = entry.second->GetScalars();
    vtkDataArray* groupVectors = entry.second->GetVectors();
    vtkDataArray* groupTensors = entry.second->GetTensors();
    vtkDataArray* groupTCoords = entry.second->GetTCoords();
    vtkDataArray* groupTangents = entry.second->GetTangents();
    vtkDataArray* groupGlobalIds = entry.second->GetGlobalIds();
    vtkAbstractArray* groupPedigreeIds = entry.second->GetPedigreeIds();
    vtkDataArray* groupRationalWeights = entry.second->GetRationalWeights();
    vtkDataArray* groupHigherOrderDegrees = entry.second->GetHigherOrderDegrees();

    for (vtkIdType ii = 0; ii < entry.second->GetNumberOfArrays(); ++ii)
    {
      vtkAbstractArray* arr = entry.second->GetAbstractArray(ii);
      if (!arr)
      {
        continue;
      }
      arrayLocations[arr] = groupName;
      nlohmann::json arrayRecord{ { "name", arr->GetName() },
        { "tuples", arr->GetNumberOfTuples() }, { "components", arr->GetNumberOfComponents() },
        { "type", dataTypeToString(arr->GetDataType()) }, { "data", serializeArrayValues(arr) } };
      if (arr == groupScalars)
      {
        arrayRecord["default_scalars"] = true;
      }
      if (arr == groupVectors)
      {
        arrayRecord["default_vectors"] = true;
      }
      if (arr == groupTensors)
      {
        arrayRecord["default_tensors"] = true;
      }
      if (arr == groupTCoords)
      {
        arrayRecord["default_tcoords"] = true;
      }
      if (arr == groupTangents)
      {
        arrayRecord["default_tangents"] = true;
      }
      if (arr == groupGlobalIds)
      {
        arrayRecord["default_global_ids"] = true;
      }
      if (arr == groupPedigreeIds)
      {
        arrayRecord["default_pedigree_ids"] = true;
      }
      if (arr == groupRationalWeights)
      {
        arrayRecord["default_rational_weights"] = true;
      }
      if (arr == groupHigherOrderDegrees)
      {
        arrayRecord["default_higher_order_degrees"] = true;
      }
      arraysInGroup.push_back(arrayRecord);
    }
    if (!arraysInGroup.empty())
    {
      arrayGroups[groupName] = arraysInGroup;
    }
  }

  auto attributes = nlohmann::json::array();
  auto* shapeAtt = grid->GetShapeAttribute();
  auto& rangeCache = grid->GetRangeCache();
  for (const auto cellAttId : grid->GetCellAttributeIds())
  {
    auto* cellAtt = grid->GetCellAttributeById(cellAttId);
    if (cellAtt)
    {
      nlohmann::json record{ { "name", cellAtt->GetName().Data() },
        { "space", cellAtt->GetSpace().Data() },
        { "components", cellAtt->GetNumberOfComponents() } };
      if (cellAtt == shapeAtt)
      {
        record["shape"] = true;
      }
      nlohmann::json rangeInfo;
      if (getCachedRange(rangeCache, cellAtt, rangeInfo))
      {
        record["range"] = rangeInfo;
      }
      attributes.push_back(record);
    }
  }

  auto cellTypes = nlohmann::json::array();
#if 0
  for (const auto& cellType : grid->CellTypeArray())
  {
    auto entry = nlohmann::json::object({ { "type", cellType.Data() } });
    cellTypes.push_back(entry);
  }
#endif
  // Now provide vtkCellMetadata subclasses with a chance to add to \a cellTypes.
  vtkNew<vtkCellGridIOQuery> query;
  query->PrepareToSerialize(cellTypes, attributes, arrayLocations);
  if (!grid->Query(query))
  {
    vtkErrorMacro("Could not prepare cell metadata.");
    return false;
  }

  auto schemaName = grid->GetSchemaName();
  if (!schemaName.IsValid())
  {
    schemaName = "dg leaf";
  }
  // clang-format off
  data = {
    { "data-type", "cell-grid" }, { "arrays", arrayGroups }, { "attributes", attributes },
    { "cell-types", cellTypes },
    { "format-version", 1 }, // A version number for the file format (i.e.., JSON)
    { "schema-name", schemaName.Data() }, // A name for the schema (key/value structure) of this file's content.
    { "schema-version", grid->GetSchemaVersion() }, // A version number for the file's schema.
    { "content-version", grid->GetContentVersion() } // A version number for the file's content (key/value data).
  };
  // clang-format on

  return true;
}

void vtkCellGridWriter::WriteData()
{
  if (!this->FileName || !this->FileName[0])
  {
    vtkErrorMacro("No filename set.");
    return;
  }

  auto* grid = this->GetInput();
  if (!grid)
  {
    vtkErrorMacro("No input dataset to write to \"" << this->FileName << "\"");
    return;
  }

  std::ofstream output(this->FileName);
  if (!output.good())
  {
    vtkErrorMacro("Could not open \"" << this->FileName << "\" for writing.");
    return;
  }

  nlohmann::json data;
  if (this->ToJSON(data, grid))
  {
    switch (this->FileFormat)
    {
      default:
      case Format::PlainText:
        output << data;
        break;
      case Format::MessagePack:
      {
        output << "vtkCellGrid\nMessagePack\nv1\n";
        std::vector<char> mpack;
        nlohmann::json::to_msgpack(data, mpack);
        auto writeSize = mpack.size() / sizeof(std::istream::char_type) +
          (mpack.size() % sizeof(std::istream::char_type) ? 1 : 0);
        output.write(reinterpret_cast<std::ostream::char_type*>(mpack.data()), writeSize);
      }
      break;
    }
    output.close();
  }
  else
  {
    vtkErrorMacro("Could not write JSON to \"" << this->FileName << "\".");
    return; /* 0 */
  }
}

VTK_ABI_NAMESPACE_END