stroid/tools/stroid.cpp

#include <iostream>
#include <fstream>
#include <string>
#include <map>
#include <memory>
#include <optional>
#include <ranges>
#include <algorithm>
#include <print>

// ReSharper disable once CppUnusedIncludeDirective
#include "mfem.hpp"
#include "stroid/stroid.h"
#include "fourdst/config/config.h"
#include "CLI/CLI.hpp"
#include "magic_enum/magic_enum.hpp"

bool validate_filename_extension(const std::string& filename, const std::string& expected_ext) {
    const auto pos = filename.rfind('.');
    if (pos == std::string::npos) return false;
    return filename.substr(pos + 1) == expected_ext;
}

enum class MESH_FORMATS : int {
    VTK,
    VTU,
    MFEM,
    NETGEN,
    INFO,
    PARAVIEW
};

struct OUTPUT_CONFIG {
    struct VTK_SETTINGS {
        int ref = 0;
        int field_data = 0;
    } vtk;

    struct VTU_SETTINGS {
        int ref = 1;
        mfem::VTKFormat format = mfem::VTKFormat::ASCII;
        bool high_order_output = true;
        int compression_level = 0;
        bool bdr_elements = false;
    } vtu;

    struct MFEM_SETTINGS {
        std::string comments = "Generated by stroid";
    } mfem;
};

int main(int argc, char** argv) {
    fourdst::config::Config<stroid::config::MeshConfig> cfg;
    OUTPUT_CONFIG out_cfg;
    auto selected_format = MESH_FORMATS::MFEM;

    CLI::App app{"stroid - A tool for generating multi-block meshes for stellar modeling"};

    app.footer(
      "\nEXAMPLES:\n"
      "|   stroid generate -c config.toml -o star_mesh.vtu vtu\n"
      "|   stroid generate --config config.toml --view --no-save\n"
      "|   stroid info --version\n"
      "|   stroid info --default\n"
    );

    auto* generate = app.add_subcommand("generate", "Generate a multi-block mesh");
    auto* info = app.add_subcommand("info", "Access information about stroid");

    std::optional<std::string> config_filename;
    std::string output_filename = "stroid.mesh";
    bool view_mesh = false;
    bool no_save = false;
    std::string glvis_host = "localhost";
    int glvis_port = 19916;

    generate->add_option("-c,--config", config_filename, "Path to configuration file")->check(CLI::ExistingFile);
    generate->add_flag("-v,--view", view_mesh, "View the generated mesh using GLVis");
    generate->add_flag("-n,--no-save", no_save, "Do not save the generated mesh to a file");
    generate->add_option("--glvis-host", glvis_host, "GLVis server host")->capture_default_str();
    generate->add_option("--glvis-port", glvis_port, "GLVis server port")->capture_default_str();
    generate->add_option("-o,--output", output_filename, "Output filename base")->capture_default_str();

    for (auto [value, name_view] : magic_enum::enum_entries<MESH_FORMATS>()) {
        std::string name{name_view};
        std::ranges::transform(name, name.begin(), ::tolower);

        auto* sub = generate->add_subcommand(name, "Output as " + name);

        switch (value) {
            case MESH_FORMATS::VTK:
                sub->add_option("--ref", out_cfg.vtk.ref, "Refinement level");
                sub->add_option("--field-data", out_cfg.vtk.field_data, "Field data flag");
                break;
            case MESH_FORMATS::VTU:
                sub->add_option("--ref", out_cfg.vtu.ref, "Refinement level");
                sub->add_option("--compression", out_cfg.vtu.compression_level, "Compression level (0-9)");
                sub->add_flag("--high-order", out_cfg.vtu.high_order_output, "High order output");
                break;
            case MESH_FORMATS::MFEM:
                sub->add_option("--comments", out_cfg.mfem.comments, "Header comments");
                break;
            default: break;
        }

        sub->callback([&selected_format, value]() {
            selected_format = value;
        });
    }

    // generate->require_subcommand(1);

    info->add_flag_callback("-v,--version", []() {
        std::println("Stroid Version {}", stroid::version::toString());
    }, "Display stroid version information");

    info->add_flag_callback("-d,--default", [&cfg]() {
        cfg.save("default.toml");
        std::println("Default configuration saved to default.toml");
    }, "Save the default configuration to default.toml");

    try {
        app.parse(argc, argv);
    } catch (const CLI::ParseError &e) {
        return app.exit(e);
    }

    if (*generate) {
        if (config_filename.has_value()) {
            cfg.load(config_filename.value());
        }

        const std::unique_ptr<mfem::Mesh> mesh = stroid::topology::BuildSkeleton(cfg);
        stroid::topology::Finalize(*mesh, cfg);
        stroid::topology::PromoteToHighOrder(*mesh, cfg);
        stroid::topology::ProjectMesh(*mesh, cfg);

        if (!no_save) {
            const std::string& final_path = output_filename;

            switch (selected_format) {
                case MESH_FORMATS::MFEM: {
                    if (!validate_filename_extension(final_path, "mesh")) {
                        std::cerr << "WARNING! Saving to MFEM format without the standard '.mesh' extension. File will be called " << final_path << std::endl;
                    }
                    std::ofstream ofs(final_path);
                    mesh->Print(ofs);
                    break;
                }
                case MESH_FORMATS::VTU: {
                    if (!validate_filename_extension(final_path, "vtu")) {
                        std::cerr << "WARNING! Saving to VTU format without the standard '.vtu' extension. File will be called " << final_path << std::endl;
                    }
                    std::ofstream ofs(final_path);
                    mesh->PrintVTU(ofs,
                                        out_cfg.vtu.ref,
                                        out_cfg.vtu.format,
                                        out_cfg.vtu.high_order_output,
                                        out_cfg.vtu.compression_level,
                                        out_cfg.vtu.bdr_elements);
                    break;
                }
                case MESH_FORMATS::VTK: {
                    if (!validate_filename_extension(final_path, "vtk")) {
                        std::cerr << "WARNING! Saving to VTK format without the standard '.vtk' extension. File will be called " << final_path << std::endl;
                    }
                    std::ofstream ofs(final_path);
                    mesh->PrintVTK(ofs, out_cfg.vtk.ref, out_cfg.vtk.field_data);
                    break;
                }
                case MESH_FORMATS::NETGEN: {
                    if (!validate_filename_extension(final_path, "nmesh")) {
                        std::cerr << "WARNING! Saving to Netgen format without the standard '.nmesh' extension. File will be called " << final_path << std::endl;
                    }
                    std::ofstream ofs(final_path);
                    mesh->PrintXG(ofs);
                    break;
                }
                case MESH_FORMATS::INFO: {
                    mesh->PrintCharacteristics();
                }
                case MESH_FORMATS::PARAVIEW: {
                    stroid::IO::SaveVTU(*mesh, final_path);
                }
            }
        }

        if (view_mesh) {
            stroid::IO::ViewMesh(*mesh,
                                 "Spheroidal Mesh - Colored by Element ID",
                                 stroid::IO::VISUALIZATION_MODE::ELEMENT_ID,
                                 glvis_host,
                                 glvis_port);
        }
    } else if (!*info) {
        std::println("Usage: {} [generate|info] --help", argv[0]);
    }

    return 0;
}