#include <output/WriterBase.hpp>

#include <mesh/IMesh.hpp>
#include <mesh/ItemArrayVariant.hpp>
#include <mesh/ItemValueVariant.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 IMesh>& mesh,
  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::shared_ptr<const IConnectivity> {
    switch (mesh->dimension()) {
    case 1: {
      return dynamic_cast<const Mesh<Connectivity<1>>&>(*mesh).shared_connectivity();
    }
    case 2: {
      return dynamic_cast<const Mesh<Connectivity<2>>&>(*mesh).shared_connectivity();
    }
    case 3: {
      return dynamic_cast<const Mesh<Connectivity<3>>&>(*mesh).shared_connectivity();
    }
    default: {
      throw UnexpectedError("invalid dimension");
    }
    }
  }();

  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 IMesh>& mesh,
                       const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list) const
{
  Assert(named_discrete_data_list.size() > 0);

  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);

      std::shared_ptr<const IMesh> discrete_function_mesh =
        std::visit([](auto&& f) { return f.mesh(); },
                   named_discrete_function.discreteFunctionVariant()->discreteFunction());

      if (mesh != discrete_function_mesh) {
        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 IMesh>
WriterBase::_getMesh(const std::vector<std::shared_ptr<const INamedDiscreteData>>& named_discrete_data_list) const
{
  Assert(named_discrete_data_list.size() > 0);

  std::map<std::shared_ptr<const IMesh>, std::string> mesh_set;
  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::shared_ptr mesh = std::visit([&](auto&& f) { return f.mesh(); },
                                        named_discrete_function.discreteFunctionVariant()->discreteFunction());
      mesh_set[mesh]       = named_discrete_function.name();

      switch (mesh->dimension()) {
      case 1: {
        connectivity_set[dynamic_cast<const Mesh<Connectivity<1>>&>(*mesh).shared_connectivity()] =
          named_discrete_function.name();
        break;
      }
      case 2: {
        connectivity_set[dynamic_cast<const Mesh<Connectivity<2>>&>(*mesh).shared_connectivity()] =
          named_discrete_function.name();
        break;
      }
      case 3: {
        connectivity_set[dynamic_cast<const Mesh<Connectivity<3>>&>(*mesh).shared_connectivity()] =
          named_discrete_function.name();
        break;
      }
      default: {
        throw UnexpectedError("invalid dimension");
      }
      }
      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_set.size() != 1) {
    if (mesh_set.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_set) {
        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_set.begin()->first;
}

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_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 IMesh> 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 IMesh> 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 IMesh> 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 IMesh>& 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 IMesh>& 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 IMesh>& 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 IMesh>& mesh) const
{
  writeMesh(*mesh);
}

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