From 9f890422eeb10d867e97864d0424247f8e4dd053 Mon Sep 17 00:00:00 2001
From: Abel Soares Siqueira <abel.siqueira@esciencecenter.nl>
Date: Wed, 4 Dec 2024 13:38:17 +0100
Subject: [PATCH] Remove graph dependency from src/constraints/consumer.jl

Part of #942
Closes #943
---
 src/constraints/consumer.jl | 80 +++++++++++++++++++++++++++----------
 src/create-model.jl         |  4 +-
 test/test-case-studies.jl   | 10 ++++-
 3 files changed, 69 insertions(+), 25 deletions(-)

diff --git a/src/constraints/consumer.jl b/src/constraints/consumer.jl
index 68376b45..9b14a9e8 100644
--- a/src/constraints/consumer.jl
+++ b/src/constraints/consumer.jl
@@ -11,29 +11,65 @@ add_consumer_constraints!(model,
 
 Adds the consumer asset constraints to the model.
 """
-
-function add_consumer_constraints!(model, constraints, graph)
+function add_consumer_constraints!(connection, model, constraints)
     cons = constraints[:balance_consumer]
-    incoming_flow_highest_in_out_resolution = cons.expressions[:incoming]
-    outgoing_flow_highest_in_out_resolution = cons.expressions[:outgoing]
+
+    # TODO: Store the name of the table in the TulipaConstraint
+    table = _create_consumer_table(connection, "cons_balance_consumer")
 
     # - Balance constraint (using the lowest temporal resolution)
-    model[:consumer_balance] = [
-        @constraint(
-            model,
-            incoming_flow_highest_in_out_resolution[row.index] -
-            outgoing_flow_highest_in_out_resolution[row.index] -
-            profile_aggregation(
-                Statistics.mean,
-                graph[row.asset].rep_periods_profiles,
-                row.year,
-                row.year,
-                ("demand", row.rep_period),
-                row.time_block_start:row.time_block_end,
-                1.0,
-            ) * graph[row.asset].peak_demand[row.year] in
-            graph[row.asset].consumer_balance_sense,
-            base_name = "consumer_balance[$(row.asset),$(row.year),$(row.rep_period),$(row.time_block_start):$(row.time_block_end)]"
-        ) for row in eachrow(cons.indices)
-    ]
+    attach_constraint!(
+        model,
+        cons,
+        :consumer_balance,
+        [
+            begin
+                consumer_balance_sense = if ismissing(row.consumer_balance_sense)
+                    MathOptInterface.EqualTo(0.0)
+                else
+                    MathOptInterface.GreaterThan(0.0)
+                end
+                @constraint(
+                    model,
+                    incoming_flow - outgoing_flow - row.demand_agg * row.peak_demand in
+                    consumer_balance_sense,
+                    base_name = "consumer_balance[$(row.asset),$(row.year),$(row.rep_period),$(row.time_block_start):$(row.time_block_end)]"
+                )
+            end for (row, incoming_flow, outgoing_flow) in
+            zip(table, cons.expressions[:incoming], cons.expressions[:outgoing])
+        ],
+    )
+
+    return
+end
+
+function _create_consumer_table(connection, cons::String)
+    return DuckDB.query(
+        connection,
+        "SELECT
+            cons.*,
+            ANY_VALUE(asset.type) AS type,
+            ANY_VALUE(asset.consumer_balance_sense) AS consumer_balance_sense,
+            ANY_VALUE(asset_milestone.peak_demand) AS peak_demand,
+            COALESCE(MEAN(profile.value), 1) AS demand_agg
+        FROM $cons AS cons
+        LEFT JOIN asset
+            ON cons.asset = asset.asset
+        LEFT JOIN asset_milestone
+            ON cons.asset = asset_milestone.asset
+            AND cons.year = asset_milestone.milestone_year
+        LEFT JOIN assets_profiles
+            ON cons.asset = assets_profiles.asset
+            AND cons.year = assets_profiles.commission_year
+            AND assets_profiles.profile_type = 'demand'
+        LEFT JOIN profiles_rep_periods AS profile
+            ON profile.profile_name = assets_profiles.profile_name
+            AND cons.year = profile.year
+            AND cons.rep_period = profile.rep_period
+            AND cons.time_block_start <= profile.timestep
+            AND profile.timestep <= cons.time_block_end
+        GROUP BY cons.*
+        ORDER BY cons.index -- order is important
+        ",
+    )
 end
diff --git a/src/create-model.jl b/src/create-model.jl
index eb19b574..1c80975a 100644
--- a/src/create-model.jl
+++ b/src/create-model.jl
@@ -36,7 +36,7 @@ function create_model!(energy_problem; kwargs...)
 end
 
 """
-    model = create_model(graph, representative_periods, dataframes, timeframe, groups; write_lp_file = false, enable_names = true)
+    model = create_model(connection, graph, representative_periods, dataframes, timeframe, groups; write_lp_file = false, enable_names = true)
 
 Create the energy model given the `graph`, `representative_periods`, dictionary of `dataframes` (created by [`construct_dataframes`](@ref)), timeframe, and groups.
 """
@@ -122,7 +122,7 @@ function create_model(
 
     @timeit to "add_energy_constraints!" add_energy_constraints!(model, constraints, graph)
 
-    @timeit to "add_consumer_constraints!" add_consumer_constraints!(model, constraints, graph)
+    @timeit to "add_consumer_constraints!" add_consumer_constraints!(connection, model, constraints)
 
     @timeit to "add_storage_constraints!" add_storage_constraints!(
         model,
diff --git a/test/test-case-studies.jl b/test/test-case-studies.jl
index fc8cdb1b..b00b2d01 100644
--- a/test/test-case-studies.jl
+++ b/test/test-case-studies.jl
@@ -85,8 +85,16 @@ end
     dir = joinpath(INPUT_FOLDER, "Tiny")
     connection = DBInterface.connect(DuckDB.DB)
     _read_csv_folder(connection, dir)
+    DuckDB.execute( # Make it infeasible
+        connection,
+        "UPDATE asset_milestone
+            SET peak_demand = -1
+            WHERE
+                asset = 'demand'
+                AND milestone_year = 2030
+        ",
+    )
     energy_problem = EnergyProblem(connection)
-    energy_problem.graph["demand"].peak_demand[2030] = -1 # make it infeasible
     create_model!(energy_problem)
     @test_logs (:warn, "Model status different from optimal") solve_model!(energy_problem)
     @test energy_problem.termination_status == JuMP.INFEASIBLE