From ccf67fc1d3cb0cf5d79cd64a8a6eb30cb88e480b Mon Sep 17 00:00:00 2001
From: pedroripper <pedros.ripper@gmail.com>
Date: Tue, 3 Dec 2024 10:54:57 -0300
Subject: [PATCH] Fix OpenBinary reader when using header

---
 src/OpenBinary/reader.jl      |   6 +-
 test/OpenBinary/use_header.jl | 144 ++++++++++++++++++++++++++++++++++
 test/runtests.jl              |   3 +
 3 files changed, 152 insertions(+), 1 deletion(-)
 create mode 100644 test/OpenBinary/use_header.jl

diff --git a/src/OpenBinary/reader.jl b/src/OpenBinary/reader.jl
index 920b43f2..674e0fe6 100644
--- a/src/OpenBinary/reader.jl
+++ b/src/OpenBinary/reader.jl
@@ -301,6 +301,9 @@ function PSRI.open(
             end
         end
 
+        queried_agents = length(header)
+        agent_names = header
+
         if !allow_empty && isempty(indices)
             if isempty(header)
                 error(
@@ -314,12 +317,13 @@ function PSRI.open(
             end
         end
     else
+        queried_agents = total_agents
         indices = collect(1:total_agents)
     end
 
     data_buffer = zeros(Float32, total_agents)
 
-    data = zeros(Float64, total_agents)
+    data = zeros(Float64, queried_agents)
 
     if initial_stage != Dates.Date(1900, 1, 1)
         _year = initial_year
diff --git a/test/OpenBinary/use_header.jl b/test/OpenBinary/use_header.jl
new file mode 100644
index 00000000..931dada6
--- /dev/null
+++ b/test/OpenBinary/use_header.jl
@@ -0,0 +1,144 @@
+function use_header()
+    file_path = joinpath(".", "data")
+
+    STAGES = 2
+    SCENARIOS = 2
+    AGENTS = ["X", "Y", "Z"]
+    UNIT = "MW"
+
+    iow = PSRI.open(
+        PSRI.OpenBinary.Writer,
+        file_path;
+        is_hourly = true,
+        scenarios = SCENARIOS,
+        stages = STAGES,
+        agents = AGENTS,
+        unit = UNIT,
+        # optional:
+        initial_stage = 2,
+        initial_year = 2006,
+        stage_type = PSRI.STAGE_MONTH,
+    )
+
+    for t in 1:STAGES, s in 1:SCENARIOS
+        for b in 1:PSRI.blocks_in_stage(iow, t)
+            X = 10_000.0 * t + 1000.0 * s + b
+            Y = b + 0.0
+            Z = 10.0 * t + s
+            PSRI.write_registry(
+                iow,
+                [X, Y, Z],
+                t,
+                s,
+                b,
+            )
+        end
+    end
+
+    PSRI.close(iow)
+
+    # All agents
+
+    ior = PSRI.open(
+        PSRI.OpenBinary.Reader,
+        file_path;
+        use_header = true,
+        header = ["X", "Y", "Z"],
+    )
+
+    @test PSRI.agent_names(ior) == ["X", "Y", "Z"]
+
+    for t in 1:1, s in 1:1
+        for b in 1:PSRI.blocks_in_stage(ior, t)
+            X = 10_000.0 * t + 1000.0 * s + b
+            Y = b + 0.0
+            Z = 10.0 * t + s
+            ref = [X, Y, Z]
+            for agent in 1:3
+                @test ior[agent] == ref[agent]
+            end
+            PSRI.next_registry(ior)
+        end
+    end
+
+    PSRI.close(ior)
+
+    # # Only X
+
+    ior = PSRI.open(
+        PSRI.OpenBinary.Reader,
+        file_path;
+        use_header = true,
+        header = ["X"],
+    )
+
+    @test PSRI.agent_names(ior) == ["X"]
+
+    for t in 1:1, s in 1:1
+        for b in 1:PSRI.blocks_in_stage(ior, t)
+            X = 10_000.0 * t + 1000.0 * s + b
+            ref = [X]
+            @test ior[1] == ref[1]
+            PSRI.next_registry(ior)
+        end
+    end
+
+    PSRI.close(ior)
+    ior = nothing
+
+    # Only Y
+
+    ior = PSRI.open(
+        PSRI.OpenBinary.Reader,
+        file_path;
+        use_header = true,
+        header = ["Y"],
+    )
+
+    @test PSRI.agent_names(ior) == ["Y"]
+
+    for t in 1:1, s in 1:1
+        for b in 1:PSRI.blocks_in_stage(ior, t)
+            Y = b + 0.0
+            ref = [Y]
+            @test ior[1] == ref[1]
+            PSRI.next_registry(ior)
+        end
+    end
+
+    PSRI.close(ior)
+    ior = nothing
+
+    # All agents reverse
+
+    ior = PSRI.open(
+        PSRI.OpenBinary.Reader,
+        file_path;
+        use_header = true,
+        header = ["Z", "Y", "X"],
+    )
+
+    @test PSRI.agent_names(ior) == ["Z", "Y", "X"]
+
+    for t in 1:1, s in 1:1
+        for b in 1:PSRI.blocks_in_stage(ior, t)
+            X = 10_000.0 * t + 1000.0 * s + b
+            Y = b + 0.0
+            Z = 10.0 * t + s
+            ref = [Z, Y, X]
+            for agent in 1:3
+                @test ior[agent] == ref[agent]
+            end
+            PSRI.next_registry(ior)
+        end
+    end
+
+    PSRI.close(ior)
+    ior = nothing
+
+    rm(file_path * ".bin")
+    rm(file_path * ".hdr")
+    return
+end
+
+use_header()
diff --git a/test/runtests.jl b/test/runtests.jl
index 61838460..59d4f5d4 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -35,6 +35,9 @@ end
         @testset "Write file partially" begin
             @time include("OpenBinary/incomplete_file.jl")
         end
+        @testset "Use header" begin
+            @time include("OpenBinary/use_header.jl")
+        end
     end
     @testset "ReaderMapper" begin
         @time include("reader_mapper.jl")