From 17f2b6ec7abc1faa95f8421b7274ede641a540f5 Mon Sep 17 00:00:00 2001
From: Huite Bootsma <huitebootsma@gmail.com>
Date: Fri, 17 Nov 2023 19:20:57 +0100
Subject: [PATCH] Allow outputting of multiple systems for exported_levels

---
 core/src/bmi.jl |  5 +----
 core/src/io.jl  | 25 +++++++++++++++----------
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/core/src/bmi.jl b/core/src/bmi.jl
index 2a0788b59..15d3a9fa3 100644
--- a/core/src/bmi.jl
+++ b/core/src/bmi.jl
@@ -513,10 +513,7 @@ end
 """Interpolate the levels and save them to SavedValues"""
 function save_exported_levels(u, t, integrator)
     update_exporter_levels!(integrator)
-    # TODO: multiple systems. Although at the point, shouldn't we
-    # just write to disk instead of using SavedValues?
-    exporter = first(values(integrator.p.level_exporters))
-    copy(exporter.level)
+    return vcat([exporter.level for exporter in values(integrator.p.level_exporters)]...)
 end
 
 "Load updates from 'Basin / time' into the parameters"
diff --git a/core/src/io.jl b/core/src/io.jl
index 4f5be008b..00e7ce4cb 100644
--- a/core/src/io.jl
+++ b/core/src/io.jl
@@ -231,17 +231,22 @@ function exported_levels_table(model::Model)::NamedTuple
     (; config, saved, integrator) = model
     (; t, saveval) = saved.exported_levels
 
-    name, exporter = first(integrator.p.level_exporters)
-    nelem = length(exporter.basin_index)
-    unique_elem_id = collect(1:nelem)
-    ntsteps = length(t)
-
-    time = repeat(datetime_since.(t, config.starttime); inner = nelem)
-    elem_id = repeat(unique_elem_id; outer = ntsteps)
-    levels = FlatVector(saveval)
-    names = fill(name, length(time))
+    # The level exporter may contain multiple named systems, but the
+    # saved levels are flat.
+    time = DateTime[]
+    name = String[]
+    element_id = Int[]
+    for (unique_name, exporter) in integrator.p.level_exporters
+        nelem = length(exporter.basin_index)
+        unique_elem_id = collect(1:nelem)
+        ntsteps = length(t)
+        append!(time, repeat(datetime_since.(t, config.starttime); inner = nelem))
+        append!(element_id, repeat(unique_elem_id; outer = ntsteps))
+        append!(name, fill(unique_name, length(time)))
+    end
 
-    return (; time, names, elem_id, levels)
+    level = FlatVector(saveval)
+    return (; time, name, element_id, level)
 end
 
 "Write a result table to disk as an Arrow file"