From 06422de3265970858cc259fc1a0f31a3e29f4801 Mon Sep 17 00:00:00 2001
From: Adam Green <adam.green@adgefficiency.com>
Date: Sun, 19 May 2024 13:08:37 +1200
Subject: [PATCH] refactor: constraints

---
 energypylinear/constraints.py | 73 ++++++++++++++++-------------------
 1 file changed, 33 insertions(+), 40 deletions(-)

diff --git a/energypylinear/constraints.py b/energypylinear/constraints.py
index f76bc55a..57f22464 100644
--- a/energypylinear/constraints.py
+++ b/energypylinear/constraints.py
@@ -15,9 +15,8 @@ class ConstraintTerm:
 
     The sum of terms creates either side of a constraint - left-hand side (LHS) or right-hand side (RHS).
 
-    Examples:
-
     ```
+    Constraint = LHS sense RHS
     Constraint = ConstraintTerm + ConstraintTerm <= 2.0
     ```
 
@@ -218,46 +217,29 @@ def add_custom_constraint(
     if isinstance(constraint, dict):
         constraint = epl.Constraint(**constraint)
 
-    if constraint.interval_aggregation == "sum":
-        lhs: list = []
-        rhs: list = []
-        for i in interval_data.idx:
-            _add_terms(
-                constraint.lhs,
-                lhs,
-                ivars,
-                interval_data,
-                i,
-                divide_constant_by_idx_len=True,
-            )
-            _add_terms(
-                constraint.rhs,
-                rhs,
-                ivars,
-                interval_data,
-                i,
-                divide_constant_by_idx_len=True,
-            )
-            assert len(lhs) > 0
-            assert len(rhs) > 0
-
-        optimizer.constrain(
-            pulp.LpConstraint(
-                e=pulp.lpSum(lhs) - pulp.lpSum(rhs),
-                sense=sense_mapper[constraint.sense],
-                rhs=0.0,
-            )
+    lhs: list = []
+    rhs: list = []
+    for i in interval_data.idx:
+        _add_terms(
+            constraint.lhs,
+            lhs,
+            ivars,
+            interval_data,
+            i,
+            divide_constant_by_idx_len=constraint.interval_aggregation == "sum",
         )
+        _add_terms(
+            constraint.rhs,
+            rhs,
+            ivars,
+            interval_data,
+            i,
+            divide_constant_by_idx_len=constraint.interval_aggregation == "sum",
+        )
+        assert len(lhs) > 0
+        assert len(rhs) > 0
 
-    else:
-        for i in interval_data.idx:
-            lhs = []
-            rhs = []
-            _add_terms(constraint.lhs, lhs, ivars, interval_data, i)
-            _add_terms(constraint.rhs, rhs, ivars, interval_data, i)
-            assert len(lhs) > 0
-            assert len(rhs) > 0
-
+        if constraint.interval_aggregation != "sum":
             optimizer.constrain(
                 pulp.LpConstraint(
                     e=pulp.lpSum(lhs) - pulp.lpSum(rhs),
@@ -265,3 +247,14 @@ def add_custom_constraint(
                     rhs=0.0,
                 )
             )
+            lhs = []
+            rhs = []
+
+    if constraint.interval_aggregation == "sum":
+        optimizer.constrain(
+            pulp.LpConstraint(
+                e=pulp.lpSum(lhs) - pulp.lpSum(rhs),
+                sense=sense_mapper[constraint.sense],
+                rhs=0.0,
+            )
+        )