From 1a67791cd943e1976254d9c640484147374bdbde Mon Sep 17 00:00:00 2001
From: Steph Prince <40640337+stephprince@users.noreply.github.com>
Date: Tue, 3 Sep 2024 10:54:00 -0700
Subject: [PATCH] add example to track container index

---
 src/nwb/NWBFile.cpp                     |  8 ++++++--
 src/nwb/NWBFile.hpp                     |  6 +++++-
 tests/examples/testWorkflowExamples.cpp | 11 +++++++----
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/src/nwb/NWBFile.cpp b/src/nwb/NWBFile.cpp
index 5d519944..007bee3c 100644
--- a/src/nwb/NWBFile.cpp
+++ b/src/nwb/NWBFile.cpp
@@ -23,6 +23,8 @@ using namespace AQNWB::NWB;
 
 constexpr SizeType CHUNK_XSIZE = 2048;
 
+std::vector<SizeType> NWBFile::emptyContainerIndexes = {};
+
 // NWBFile
 
 NWBFile::NWBFile(const std::string& idText, std::shared_ptr<BaseIO> io)
@@ -93,7 +95,8 @@ Status NWBFile::createFileStructure()
 Status NWBFile::createElectricalSeries(
     std::vector<Types::ChannelVector> recordingArrays,
     const BaseDataType& dataType,
-    RecordingContainers* recordingContainers)
+    RecordingContainers* recordingContainers,
+    std::vector<SizeType>& containerIndexes)
 {
   if (!io->canModifyObjects()) {
     return Status::Failure;
@@ -132,7 +135,8 @@ Status NWBFile::createElectricalSeries(
         SizeArray {0, channelVector.size()},
         SizeArray {CHUNK_XSIZE, 0});
     electricalSeries->initialize();
-    recordingContainers->addData(std::move(electricalSeries));
+    recordingContainers->addContainer(std::move(electricalSeries));
+    containerIndexes.push_back(recordingContainers->containers.size() - 1);
 
     // Add electrode information to electrode table (does not write to datasets
     // yet)
diff --git a/src/nwb/NWBFile.hpp b/src/nwb/NWBFile.hpp
index 61ecca04..277f45ab 100644
--- a/src/nwb/NWBFile.hpp
+++ b/src/nwb/NWBFile.hpp
@@ -69,13 +69,16 @@ class NWBFile
    *                        record from. A separate ElectricalSeries will be
    *                        created for each ChannelVector.
    * @param recordingContainers The container to store the created TimeSeries.
+   * @param containerIndexes The indexes of the containers added to
+   * recordingContainers
    * @param dataType The data type of the elements in the data block.
    * @return Status The status of the object creation operation.
    */
   Status createElectricalSeries(
       std::vector<Types::ChannelVector> recordingArrays,
       const BaseDataType& dataType = BaseDataType::I16,
-      RecordingContainers* recordingContainers = nullptr);
+      RecordingContainers* recordingContainers = nullptr,
+      std::vector<SizeType>& containerIndexes = emptyContainerIndexes);
 
 protected:
   /**
@@ -120,6 +123,7 @@ class NWBFile
 
   const std::string identifierText;
   std::shared_ptr<BaseIO> io;
+  static std::vector<SizeType> emptyContainerIndexes;
 };
 
 }  // namespace AQNWB::NWB
\ No newline at end of file
diff --git a/tests/examples/testWorkflowExamples.cpp b/tests/examples/testWorkflowExamples.cpp
index 759ec4da..b04fd73f 100644
--- a/tests/examples/testWorkflowExamples.cpp
+++ b/tests/examples/testWorkflowExamples.cpp
@@ -48,8 +48,11 @@ TEST_CASE("workflowExamples")
     // [example_workflow_nwbfile_snippet]
 
     // [example_workflow_datasets_snippet]
-    nwbfile->createElectricalSeries(
-        mockRecordingArrays, BaseDataType::I16, recordingContainers.get());
+    std::vector<SizeType> containerIndexes;
+    nwbfile->createElectricalSeries(mockRecordingArrays,
+                                    BaseDataType::I16,
+                                    recordingContainers.get(),
+                                    containerIndexes);
     // [example_workflow_datasets_snippet]
 
     // [example_workflow_start_snippet]
@@ -60,7 +63,7 @@ TEST_CASE("workflowExamples")
     bool isRecording = true;
     while (isRecording) {
       // write data to the file for each channel
-      for (SizeType i = 0; i < mockRecordingArrays.size(); ++i) {
+      for (SizeType i = 0; i < containerIndexes.size(); ++i) {
         const auto& channelVector = mockRecordingArrays[i];
         for (const auto& channel : channelVector) {
           // copy data into buffer
@@ -80,7 +83,7 @@ TEST_CASE("workflowExamples")
               dataBuffer.size(), channel.getBitVolts(), dataBuffer.data());
 
           // [example_workflow_write_snippet]
-          recordingContainers->writeTimeseriesData(i,
+          recordingContainers->writeTimeseriesData(containerIndexes[i],
                                                    channel,
                                                    dataShape,
                                                    positionOffset,