WriterBase.cpp

#include <output/WriterBase.hpp>

#include <mesh/ItemArrayVariant.hpp>
#include <mesh/ItemValueVariant.hpp>
#include <mesh/Mesh.hpp>
#include <mesh/MeshVariant.hpp>
#include <output/NamedDiscreteFunction.hpp>
#include <output/NamedItemArrayVariant.hpp>
#include <output/NamedItemValueVariant.hpp>
#include <output/OutputNamedItemValueSet.hpp>
#include <scheme/DiscreteFunctionP0.hpp>
#include <scheme/DiscreteFunctionP0Vector.hpp>
#include <scheme/DiscreteFunctionVariant.hpp>
#include <scheme/IDiscreteFunctionDescriptor.hpp>
#include <utils/Exceptions.hpp>

template <typename DiscreteFunctionType>
void
WriterBase::_registerDiscreteFunction(const std::string& name,
                                      const DiscreteFunctionType& discrete_function,
                                      OutputNamedItemDataSet& named_item_data_set)
{
  if constexpr (is_discrete_function_P0_v<DiscreteFunctionType>) {
    named_item_data_set.add(NamedItemData{name, discrete_function.cellValues()});
  } else {
    static_assert(is_discrete_function_P0_vector_v<DiscreteFunctionType>, "unexpected discrete function type");
    named_item_data_set.add(NamedItemData{name, discrete_function.cellArrays()});
  }
}

void
WriterBase::_checkConnectivity(
  const std::shared_ptr<const MeshVariant>& mesh_v,
  const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list) const
{
  Assert(named_discrete_data_list.size() > 0);

  std::shared_ptr<const IConnectivity> connectivity =
    std::visit([&](auto&& p_mesh) -> std::shared_ptr<const IConnectivity> { return p_mesh->shared_connectivity(); },
               mesh_v->variant());

  for (size_t i = 0; i < named_discrete_data_list.size(); ++i) {
    const auto& named_discrete_data = named_discrete_data_list[i];

    if (named_discrete_data->type() == INamedDiscreteData::Type::item_value) {
      const NamedItemValueVariant& named_item_value_variant =
        dynamic_cast<const NamedItemValueVariant&>(*named_discrete_data);

      std::visit(
        [&](auto&& item_value) {
          if (item_value.connectivity_ptr() != connectivity) {
            std::ostringstream error_msg;
            error_msg << "The variable " << rang::fgB::yellow << named_item_value_variant.name() << rang::fg::reset
                      << " is not defined on the provided same connectivity as the mesh\n";
            throw NormalError(error_msg.str());
          }
        },
        named_item_value_variant.itemValueVariant()->itemValue());
    }
  }
}

void
WriterBase::_checkSignature(
  const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list) const
{
  using NameTypeMap = std::map<std::string, std::string>;
  NameTypeMap name_type_map;

  for (auto named_discrete_data : named_discrete_data_list) {
    switch (named_discrete_data->type()) {
    case INamedDiscreteData::Type::item_value: {
      const NamedItemValueVariant& named_item_value = dynamic_cast<const NamedItemValueVariant&>(*named_discrete_data);
      std::visit(
        [&](auto&& item_value) {
          using ItemValueT = std::decay_t<decltype(item_value)>;
          using DataType   = std::decay_t<typename ItemValueT::data_type>;

          std::ostringstream type_name;
          type_name << "item_value(" << dataTypeName(ast_node_data_type_from<DataType>) << ')';

          name_type_map[named_discrete_data->name()] = type_name.str();
        },
        named_item_value.itemValueVariant()->itemValue());
      break;
    }
    case INamedDiscreteData::Type::item_array: {
      const NamedItemArrayVariant& named_item_array = dynamic_cast<const NamedItemArrayVariant&>(*named_discrete_data);
      std::visit(
        [&](auto&& item_value) {
          using ItemValueT = std::decay_t<decltype(item_value)>;
          using DataType   = std::decay_t<typename ItemValueT::data_type>;

          std::ostringstream type_name;
          type_name << "item_array(" << dataTypeName(ast_node_data_type_from<DataType>) << ')';

          name_type_map[named_discrete_data->name()] = type_name.str();
        },
        named_item_array.itemArrayVariant()->itemArray());
      break;
    }
    case INamedDiscreteData::Type::discrete_function: {
      const NamedDiscreteFunction& named_discrete_function =
        dynamic_cast<const NamedDiscreteFunction&>(*named_discrete_data);
      std::visit(
        [&](auto&& discrete_function) {
          std::ostringstream type_name;
          type_name << "Vh(" << dataTypeName(discrete_function.dataType()) << ')';

          name_type_map[named_discrete_data->name()] = type_name.str();
        },
        named_discrete_function.discreteFunctionVariant()->discreteFunction());
      break;
    }
    }
  }

  std::string signature{"|"};
  for (auto&& [name, type_name] : name_type_map) {
    signature += name + std::string{"-"} + type_name + std::string{"|"};
  }

  if (m_signature.has_value()) {
    if (m_signature.value() != signature) {
      throw NormalError("output variable list changed");
    }
  } else {
    m_signature = signature;
  }
}

void
WriterBase::_checkMesh(const std::shared_ptr<const MeshVariant>& mesh_v,
                       const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list) const
{
  Assert(named_discrete_data_list.size() > 0);

  const size_t mesh_id = std::visit([](auto&& p_mesh) { return p_mesh->id(); }, mesh_v->variant());

  for (size_t i = 0; i < named_discrete_data_list.size(); ++i) {
    const auto& named_discrete_data = named_discrete_data_list[i];

    if (named_discrete_data->type() == INamedDiscreteData::Type::discrete_function) {
      const NamedDiscreteFunction& named_discrete_function =
        dynamic_cast<const NamedDiscreteFunction&>(*named_discrete_data);

      const size_t discrete_function_mesh_id =
        std::visit([](auto&& f) { return f.meshVariant()->id(); },
                   named_discrete_function.discreteFunctionVariant()->discreteFunction());

      if (mesh_id != discrete_function_mesh_id) {
        std::ostringstream error_msg;
        error_msg << "The variable " << rang::fgB::yellow << named_discrete_function.name() << rang::fg::reset
                  << " is not defined on the provided mesh\n";
        throw NormalError(error_msg.str());
      }
    }
  }
}

std::shared_ptr<const MeshVariant>
WriterBase::_getMesh(const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list) const
{
  Assert(named_discrete_data_list.size() > 0);

  std::shared_ptr<const MeshVariant> mesh_v;

  std::map<size_t, std::string> mesh_id_to_function_name_map;
  std::map<std::shared_ptr<const IConnectivity>, std::string> connectivity_set;

  for (size_t i = 0; i < named_discrete_data_list.size(); ++i) {
    const auto& named_discrete_data = named_discrete_data_list[i];

    switch (named_discrete_data->type()) {
    case INamedDiscreteData::Type::discrete_function: {
      const NamedDiscreteFunction& named_discrete_function =
        dynamic_cast<const NamedDiscreteFunction&>(*named_discrete_data);

      std::visit(
        [&](auto&& f) {
          mesh_id_to_function_name_map[f.meshVariant()->id()] = named_discrete_function.name();
          if (not mesh_v) {
            mesh_v = f.meshVariant();
          }
          connectivity_set[f.meshVariant()->shared_connectivity()] = named_discrete_function.name();
        },
        named_discrete_function.discreteFunctionVariant()->discreteFunction());

      break;
    }
    case INamedDiscreteData::Type::item_value: {
      const NamedItemValueVariant& named_item_value_variant =
        dynamic_cast<const NamedItemValueVariant&>(*named_discrete_data);

      std::visit([&](
                   auto&&
                     item_value) { connectivity_set[item_value.connectivity_ptr()] = named_item_value_variant.name(); },
                 named_item_value_variant.itemValueVariant()->itemValue());
      break;
    }
    case INamedDiscreteData::Type::item_array: {
      const NamedItemArrayVariant& named_item_array_variant =
        dynamic_cast<const NamedItemArrayVariant&>(*named_discrete_data);

      std::visit([&](
                   auto&&
                     item_array) { connectivity_set[item_array.connectivity_ptr()] = named_item_array_variant.name(); },
                 named_item_array_variant.itemArrayVariant()->itemArray());
      break;
    }
    }
  }

  if (mesh_id_to_function_name_map.size() != 1) {
    if (mesh_id_to_function_name_map.size() == 0) {
      throw NormalError("cannot find any mesh associated to output quantities");
    } else {
      std::ostringstream error_msg;
      error_msg << "cannot save data using different " << rang::fgB::red << "meshes" << rang::fg::reset
                << " in the same file!\n";
      error_msg << rang::fgB::yellow << "note:" << rang::fg::reset
                << "the following variables are defined on different meshes:";
      for (const auto& [mesh, name] : mesh_id_to_function_name_map) {
        error_msg << "\n- " << name;
      }
      throw NormalError(error_msg.str());
    }
  }

  if (connectivity_set.size() > 1) {
    std::ostringstream error_msg;
    error_msg << "cannot save data using different " << rang::fgB::red << "connectivities" << rang::fg::reset
              << " in the same file!\n";
    error_msg << rang::fgB::yellow << "note:" << rang::fg::reset
              << "the following variables are defined on different connectivities:";
    for (const auto& [connectivity, name] : connectivity_set) {
      error_msg << "\n- " << name;
    }
    throw NormalError(error_msg.str());
  }

  return mesh_v;
}

OutputNamedItemDataSet
WriterBase::_getOutputNamedItemDataSet(
  const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list) const
{
  OutputNamedItemDataSet named_item_data_set;

  for (auto& named_discrete_data : named_discrete_data_list) {
    switch (named_discrete_data->type()) {
    case INamedDiscreteData::Type::discrete_function: {
      const NamedDiscreteFunction& named_discrete_function =
        dynamic_cast<const NamedDiscreteFunction&>(*named_discrete_data);

      const std::string& name = named_discrete_function.name();

      const DiscreteFunctionVariant& discrete_function_variant = *named_discrete_function.discreteFunctionVariant();

      std::visit([&](auto&& f) { WriterBase::_registerDiscreteFunction(name, f, named_item_data_set); },
                 discrete_function_variant.discreteFunction());
      break;
    }
    case INamedDiscreteData::Type::item_value: {
      const NamedItemValueVariant& named_item_value_variant =
        dynamic_cast<const NamedItemValueVariant&>(*named_discrete_data);

      const std::string& name = named_item_value_variant.name();

      const ItemValueVariant& item_value_variant = *named_item_value_variant.itemValueVariant();

      std::visit([&](auto&& item_value) { named_item_data_set.add(NamedItemData{name, item_value}); },
                 item_value_variant.itemValue());
      break;
    }
    case INamedDiscreteData::Type::item_array: {
      const NamedItemArrayVariant& named_item_array_variant =
        dynamic_cast<const NamedItemArrayVariant&>(*named_discrete_data);

      const std::string& name = named_item_array_variant.name();

      const ItemArrayVariant& item_value_variant = *named_item_array_variant.itemArrayVariant();

      std::visit(
        [&](auto&& item_array) {
          using ItemArrayType = std::decay_t<decltype(item_array)>;
          using DataType      = std::decay_t<typename ItemArrayType::data_type>;
          if constexpr (std::is_arithmetic_v<DataType>) {
            named_item_data_set.add(NamedItemData{name, item_array});
          } else {
            throw NormalError("can only write item_array containing scalar values");
          }
        },
        item_value_variant.itemArray());
      break;
    }
    }
  }
  return named_item_data_set;
}

WriterBase::WriterBase(const std::string& base_filename, const double& time_period)
  : m_base_filename{base_filename}, m_period_manager(time_period), m_signature(std::nullopt)
{}

WriterBase::WriterBase(const std::string& base_filename)
  : m_base_filename{base_filename}, m_period_manager{std::nullopt}, m_signature(std::nullopt)
{}

void
WriterBase::write(const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list) const
{
  if (m_period_manager.has_value()) {
    throw NormalError("this writer requires time value");
  } else {
    std::shared_ptr<const MeshVariant> mesh = _getMesh(named_discrete_data_list);
    this->_write(*mesh, named_discrete_data_list);
  }
}

void
WriterBase::writeIfNeeded(const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list,
                          double time) const
{
  if (m_period_manager.has_value()) {
    const double last_time = m_period_manager->getLastTime();
    if (time == last_time)
      return;   // output already performed

    if (time >= m_period_manager->nextTime()) {
      std::shared_ptr<const MeshVariant> mesh = _getMesh(named_discrete_data_list);
      this->_checkSignature(named_discrete_data_list);
      this->_writeAtTime(*mesh, named_discrete_data_list, time);
      m_period_manager->setSaveTime(time);
    }
  } else {
    throw NormalError("this writer does not allow time value");
  }
}

void
WriterBase::writeForced(const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list,
                        double time) const
{
  if (m_period_manager.has_value()) {
    if (time == m_period_manager->getLastTime())
      return;   // output already performed

    std::shared_ptr<const MeshVariant> mesh = _getMesh(named_discrete_data_list);
    this->_checkSignature(named_discrete_data_list);
    this->_writeAtTime(*mesh, named_discrete_data_list, time);
    m_period_manager->setSaveTime(time);
  } else {
    throw NormalError("this writer does not allow time value");
  }
}

void
WriterBase::writeOnMesh(const std::shared_ptr<const MeshVariant>& mesh,
                        const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list) const
{
  if (m_period_manager.has_value()) {
    throw NormalError("this writer requires time value");
  } else {
    this->_checkMesh(mesh, named_discrete_data_list);
    this->_checkConnectivity(mesh, named_discrete_data_list);
    this->_write(*mesh, named_discrete_data_list);
  }
}

void
WriterBase::writeOnMeshIfNeeded(const std::shared_ptr<const MeshVariant>& mesh,
                                const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list,
                                double time) const
{
  if (m_period_manager.has_value()) {
    if (time == m_period_manager->getLastTime()) {
      return;   // output already performed
    }

    this->_checkMesh(mesh, named_discrete_data_list);
    this->_checkConnectivity(mesh, named_discrete_data_list);
    this->_writeAtTime(*mesh, named_discrete_data_list, time);
    m_period_manager->setSaveTime(time);
  } else {
    throw NormalError("this writer does not allow time value");
  }
}

void
WriterBase::writeOnMeshForced(const std::shared_ptr<const MeshVariant>& mesh,
                              const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list,
                              double time) const
{
  if (m_period_manager.has_value()) {
    if (time == m_period_manager->getLastTime()) {
      return;   // output already performed
    }
    this->_checkMesh(mesh, named_discrete_data_list);
    this->_checkConnectivity(mesh, named_discrete_data_list);
    this->_writeAtTime(*mesh, named_discrete_data_list, time);
    m_period_manager->setSaveTime(time);
  } else {
    throw NormalError("this writer does not allow time value");
  }
}

void
WriterBase::writeMesh(const std::shared_ptr<const MeshVariant>& mesh) const
{
  writeMesh(*mesh);
}

void
WriterBase::writeMesh(const MeshVariant& mesh) const
{
  if (m_period_manager.has_value()) {
    throw NormalError("write_mesh requires a writer without time period");
  } else {
    this->_writeMesh(mesh);
  }
}