#include "io/hdf5/hdf5_model_writer.h"

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>

#include "fesa/build_info.h"

namespace fesa::hdf5_internal {

struct NodeWriteRow {
  std::uint64_t internal_node_id;
  const char* instance_name;
  const char* source_label;
  double coordinates[3];
};

struct ElementWriteRow {
  std::uint64_t internal_element_id;
  const char* instance_name;
  const char* source_label;
  std::uint64_t node_internal_ids[2];
  double local_axes[9];
};

struct ShellElementWriteRow {
  std::uint64_t internal_element_id;
  const char* instance_name;
  const char* source_label;
  const char* source_element_type;
  const char* internal_formulation;
  std::uint64_t node_internal_ids[4];
  std::uint64_t shell_section_internal_id;
  std::uint64_t material_internal_id;
};

struct ShellMaterialWriteRow {
  std::uint64_t internal_material_id;
  const char* name;
  double youngs_modulus;
  double poisson_ratio;
};

struct ShellSectionWriteRow {
  std::uint64_t internal_section_id;
  const char* source_file;
  std::uint64_t source_line;
  const char* source_elset;
  std::uint64_t material_internal_id;
  double thickness;
};

void WriteMetadata(const hid_t file, const Domain& domain) {
  auto metadata = CreateGroup(file, "/metadata");
  WriteUint64Attribute(metadata.Get(), "schema_version", 0U);
  WriteStringAttribute(metadata.Get(), "solver_version",
                       std::string{SolverVersion()});
  WriteStringAttribute(metadata.Get(), "source_input_identity",
                       SourceInputIdentity(domain));
  WriteStringAttribute(metadata.Get(), "unit_system_label",
                       "user-consistent-unspecified");
  if (IsShellDomain(domain)) {
    WriteStringAttribute(metadata.Get(), "feature_id",
                         "linear-static-mitc4-shell");
    WriteStringAttribute(
        metadata.Get(), "coordinate_convention",
        "global-cartesian; shell-local=(e1,e2,e3); positive-thickness=+zeta");
    WriteStringAttribute(metadata.Get(), "internal_formulation", "FESA-MITC4");
    WriteStringAttribute(metadata.Get(), "integration_rule",
                         "2x2x2-gauss; mitc4-edge-midpoint-shear");
  } else {
    WriteStringAttribute(metadata.Get(), "feature_id",
                         "linear-static-3d-euler-beam");
    WriteStringAttribute(metadata.Get(), "coordinate_convention",
                         "global-cartesian; beam-local=(t,n1,t-cross-n1)");
    WriteStringAttribute(metadata.Get(), "element_formulation",
                         "B33-3D-Euler-Bernoulli");
  }
  WriteStringAttribute(metadata.Get(), "step_name", kStepName);
  WriteUint64Attribute(metadata.Get(), "frame_index", 0U);
}

void WriteNodes(const hid_t file, const Domain& domain) {
  std::vector<NodeWriteRow> rows;
  rows.reserve(domain.Nodes().size());
  for (std::size_t index = 0U; index < domain.Nodes().size(); ++index) {
    const Node& node = domain.Nodes()[index];
    rows.push_back(
        {static_cast<std::uint64_t>(index),
         node.source_id.instance_name.c_str(),
         node.source_id.source_label_text.c_str(),
         {node.coordinates[0U], node.coordinates[1U], node.coordinates[2U]}});
  }

  auto string_type = MakeUtf8StringType();
  const hsize_t coordinate_dimensions[] = {3U};
  Hdf5Handle file_coordinates{
      RequireHdf5Id(H5Tarray_create2(H5T_IEEE_F64LE, 1, coordinate_dimensions),
                    "Unable to create the node coordinate file type."),
      H5Tclose};
  Hdf5Handle memory_coordinates{
      RequireHdf5Id(
          H5Tarray_create2(H5T_NATIVE_DOUBLE, 1, coordinate_dimensions),
          "Unable to create the node coordinate memory type."),
      H5Tclose};
  Hdf5Handle file_type{
      RequireHdf5Id(H5Tcreate(H5T_COMPOUND, sizeof(NodeWriteRow)),
                    "Unable to create the node file type."),
      H5Tclose};
  Hdf5Handle memory_type{
      RequireHdf5Id(H5Tcreate(H5T_COMPOUND, sizeof(NodeWriteRow)),
                    "Unable to create the node memory type."),
      H5Tclose};
  RequireHdf5(H5Tinsert(file_type.Get(), "internal_node_id",
                        HOFFSET(NodeWriteRow, internal_node_id), H5T_STD_U64LE),
              "Unable to define the node internal ID field.");
  RequireHdf5(
      H5Tinsert(file_type.Get(), "instance_name",
                HOFFSET(NodeWriteRow, instance_name), string_type.Get()),
      "Unable to define the node instance field.");
  RequireHdf5(H5Tinsert(file_type.Get(), "source_label",
                        HOFFSET(NodeWriteRow, source_label), string_type.Get()),
              "Unable to define the node source-label field.");
  RequireHdf5(
      H5Tinsert(file_type.Get(), "coordinates",
                HOFFSET(NodeWriteRow, coordinates), file_coordinates.Get()),
      "Unable to define the node coordinate field.");
  RequireHdf5(
      H5Tinsert(memory_type.Get(), "internal_node_id",
                HOFFSET(NodeWriteRow, internal_node_id), H5T_NATIVE_UINT64),
      "Unable to define the node internal ID memory field.");
  RequireHdf5(
      H5Tinsert(memory_type.Get(), "instance_name",
                HOFFSET(NodeWriteRow, instance_name), string_type.Get()),
      "Unable to define the node instance memory field.");
  RequireHdf5(H5Tinsert(memory_type.Get(), "source_label",
                        HOFFSET(NodeWriteRow, source_label), string_type.Get()),
              "Unable to define the node source-label memory field.");
  RequireHdf5(
      H5Tinsert(memory_type.Get(), "coordinates",
                HOFFSET(NodeWriteRow, coordinates), memory_coordinates.Get()),
      "Unable to define the node coordinate memory field.");

  auto dataset =
      WriteCompoundDataset(file, "/model/nodes", rows.size(), file_type.Get(),
                           memory_type.Get(), rows.data());
  WriteStringAttribute(dataset.Get(), "coordinate_system", "global-cartesian");
  WriteStringAttribute(dataset.Get(), "units_label", "length");
}

void WriteBeamElements(const hid_t file, const Domain& domain,
                       const std::vector<AxisSet>& axes) {
  std::vector<ElementWriteRow> rows;
  rows.reserve(domain.BeamElements().Size());
  for (std::size_t index = 0U; index < domain.BeamElements().Size(); ++index) {
    const auto& element = domain.BeamElements()[index];
    ElementWriteRow row{static_cast<std::uint64_t>(index),
                        element.source_id.instance_name.c_str(),
                        element.source_id.source_label_text.c_str(),
                        {static_cast<std::uint64_t>(element.node_indices[0U]),
                         static_cast<std::uint64_t>(element.node_indices[1U])},
                        {}};
    std::copy(axes[index].begin(), axes[index].end(), row.local_axes);
    rows.push_back(row);
  }

  auto string_type = MakeUtf8StringType();
  const hsize_t node_dimensions[] = {2U};
  const hsize_t axis_dimensions[] = {3U, 3U};
  Hdf5Handle file_nodes{
      RequireHdf5Id(H5Tarray_create2(H5T_STD_U64LE, 1, node_dimensions),
                    "Unable to create the element connectivity file type."),
      H5Tclose};
  Hdf5Handle memory_nodes{
      RequireHdf5Id(H5Tarray_create2(H5T_NATIVE_UINT64, 1, node_dimensions),
                    "Unable to create the element connectivity memory type."),
      H5Tclose};
  Hdf5Handle file_axes{
      RequireHdf5Id(H5Tarray_create2(H5T_IEEE_F64LE, 2, axis_dimensions),
                    "Unable to create the local-axis file type."),
      H5Tclose};
  Hdf5Handle memory_axes{
      RequireHdf5Id(H5Tarray_create2(H5T_NATIVE_DOUBLE, 2, axis_dimensions),
                    "Unable to create the local-axis memory type."),
      H5Tclose};
  Hdf5Handle file_type{
      RequireHdf5Id(H5Tcreate(H5T_COMPOUND, sizeof(ElementWriteRow)),
                    "Unable to create the element file type."),
      H5Tclose};
  Hdf5Handle memory_type{
      RequireHdf5Id(H5Tcreate(H5T_COMPOUND, sizeof(ElementWriteRow)),
                    "Unable to create the element memory type."),
      H5Tclose};
  const auto insert_fields = [&](const hid_t type, const hid_t integer_type,
                                 const hid_t node_type, const hid_t axes_type) {
    RequireHdf5(
        H5Tinsert(type, "internal_element_id",
                  HOFFSET(ElementWriteRow, internal_element_id), integer_type),
        "Unable to define the element internal ID field.");
    RequireHdf5(
        H5Tinsert(type, "instance_name",
                  HOFFSET(ElementWriteRow, instance_name), string_type.Get()),
        "Unable to define the element instance field.");
    RequireHdf5(
        H5Tinsert(type, "source_label", HOFFSET(ElementWriteRow, source_label),
                  string_type.Get()),
        "Unable to define the element source-label field.");
    RequireHdf5(
        H5Tinsert(type, "node_internal_ids",
                  HOFFSET(ElementWriteRow, node_internal_ids), node_type),
        "Unable to define the element connectivity field.");
    RequireHdf5(H5Tinsert(type, "local_axes",
                          HOFFSET(ElementWriteRow, local_axes), axes_type),
                "Unable to define the element local-axis field.");
  };
  insert_fields(file_type.Get(), H5T_STD_U64LE, file_nodes.Get(),
                file_axes.Get());
  insert_fields(memory_type.Get(), H5T_NATIVE_UINT64, memory_nodes.Get(),
                memory_axes.Get());
  auto dataset =
      WriteCompoundDataset(file, "/model/elements", rows.size(),
                           file_type.Get(), memory_type.Get(), rows.data());
  WriteStringAttribute(dataset.Get(), "formulation", "B33-3D-Euler-Bernoulli");
}

void WriteShellElements(const hid_t file, const Domain& domain) {
  std::vector<ShellElementWriteRow> rows;
  rows.reserve(domain.ShellElements().Size());
  for (std::size_t index = 0U; index < domain.ShellElements().Size(); ++index) {
    const auto& element = domain.ShellElements()[index];
    rows.push_back({static_cast<std::uint64_t>(index),
                    element.source_id.instance_name.c_str(),
                    element.source_id.source_label_text.c_str(),
                    ShellSourceTypeName(element.source_type),
                    kMitc4InternalFormulation.data(),
                    {static_cast<std::uint64_t>(element.node_indices[0U]),
                     static_cast<std::uint64_t>(element.node_indices[1U]),
                     static_cast<std::uint64_t>(element.node_indices[2U]),
                     static_cast<std::uint64_t>(element.node_indices[3U])},
                    static_cast<std::uint64_t>(element.section_index),
                    static_cast<std::uint64_t>(element.material_index)});
  }

  auto string_type = MakeUtf8StringType();
  const hsize_t node_dimensions[] = {kShellNodeCount};
  Hdf5Handle file_nodes{
      RequireHdf5Id(H5Tarray_create2(H5T_STD_U64LE, 1, node_dimensions),
                    "Unable to create shell connectivity file type."),
      H5Tclose};
  Hdf5Handle memory_nodes{
      RequireHdf5Id(H5Tarray_create2(H5T_NATIVE_UINT64, 1, node_dimensions),
                    "Unable to create shell connectivity memory type."),
      H5Tclose};
  Hdf5Handle file_type{
      RequireHdf5Id(H5Tcreate(H5T_COMPOUND, sizeof(ShellElementWriteRow)),
                    "Unable to create shell element file type."),
      H5Tclose};
  Hdf5Handle memory_type{
      RequireHdf5Id(H5Tcreate(H5T_COMPOUND, sizeof(ShellElementWriteRow)),
                    "Unable to create shell element memory type."),
      H5Tclose};
  const auto insert_fields = [&](const hid_t type, const hid_t integer_type,
                                 const hid_t node_type) {
    RequireHdf5(H5Tinsert(type, "internal_element_id",
                          HOFFSET(ShellElementWriteRow, internal_element_id),
                          integer_type),
                "Unable to define shell element ID field.");
    RequireHdf5(H5Tinsert(type, "instance_name",
                          HOFFSET(ShellElementWriteRow, instance_name),
                          string_type.Get()),
                "Unable to define shell element instance field.");
    RequireHdf5(H5Tinsert(type, "source_label",
                          HOFFSET(ShellElementWriteRow, source_label),
                          string_type.Get()),
                "Unable to define shell element source-label field.");
    RequireHdf5(H5Tinsert(type, "source_element_type",
                          HOFFSET(ShellElementWriteRow, source_element_type),
                          string_type.Get()),
                "Unable to define shell source-type field.");
    RequireHdf5(H5Tinsert(type, "internal_formulation",
                          HOFFSET(ShellElementWriteRow, internal_formulation),
                          string_type.Get()),
                "Unable to define shell formulation field.");
    RequireHdf5(
        H5Tinsert(type, "node_internal_ids",
                  HOFFSET(ShellElementWriteRow, node_internal_ids), node_type),
        "Unable to define shell connectivity field.");
    RequireHdf5(
        H5Tinsert(type, "shell_section_internal_id",
                  HOFFSET(ShellElementWriteRow, shell_section_internal_id),
                  integer_type),
        "Unable to define shell section ID field.");
    RequireHdf5(H5Tinsert(type, "material_internal_id",
                          HOFFSET(ShellElementWriteRow, material_internal_id),
                          integer_type),
                "Unable to define shell material ID field.");
  };
  insert_fields(file_type.Get(), H5T_STD_U64LE, file_nodes.Get());
  insert_fields(memory_type.Get(), H5T_NATIVE_UINT64, memory_nodes.Get());
  auto dataset =
      WriteCompoundDataset(file, "/model/elements", rows.size(),
                           file_type.Get(), memory_type.Get(), rows.data());
  WriteStringAttribute(dataset.Get(), "formulation", "FESA-MITC4");
}

void WriteShellMaterials(const hid_t file, const Domain& domain) {
  std::vector<ShellMaterialWriteRow> rows;
  rows.reserve(domain.LinearElasticMaterials().Size());
  for (std::size_t index = 0U; index < domain.LinearElasticMaterials().Size();
       ++index) {
    const auto& material = domain.LinearElasticMaterials()[index];
    rows.push_back({static_cast<std::uint64_t>(index), material.name.c_str(),
                    material.youngs_modulus, material.poisson_ratio});
  }
  auto string_type = MakeUtf8StringType();
  Hdf5Handle file_type{
      RequireHdf5Id(H5Tcreate(H5T_COMPOUND, sizeof(ShellMaterialWriteRow)),
                    "Unable to create shell material file type."),
      H5Tclose};
  Hdf5Handle memory_type{
      RequireHdf5Id(H5Tcreate(H5T_COMPOUND, sizeof(ShellMaterialWriteRow)),
                    "Unable to create shell material memory type."),
      H5Tclose};
  const auto insert_fields = [&](const hid_t type, const hid_t integer_type,
                                 const hid_t double_type) {
    RequireHdf5(H5Tinsert(type, "internal_material_id",
                          HOFFSET(ShellMaterialWriteRow, internal_material_id),
                          integer_type),
                "Unable to define shell material ID field.");
    RequireHdf5(H5Tinsert(type, "name", HOFFSET(ShellMaterialWriteRow, name),
                          string_type.Get()),
                "Unable to define shell material name field.");
    RequireHdf5(
        H5Tinsert(type, "E", HOFFSET(ShellMaterialWriteRow, youngs_modulus),
                  double_type),
        "Unable to define shell material E field.");
    RequireHdf5(
        H5Tinsert(type, "nu", HOFFSET(ShellMaterialWriteRow, poisson_ratio),
                  double_type),
        "Unable to define shell material nu field.");
  };
  insert_fields(file_type.Get(), H5T_STD_U64LE, H5T_IEEE_F64LE);
  insert_fields(memory_type.Get(), H5T_NATIVE_UINT64, H5T_NATIVE_DOUBLE);
  auto dataset =
      WriteCompoundDataset(file, "/model/shell/materials", rows.size(),
                           file_type.Get(), memory_type.Get(), rows.data());
  WriteStringAttribute(dataset.Get(), "component_unit_dimensions",
                       "force/length^2,1");
}

void WriteShellSections(const hid_t file, const Domain& domain) {
  std::vector<std::string> source_files;
  source_files.reserve(domain.ShellSections().Size());
  for (std::size_t index = 0U; index < domain.ShellSections().Size(); ++index) {
    const auto& section = domain.ShellSections()[index];
    source_files.push_back(NormalizedPathString(section.location.file));
  }
  std::vector<ShellSectionWriteRow> rows;
  rows.reserve(domain.ShellSections().Size());
  for (std::size_t index = 0U; index < domain.ShellSections().Size(); ++index) {
    const auto& section = domain.ShellSections()[index];
    rows.push_back({static_cast<std::uint64_t>(index),
                    source_files[index].c_str(),
                    static_cast<std::uint64_t>(section.location.line),
                    section.name.c_str(),
                    static_cast<std::uint64_t>(section.material_index),
                    section.thickness});
  }
  auto string_type = MakeUtf8StringType();
  Hdf5Handle file_type{
      RequireHdf5Id(H5Tcreate(H5T_COMPOUND, sizeof(ShellSectionWriteRow)),
                    "Unable to create shell section file type."),
      H5Tclose};
  Hdf5Handle memory_type{
      RequireHdf5Id(H5Tcreate(H5T_COMPOUND, sizeof(ShellSectionWriteRow)),
                    "Unable to create shell section memory type."),
      H5Tclose};
  const auto insert_fields = [&](const hid_t type, const hid_t integer_type,
                                 const hid_t double_type) {
    RequireHdf5(H5Tinsert(type, "internal_section_id",
                          HOFFSET(ShellSectionWriteRow, internal_section_id),
                          integer_type),
                "Unable to define shell section ID field.");
    RequireHdf5(H5Tinsert(type, "source_file",
                          HOFFSET(ShellSectionWriteRow, source_file),
                          string_type.Get()),
                "Unable to define shell section source-file field.");
    RequireHdf5(
        H5Tinsert(type, "source_line",
                  HOFFSET(ShellSectionWriteRow, source_line), integer_type),
        "Unable to define shell section source-line field.");
    RequireHdf5(H5Tinsert(type, "source_elset",
                          HOFFSET(ShellSectionWriteRow, source_elset),
                          string_type.Get()),
                "Unable to define shell section ELSET field.");
    RequireHdf5(H5Tinsert(type, "material_internal_id",
                          HOFFSET(ShellSectionWriteRow, material_internal_id),
                          integer_type),
                "Unable to define shell section material ID field.");
    RequireHdf5(
        H5Tinsert(type, "thickness", HOFFSET(ShellSectionWriteRow, thickness),
                  double_type),
        "Unable to define shell section thickness field.");
  };
  insert_fields(file_type.Get(), H5T_STD_U64LE, H5T_IEEE_F64LE);
  insert_fields(memory_type.Get(), H5T_NATIVE_UINT64, H5T_NATIVE_DOUBLE);
  auto dataset =
      WriteCompoundDataset(file, "/model/shell/sections", rows.size(),
                           file_type.Get(), memory_type.Get(), rows.data());
  WriteStringAttribute(dataset.Get(), "thickness_unit_dimension", "length");
  WriteStringAttribute(dataset.Get(), "layering", "centered-single-layer");
}

void WriteShellModelData(const hid_t file, const Domain& domain,
                         const WriterModelData& model_data) {
  std::vector<double> directors;
  directors.reserve(domain.Nodes().size() * 3U);
  std::vector<double> frames;
  frames.reserve(domain.Nodes().size() * 9U);
  for (const auto& frame : domain.ShellNodeInitialFrames()) {
    directors.insert(directors.end(), frame.director.begin(),
                     frame.director.end());
    frames.insert(frames.end(), frame.tangent_a.begin(), frame.tangent_a.end());
    frames.insert(frames.end(), frame.tangent_b.begin(), frame.tangent_b.end());
    frames.insert(frames.end(), frame.director.begin(), frame.director.end());
  }
  WriteModelDoubleDataset(file, "/model/shell/nodal_director",
                          {static_cast<hsize_t>(domain.Nodes().size()), 3U},
                          directors.data(), directors.size(), "D1,D2,D3",
                          "1,1,1", "global-cartesian", "nodal");
  WriteModelDoubleDataset(file, "/model/shell/nodal_frame",
                          {static_cast<hsize_t>(domain.Nodes().size()), 3U, 3U},
                          frames.data(), frames.size(), "X,Y,Z", "1,1,1",
                          "global-cartesian", "nodal-frame");
  {
    Hdf5Handle dataset{
        RequireHdf5Id(H5Dopen2(file, "/model/shell/nodal_frame", H5P_DEFAULT),
                      "Unable to reopen the nodal-frame dataset."),
        H5Dclose};
    WriteStringAttribute(dataset.Get(), "axis_names", "A,B,D");
  }

  WriteShellMaterials(file, domain);
  WriteShellSections(file, domain);
  const std::vector<hsize_t> nodal_dimensions{
      static_cast<hsize_t>(domain.Nodes().size()), kDofsPerNode};
  WriteUint8Dataset(file, "/model/nodal_constraint_mask", nodal_dimensions,
                    model_data.constraint_mask.data(),
                    model_data.constraint_mask.size());
  WriteModelDoubleDataset(
      file, "/model/prescribed_displacement", nodal_dimensions,
      model_data.prescribed_displacement.data(),
      model_data.prescribed_displacement.size(), "UX,UY,UZ,URX,URY,URZ",
      "length,length,length,radian,radian,radian", "global-cartesian",
      "nodal-prescribed-value");

  const double gauss = 1.0 / std::sqrt(3.0);
  const std::array<double, 8> locations{-gauss, -gauss, gauss,  -gauss,
                                        gauss,  gauss,  -gauss, gauss};
  WriteModelDoubleDataset(file, "/model/shell/midsurface_locations", {4U, 2U},
                          locations.data(), locations.size(), "XI,ETA", "1,1",
                          "shell-natural", "midsurface-location");
  const std::array<double, 3> section_positions{-1.0, 0.0, 1.0};
  WriteModelDoubleDataset(file, "/model/shell/section_positions", {3U},
                          section_positions.data(), section_positions.size(),
                          "ZETA", "1", "shell-natural", "section-position");
  {
    Hdf5Handle dataset{
        RequireHdf5Id(
            H5Dopen2(file, "/model/shell/section_positions", H5P_DEFAULT),
            "Unable to reopen shell section positions."),
        H5Dclose};
    WriteStringAttribute(dataset.Get(), "position_names", "BOTTOM,MIDDLE,TOP");
  }
}

void WriteModel(const hid_t file, const Domain& domain,
                const WriterModelData& model_data) {
  WriteMetadata(file, domain);
  WriteNodes(file, domain);
  if (IsShellDomain(domain)) {
    WriteShellElements(file, domain);
    WriteShellModelData(file, domain, model_data);
  } else {
    WriteBeamElements(file, domain, model_data.beam_local_axes);
  }
}

}  // namespace fesa::hdf5_internal