4 files changed, 98 insertions, 19 deletions

diff --git a/src/openvic-simulation/dataloader/NodeTools.hpp b/src/openvic-simulation/dataloader/NodeTools.hpp
index 51e3e82..43c5092 100644
--- a/src/openvic-simulation/dataloader/NodeTools.hpp
+++ b/src/openvic-simulation/dataloader/NodeTools.hpp
@@ -287,11 +287,14 @@ namespace OpenVic {
 			ret &= add_key_map_entries(FWD(key_map), FWD(args)...);
 			return ret;
 		}
+
 		template<IsOrderedMap Map>
 		NodeCallback auto expect_dictionary_key_map_and_length_and_default(
 			Map&& key_map, LengthCallback auto&& length_callback, KeyValueCallback auto&& default_callback
 		) {
-			return [length_callback = FWD(length_callback), default_callback = FWD(default_callback), key_map = MOV(key_map)](ast::NodeCPtr node) mutable -> bool {
+			return [length_callback = FWD(length_callback), default_callback = FWD(default_callback), key_map = MOV(key_map)](
+				ast::NodeCPtr node
+			) mutable -> bool {
 				bool ret = expect_dictionary_and_length(
 					FWD(length_callback), dictionary_keys_callback(key_map, FWD(default_callback))
 				)(node);
@@ -335,6 +338,28 @@ namespace OpenVic {
 			return expect_dictionary_key_map_and_length_and_default(FWD(key_map), FWD(length_callback), FWD(default_callback));
 		}
 
+		template<IsOrderedMap Map, typename... Args>
+		NodeCallback auto expect_dictionary_key_map_and_length(
+			Map&& key_map, LengthCallback auto&& length_callback, Args&&... args
+		) {
+			add_key_map_entries(FWD(key_map), FWD(args)...);
+			return expect_dictionary_key_map_and_length(FWD(key_map), FWD(length_callback));
+		}
+
+		template<IsOrderedMap Map, typename... Args>
+		NodeCallback auto expect_dictionary_key_map_and_default(
+			Map&& key_map, KeyValueCallback auto&& default_callback, Args&&... args
+		) {
+			add_key_map_entries(FWD(key_map), FWD(args)...);
+			return expect_dictionary_key_map_and_default(FWD(key_map), FWD(default_callback));
+		}
+
+		template<IsOrderedMap Map, typename... Args>
+		NodeCallback auto expect_dictionary_key_map(Map&& key_map, Args&&... args) {
+			add_key_map_entries(FWD(key_map), FWD(args)...);
+			return expect_dictionary_key_map(FWD(key_map));
+		}
+
 		template<StringMapCase Case = StringMapCaseSensitive, typename... Args>
 		NodeCallback auto expect_dictionary_keys_and_length_and_default(
 			LengthCallback auto&& length_callback, KeyValueCallback auto&& default_callback, Args&&... args
diff --git a/src/openvic-simulation/misc/Event.cpp b/src/openvic-simulation/misc/Event.cpp
index 3c85afb..64b1f64 100644
--- a/src/openvic-simulation/misc/Event.cpp
+++ b/src/openvic-simulation/misc/Event.cpp
@@ -183,7 +183,7 @@ bool EventManager::load_event_file(IssueManager const& issue_manager, ast::NodeC
 					return ret;
 				},
 				"trigger", ZERO_OR_ONE, trigger.expect_script(),
-				"mean_time_to_happen", ZERO_OR_ONE, mean_time_to_happen.expect_conditional_weight(ConditionalWeight::MONTHS),
+				"mean_time_to_happen", ZERO_OR_ONE, mean_time_to_happen.expect_conditional_weight(ConditionalWeight::TIME),
 				"immediate", ZERO_OR_MORE, immediate.expect_script()
 			)(value);
 			ret &= register_event(
diff --git a/src/openvic-simulation/scripts/ConditionalWeight.cpp b/src/openvic-simulation/scripts/ConditionalWeight.cpp
index ca42f1b..1bb83d0 100644
--- a/src/openvic-simulation/scripts/ConditionalWeight.cpp
+++ b/src/openvic-simulation/scripts/ConditionalWeight.cpp
@@ -26,23 +26,86 @@ static NodeCallback auto expect_modifier(
 }
 
 node_callback_t ConditionalWeight::expect_conditional_weight(base_key_t base_key) {
-	return expect_dictionary_keys(
-		// TODO - add days and years as options with a shared expected count of ONE_EXACTLY
-		base_key_to_string(base_key), ZERO_OR_ONE, expect_fixed_point(assign_variable_callback(base)),
-		"days", ZERO_OR_ONE, success_callback,
-		"years", ZERO_OR_ONE, success_callback,
+	key_map_t key_map;
+	bool successfully_set_up_base_keys = true;
+
+	switch (base_key) {
+	case BASE:
+		{
+			successfully_set_up_base_keys &= add_key_map_entry(
+				key_map,
+				"base", ONE_EXACTLY, expect_fixed_point(assign_variable_callback(base))
+			);
+
+			break;
+		}
+	case FACTOR:
+		{
+			successfully_set_up_base_keys &= add_key_map_entry(
+				key_map,
+				"factor", ONE_EXACTLY, expect_fixed_point(assign_variable_callback(base))
+			);
+
+			break;
+		}
+	case TIME:
+		{
+			const auto time_callback = [this](std::string_view key, Timespan (*to_timespan)(Timespan::day_t)) -> auto {
+				return [this, key, to_timespan](uint32_t value) -> bool {
+					if (base == fixed_point_t::_0()) {
+						base = fixed_point_t::parse((*to_timespan)(value).to_int());
+						return true;
+					} else {
+						Logger::error(
+							"ConditionalWeight cannot have multiple base values - trying to set base to ", value, " ", key,
+							" when it already has a value equivalent to ", base, " days!"
+						);
+						return false;
+					}
+				};
+			};
+
+			successfully_set_up_base_keys &= add_key_map_entries(
+				key_map,
+				"days", ZERO_OR_ONE, expect_uint<uint32_t>(time_callback("days", Timespan::from_days)),
+				"months", ZERO_OR_ONE, expect_uint<uint32_t>(time_callback("months", Timespan::from_months)),
+				"years", ZERO_OR_ONE, expect_uint<uint32_t>(time_callback("years", Timespan::from_years))
+			);
+
+			break;
+		}
+	default:
+		{
+			successfully_set_up_base_keys = false;
+		}
+	}
+
+	if (!successfully_set_up_base_keys) {
+		return [base_key](ast::NodeCPtr node) -> bool {
+			Logger::error(
+				"Failed to set up base keys for ConditionalWeight with base value: ", static_cast<uint32_t>(base_key)
+			);
+			return false;
+		};
+	}
+
+	return expect_dictionary_key_map(
+		std::move(key_map),
 		"modifier", ZERO_OR_MORE, expect_modifier(condition_weight_items, initial_scope, this_scope, from_scope),
 		"group", ZERO_OR_MORE, [this](ast::NodeCPtr node) -> bool {
 			condition_weight_group_t items;
+
 			const bool ret = expect_dictionary_keys(
 				"modifier", ONE_OR_MORE, expect_modifier(items, initial_scope, this_scope, from_scope)
 			)(node);
+
 			if (!items.empty()) {
 				condition_weight_items.emplace_back(std::move(items));
 				return ret;
+			} else {
+				Logger::error("ConditionalWeight group must have at least one modifier!");
+				return false;
 			}
-			Logger::error("ConditionalWeight group must have at least one modifier!");
-			return false;
 		}
 	);
 }
diff --git a/src/openvic-simulation/scripts/ConditionalWeight.hpp b/src/openvic-simulation/scripts/ConditionalWeight.hpp
index f667cd7..fcab0a6 100644
--- a/src/openvic-simulation/scripts/ConditionalWeight.hpp
+++ b/src/openvic-simulation/scripts/ConditionalWeight.hpp
@@ -14,7 +14,7 @@ namespace OpenVic {
 		using condition_weight_item_t = std::variant<condition_weight_t, condition_weight_group_t>;
 
 		enum class base_key_t : uint8_t {
-			BASE, FACTOR, MONTHS
+			BASE, FACTOR, TIME
 		};
 		using enum base_key_t;
 
@@ -31,15 +31,6 @@ namespace OpenVic {
 		ConditionalWeight(scope_type_t new_initial_scope, scope_type_t new_this_scope, scope_type_t new_from_scope);
 		ConditionalWeight(ConditionalWeight&&) = default;
 
-		static constexpr std::string_view base_key_to_string(base_key_t base_key) {
-			switch (base_key) {
-				case base_key_t::BASE: return "base";
-				case base_key_t::FACTOR: return "factor";
-				case base_key_t::MONTHS: return "months"; // TODO - add functionality for days or months or years
-				default: return "INVALID BASE KEY";
-			}
-		}
-
 		NodeTools::node_callback_t expect_conditional_weight(base_key_t base_key);
 
 		bool parse_scripts(DefinitionManager const& definition_manager);