From 951976680b5f2e5cb054f9ef7de9d62ddc55bcbe Mon Sep 17 00:00:00 2001
From: Katie Mummah <mummah@lanl.gov>
Date: Fri, 1 Mar 2024 20:54:14 -0700
Subject: [PATCH] add cumulative cap handling to storage

---
 CHANGELOG.rst        |  2 +-
 src/storage.cc       |  8 +++++++-
 src/storage.h        |  8 ++++++++
 src/storage_tests.cc | 36 ++++++++++++++++++++++++++++++++++++
 4 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 8e5826cae..663c69d09 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -10,7 +10,7 @@ cycamore Change Log
 * GitHub workflows for building/testing on a PR and push to `main` (#549, #564, #573, #582, #583)
 * Add functionality for random behavior on the size (#550) and frequency (#565) of a sink
 * GitHub workflow to check that the CHANGELOG has been updated (#562) 
-* Added inventory policies to Storage through the material buy policy (#574)
+* Added inventory policies to Storage through the material buy policy (#574, #588)
 
 **Changed:** 
 
diff --git a/src/storage.cc b/src/storage.cc
index e01317983..276326388 100644
--- a/src/storage.cc
+++ b/src/storage.cc
@@ -133,12 +133,18 @@ void Storage::EnterNotify() {
   cyclus::Facility::EnterNotify();
 
   inventory_tracker.set_capacity(max_inv_size);
-  if (reorder_point < 0) {
+  if (reorder_point < 0 && cumulative_cap <= 0) {
     InitBuyPolicyParameters();
     buy_policy.Init(this, &inventory, std::string("inventory"),
                     &inventory_tracker, throughput, active_dist_,
                     dormant_dist_, size_dist_);
   }
+  else if (cumulative_cap > 0) {
+    InitBuyPolicyParameters();
+    buy_policy.Init(this, &inventory, std::string("inventory"),
+                    &inventory_tracker, throughput, cumulative_cap,
+                    dormant_dist_);
+  }
   else if (reorder_quantity > 0) {
     if (reorder_point + reorder_quantity > max_inv_size) {
       throw cyclus::ValueError(
diff --git a/src/storage.h b/src/storage.h
index 7ac7002ae..0e8f0c7b5 100644
--- a/src/storage.h
+++ b/src/storage.h
@@ -421,6 +421,14 @@ class Storage
                       "uilabel":"Reorder Quantity"}
   double reorder_quantity;
 
+  #pragma cyclus var {"default": -1,\
+                      "tooltip": "Total amount of material that can be recieved per cycle.",\
+                      "doc": "After receiving this much material cumulatively, the agent will go dormant. "\
+                      "Must be paired with dormant_buying_frequency_type and any other dormant parameters. "\
+                      "The per-time step demand is unchanged except the cycle cap is almost reached.",\
+                      "uilabel": "Cumulative Cap"}
+  double cumulative_cap;
+
   #pragma cyclus var {"tooltip":"Incoming material buffer"}
   cyclus::toolkit::ResBuf<cyclus::Material> inventory;
 
diff --git a/src/storage_tests.cc b/src/storage_tests.cc
index 0edef93b8..134f712e9 100644
--- a/src/storage_tests.cc
+++ b/src/storage_tests.cc
@@ -833,6 +833,42 @@ TEST_F(StorageTest, sS_Inventory) {
   EXPECT_EQ(5, qr.GetVal<double>("Quantity", 0));
 }
 
+TEST_F(StorageTest, CCap_Inventory) {
+  std::string config =
+    "   <in_commods> <val>spent_fuel</val> </in_commods> "
+    "   <out_commods> <val>dry_spent</val> </out_commods> "
+    "   <throughput>1</throughput> "
+    "   <cumulative_cap>2</cumulative_cap> "
+    "   <dormant_buying_frequency_type>Fixed</dormant_buying_frequency_type> "
+    "   <dormant_buying_val>2</dormant_buying_val> ";
+
+  int simdur = 9;
+
+  cyclus::MockSim sim(cyclus::AgentSpec (":cycamore:Storage"), config, simdur);
+
+  sim.AddSource("spent_fuel").capacity(5).Finalize();
+  sim.AddSink("dry_spent").Finalize();
+
+  int id = sim.Run();
+
+  std::vector<cyclus::Cond> conds;
+  conds.push_back(cyclus::Cond("Commodity", "==", std::string("spent_fuel")));
+  cyclus::QueryResult qr = sim.db().Query("Transactions", &conds);
+  int n_trans = qr.rows.size();
+  EXPECT_EQ(5, n_trans) << "expected 5 transactions, got " << n_trans;
+  // check that the transactions occur at the expected time (0, 1, 4, 5, 8)
+  EXPECT_EQ(0, qr.GetVal<int>("Time", 0));
+  EXPECT_EQ(1, qr.GetVal<int>("Time", 1));
+  EXPECT_EQ(4, qr.GetVal<int>("Time", 2));
+  EXPECT_EQ(5, qr.GetVal<int>("Time", 3));
+  EXPECT_EQ(8, qr.GetVal<int>("Time", 4));
+
+  // check that transactions are of size 1
+  qr = sim.db().Query("Resources", NULL);
+  EXPECT_EQ(1, qr.GetVal<double>("Quantity", 0));
+  EXPECT_EQ(1, qr.GetVal<double>("Quantity", 4));
+}
+
 
 } // namespace cycamore