From c6c5179c34336222a9f6e0154ab237e0612113c3 Mon Sep 17 00:00:00 2001
From: romainsacchi <r_s@me.com>
Date: Sat, 19 Oct 2024 17:03:20 +0200
Subject: [PATCH] Finish implementing incremental database

---
 premise/export.py       | 17 +++++++++++++++--
 premise/incremental.py  | 16 ++++++++++++++--
 premise/new_database.py |  3 +++
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/premise/export.py b/premise/export.py
index 24bc473d..ae67fd4b 100644
--- a/premise/export.py
+++ b/premise/export.py
@@ -602,7 +602,12 @@ def generate_new_activities(args):
 
 
 def generate_scenario_difference_file(
-    db_name, origin_db, scenarios, version, scenario_list, biosphere_name
+    db_name,
+    origin_db,
+    scenarios,
+    version,
+    scenario_list,
+    biosphere_name,
 ) -> tuple[DataFrame, list[dict], set[Any]]:
     """
     Generate a scenario difference file for a given list of databases
@@ -840,6 +845,7 @@ def generate_superstructure_db(
     version,
     scenario_list,
     file_format="excel",
+    preserve_original_column: bool = False,
 ) -> List[dict]:
     """
     Build a superstructure database from a list of databases
@@ -849,6 +855,9 @@ def generate_superstructure_db(
     :param filepath: the filepath of the new database
     :param version: the version of the new database
     :param file_format: the format of the scenario difference file. Can be "excel", "csv" or "feather".
+    :param preserve_original_column: whether to keep the original column in the scenario difference file
+    :param scenario_list: a list of external scenarios
+
     :return: a superstructure database
     """
 
@@ -871,7 +880,11 @@ def generate_superstructure_db(
     df = df.rename(columns={"from unit": "unit"})
 
     # remove the column `original`
-    df = df.drop(columns=["original"])
+    if not preserve_original_column:
+        df = df.drop(columns=["original"])
+    else:
+        scenario_list = ["original"] + scenario_list
+
     if "unit" in df.columns:
         df = df.drop(columns=["unit"])
 
diff --git a/premise/incremental.py b/premise/incremental.py
index ec9bfbfd..3b5a9204 100644
--- a/premise/incremental.py
+++ b/premise/incremental.py
@@ -1,3 +1,5 @@
+from datetime import datetime
+
 from .new_database import (
     NewDatabase,
     _update_biomass,
@@ -94,11 +96,10 @@ def update(self, sectors: [str, list, None] = None) -> None:
 
         applied_sectors = []
         for scenario in self.scenarios:
-            label = ""
             scenario["database"] = pickle.loads(pickle.dumps(self.database, -1))
             scenario_sectors = []
             for sector in sectors:
-                label += f"_{sector}"
+                label = f"(... + {sector})"
                 scenario_sectors.append(sector)
 
                 if sector == "external" and "external" not in scenario["pathway"]:
@@ -131,3 +132,14 @@ def update(self, sectors: [str, list, None] = None) -> None:
                 pbar_outer.update()
 
         print("Done!\n")
+
+    def write_increment_db_to_brightway(
+        self,
+        name: str = f"super_db_{datetime.now().strftime('%d-%m-%Y')}",
+        filepath: str = None,
+        file_format: str = "excel",
+    ) -> None:
+
+        self.write_superstructure_db_to_brightway(name, filepath, file_format, preserve_original_column=True)
+
+
diff --git a/premise/new_database.py b/premise/new_database.py
index 4b51d454..9bd83b5b 100644
--- a/premise/new_database.py
+++ b/premise/new_database.py
@@ -982,6 +982,7 @@ def write_superstructure_db_to_brightway(
         name: str = f"super_db_{datetime.now().strftime('%d-%m-%Y')}",
         filepath: str = None,
         file_format: str = "excel",
+        preserve_original_column: bool = False,
     ) -> None:
         """
         Register a super-structure database,
@@ -989,6 +990,7 @@ def write_superstructure_db_to_brightway(
         :param name: name of the super-structure database
         :param filepath: filepath of the "scenarios difference file"
         :param file_format: format of the "scenarios difference file" export. Can be "excel", "csv" or "feather".
+        :param preserve_original_column: if True, the original column names are preserved in the super-structure database.
         :return: filepath of the "scenarios difference file"
         """
 
@@ -1024,6 +1026,7 @@ def write_superstructure_db_to_brightway(
             version=self.version,
             file_format=file_format,
             scenario_list=list_scenarios,
+            preserve_original_column=preserve_original_column
         )
 
         tmp_scenario = self.scenarios[0]