From 4d58990957f5536e7aff16bbe76c3fe3a23112bc Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Sat, 13 Mar 2021 17:03:14 -0600
Subject: [PATCH] Adjusting UWUW material IDs if needed.

---
 include/openmc/dagmc.h |  6 +++++-
 src/dagmc.cpp          | 34 ++++++++++++++++++++++------------
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/include/openmc/dagmc.h b/include/openmc/dagmc.h
index cf399c936..a908f9f6c 100644
--- a/include/openmc/dagmc.h
+++ b/include/openmc/dagmc.h
@@ -14,12 +14,15 @@ extern "C" const bool DAGMC_ENABLED;
 #include "openmc/position.h"
 #include "openmc/xml_interface.h"
 
+class UWUW;
+
 namespace openmc {
 
 namespace model {
   extern std::shared_ptr<moab::DagMC> DAG;
 }
 
+
 class DAGUniverse : public Universe {
 
 public:
@@ -28,7 +31,8 @@ public:
 
   void initialize(); //!< Sets up the DAGMC instance and OpenMC internals
 
-  std::map<int32_t, int32_t> read_uwuw_materials(); //!< Reads UWUW materials and returns an ID map
+  std::shared_ptr<UWUW>
+  read_uwuw_materials(); //!< Reads UWUW materials and returns an ID map
 
   // Data Members
   std::string filename_;
diff --git a/src/dagmc.cpp b/src/dagmc.cpp
index 8a54b1546..9f2991268 100644
--- a/src/dagmc.cpp
+++ b/src/dagmc.cpp
@@ -275,13 +275,10 @@ void DAGUniverse::initialize() {
   // --- Materials ---
 
   // read any UWUW materials from the file
-  auto uwuw_id_map = read_uwuw_materials();
-
-  // create uwuw instance
-  UWUW uwuw(filename_.c_str());
+  std::shared_ptr<UWUW> uwuw = read_uwuw_materials();
 
   // check for uwuw material definitions
-  bool using_uwuw = !uwuw.material_library.empty();
+  bool using_uwuw = !uwuw->material_library.empty();
 
   // notify user if UWUW materials are going to be used
   if (using_uwuw) {
@@ -355,9 +352,9 @@ void DAGUniverse::initialize() {
       if (using_uwuw) {
         // lookup material in uwuw if present
         std::string uwuw_mat = DMD.volume_material_property_data_eh[vol_handle];
-        if (uwuw.material_library.count(uwuw_mat) != 0) {
+        if (uwuw->material_library.count(uwuw_mat) != 0) {
           // Note: material numbers are set by UWUW
-          int mat_number = uwuw.material_library.get_material(uwuw_mat).metadata["mat_number"].asInt();
+          int mat_number = uwuw->material_library.get_material(uwuw_mat).metadata["mat_number"].asInt();
           c->material_.push_back(mat_number);
         } else {
           fatal_error(fmt::format("Material with value {} not found in the "
@@ -453,12 +450,25 @@ void DAGUniverse::initialize() {
   } // end surface loop
 }
 
-std::map<int32_t, int32_t>
+std::shared_ptr<UWUW>
 DAGUniverse::read_uwuw_materials() {
 
-  UWUW uwuw(filename_.c_str());
-  const auto& mat_lib = uwuw.material_library;
-  if (mat_lib.size() == 0) return {};
+  int32_t next_material_id = 0;
+  for (const auto& m : model::materials) {
+    next_material_id = std::max(m->id_, next_material_id);
+  }
+  next_material_id++;
+
+  auto uwuw = std::make_shared<UWUW>(filename_.c_str());
+  const auto& mat_lib = uwuw->material_library;
+  if (mat_lib.size() == 0) return uwuw;
+
+  // if we're using automatic IDs, update the UWUW material metadata
+  if (adjust_material_ids_) {
+    for (auto& mat : uwuw->material_library) {
+      mat.second->metadata["mat_number"] = std::to_string(next_material_id++);
+    }
+  }
 
   std::stringstream ss;
   ss << "<?xml version=\"1.0\"?>\n";
@@ -477,7 +487,7 @@ DAGUniverse::read_uwuw_materials() {
   for (pugi::xml_node material_node : root.children("material")) {
     model::materials.push_back(std::make_unique<Material>(material_node));
   }
-  return {};
+  return uwuw;
 }
 
 int32_t create_dagmc_universe(const std::string& filename) {