From f5ff21764080c4d4c11d0869b7302c837f585e27 Mon Sep 17 00:00:00 2001
From: trevorb1 <trevor_barnes@sfu.ca>
Date: Thu, 7 Nov 2024 16:18:09 -0800
Subject: [PATCH] update total discounted cost calculation

---
 src/otoole/results/result_package.py  | 24 +++++++++++++++++++++++-
 tests/results/test_results_package.py | 26 +++++++++++++-------------
 2 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/src/otoole/results/result_package.py b/src/otoole/results/result_package.py
index 33fdb93..6917781 100644
--- a/src/otoole/results/result_package.py
+++ b/src/otoole/results/result_package.py
@@ -943,14 +943,36 @@ def total_discounted_cost(self) -> pd.DataFrame:
                     - SalvageValueStorage[r,s,y] / ((1+DiscountRateStorage[r,s])^(max{yy in YEAR} max(yy)-min{yy in YEAR} min(yy)+1))
                 )
             ) ~VALUE;
+
+        Alternatively, can be written as::
+
+            r~REGION, y~YEAR,
+            TotalDiscountedCost[r,y] :=
+            sum{t in TECHNOLOGY} TotalDiscountedCostByTechnology[r,t,y] + sum{s in STORAGE} TotalDiscountedStorageCost[r,s,y]
         """
         try:
             discounted_cost_by_technology = self["DiscountedCostByTechnology"]
+            discounted_cost_by_storage = self["DiscountedCostByStorage"]
 
         except KeyError as ex:
             raise KeyError(self._msg("TotalDiscountedCost", str(ex)))
 
-        data = discounted_cost_by_technology
+        discounted_tech = (
+            discounted_cost_by_technology.droplevel("TECHNOLOGY")
+            .reset_index()
+            .groupby(["REGION", "YEAR"])
+            .sum()
+        )
+        discounted_storage = (
+            discounted_cost_by_storage.droplevel("STORAGE")
+            .reset_index()
+            .groupby(["REGION", "YEAR"])
+            .sum()
+        )
+
+        total_discounted_cost = discounted_tech.add(discounted_storage, fill_value=0)
+
+        data = total_discounted_cost
 
         if not data.empty:
             data = data.groupby(by=["REGION", "YEAR"]).sum()
diff --git a/tests/results/test_results_package.py b/tests/results/test_results_package.py
index 052d614..9dc85b3 100644
--- a/tests/results/test_results_package.py
+++ b/tests/results/test_results_package.py
@@ -444,15 +444,15 @@ def discounted_salvage_value_storage():
 def discounted_storage_cost():
     data = pd.DataFrame(
         data=[
-            ["SIMPLICITY", "DUMMY", 2014, 111],
-            ["SIMPLICITY", "DUMMY", 2015, 222],
-            ["SIMPLICITY", "DUMMY", 2016, 333],
-            ["SIMPLICITY", "GAS_EXTRACTION", 2014, 444],
-            ["SIMPLICITY", "GAS_EXTRACTION", 2015, 555],
-            ["SIMPLICITY", "GAS_EXTRACTION", 2016, 666],
+            ["SIMPLICITY", "DAM", 2014, 11.1],
+            ["SIMPLICITY", "DAM", 2015, 22.2],
+            ["SIMPLICITY", "DAM", 2016, 33.3],
+            ["SIMPLICITY", "BATTERY", 2014, 44.4],
+            ["SIMPLICITY", "BATTERY", 2015, 55.5],
+            ["SIMPLICITY", "BATTERY", 2016, 66.6],
         ],
-        columns=["REGION", "TECHNOLOGY", "YEAR", "VALUE"],
-    ).set_index(["REGION", "TECHNOLOGY", "YEAR"])
+        columns=["REGION", "STORAGE", "YEAR", "VALUE"],
+    ).set_index(["REGION", "STORAGE", "YEAR"])
     return data
 
 
@@ -1217,21 +1217,21 @@ def test_null(self, null: ResultsPackage):
 
 class TestTotalDiscountedCost:
     def test_calculate_total_discounted_cost(
-        self,
-        discounted_technology_cost,
+        self, discounted_technology_cost, discounted_storage_cost
     ):
 
         results = {
             "DiscountedCostByTechnology": discounted_technology_cost,
+            "DiscountedCostByStorage": discounted_storage_cost,
         }
 
         package = ResultsPackage(results)
         actual = package.total_discounted_cost()
         expected = pd.DataFrame(
             data=[
-                ["SIMPLICITY", 2014, 555],
-                ["SIMPLICITY", 2015, 777],
-                ["SIMPLICITY", 2016, 999],
+                ["SIMPLICITY", 2014, 610.5],
+                ["SIMPLICITY", 2015, 854.7],
+                ["SIMPLICITY", 2016, 1098.9],
             ],
             columns=["REGION", "YEAR", "VALUE"],
         ).set_index(["REGION", "YEAR"])