From 56c84ad39988c8cb2a6c230bcbaee0b6bc3046fa Mon Sep 17 00:00:00 2001
From: p_monsalvete <pilar.monsalvete@concordia.ca>
Date: Tue, 29 Aug 2023 13:29:51 -0400
Subject: [PATCH] solved bugs in adapting montreal_custom_catalog.py to new
 catalog data

---
 .../data_models/energy_systems/content.py     |  7 ++++++
 .../data_models/energy_systems/system.py      |  5 ++--
 .../energy_systems/thermal_storage_system.py  |  8 ++++---
 .../energy_systems/montreal_custom_catalog.py | 24 +++++++++----------
 4 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/hub/catalog_factories/data_models/energy_systems/content.py b/hub/catalog_factories/data_models/energy_systems/content.py
index f6a03ea0..d3fd9068 100644
--- a/hub/catalog_factories/data_models/energy_systems/content.py
+++ b/hub/catalog_factories/data_models/energy_systems/content.py
@@ -69,4 +69,11 @@ class Content:
 
     return content
 
+  def __str__(self):
+    """Print content"""
+    _archetypes = []
+    for _archetype in self.archetypes:
+      _archetypes.append(_archetype.to_dictionary())
+    content = {'Archetypes': _archetypes}
 
+    return str(content)
diff --git a/hub/catalog_factories/data_models/energy_systems/system.py b/hub/catalog_factories/data_models/energy_systems/system.py
index 7acb930b..30ee9292 100644
--- a/hub/catalog_factories/data_models/energy_systems/system.py
+++ b/hub/catalog_factories/data_models/energy_systems/system.py
@@ -9,7 +9,8 @@ Code contributors: Saeed Ranjbar saeed.ranjbar@concordia.ca
 from typing import Union, List
 
 from hub.catalog_factories.data_models.energy_systems.generation_system import GenerationSystem
-from hub.catalog_factories.data_models.energy_systems.energy_storage_system import EnergyStorageSystem
+from hub.catalog_factories.data_models.energy_systems.thermal_storage_system import ThermalStorageSystem
+from hub.catalog_factories.data_models.energy_systems.electrical_storage_system import ElectricalStorageSystem
 from hub.catalog_factories.data_models.energy_systems.distribution_system import DistributionSystem
 from hub.catalog_factories.data_models.energy_systems.emission_system import EmissionSystem
 
@@ -95,7 +96,7 @@ class System:
     return self._emission_systems
 
   @property
-  def energy_storage_systems(self) -> Union[None, List[EnergyStorageSystem]]:
+  def energy_storage_systems(self) -> Union[None, List[ThermalStorageSystem], List[ElectricalStorageSystem]]:
     """
     Get energy storage systems
     :return: [EnergyStorageSystem]
diff --git a/hub/catalog_factories/data_models/energy_systems/thermal_storage_system.py b/hub/catalog_factories/data_models/energy_systems/thermal_storage_system.py
index 679669a4..71b4f931 100644
--- a/hub/catalog_factories/data_models/energy_systems/thermal_storage_system.py
+++ b/hub/catalog_factories/data_models/energy_systems/thermal_storage_system.py
@@ -67,9 +67,11 @@ class ThermalStorageSystem(EnergyStorageSystem):
 
   def to_dictionary(self):
     """Class content to dictionary"""
-    _layers = []
-    for _layer in self.layers:
-      _layers.append(_layer.to_dictionary())
+    _layers = None
+    if self.layers is not None:
+      _layers = []
+      for _layer in self.layers:
+        _layers.append(_layer.to_dictionary())
     content = {'Storage component': {
       'storage id': self.id,
       'name': self.name,
diff --git a/hub/catalog_factories/energy_systems/montreal_custom_catalog.py b/hub/catalog_factories/energy_systems/montreal_custom_catalog.py
index a9696c21..b2df20c2 100644
--- a/hub/catalog_factories/energy_systems/montreal_custom_catalog.py
+++ b/hub/catalog_factories/energy_systems/montreal_custom_catalog.py
@@ -162,33 +162,33 @@ class MontrealCustomCatalog(Catalog):
       name = system['name']
       demands = system['demands']['demand']
       generation_equipment = system['equipments']['generation_id']
-      _generation_equipment = None
+      _generation_equipments = None
       for equipment_archetype in self._catalog_generation_equipments:
         if int(equipment_archetype.id) == int(generation_equipment):
-          _generation_equipment = equipment_archetype
-      _storage_equipment = None
+          _generation_equipments = [equipment_archetype]
+      _storage_equipments = None
       for equipment_archetype in self._catalog_storage_equipments:
         if int(equipment_archetype.id) == int(generation_equipment):
-          _storage_equipment = equipment_archetype
+          _storage_equipments = [equipment_archetype]
       distribution_equipment = system['equipments']['distribution_id']
-      _distribution_equipment = None
+      _distribution_equipments = None
       for equipment_archetype in self._catalog_distribution_equipments:
         if int(equipment_archetype.id) == int(distribution_equipment):
-          _distribution_equipment = equipment_archetype
+          _distribution_equipments = [equipment_archetype]
       emission_equipment = system['equipments']['dissipation_id']
-      _emission_equipment = None
+      _emission_equipments = None
       for equipment_archetype in self._catalog_emission_equipments:
         if int(equipment_archetype.id) == int(emission_equipment):
-          _emission_equipment = equipment_archetype
+          _emission_equipments = [equipment_archetype]
 
       _catalog_systems.append(System(self._lod,
                                      system_id,
                                      name,
                                      demands,
-                                     [_generation_equipment],
-                                     [_distribution_equipment],
-                                     [_emission_equipment],
-                                     [_storage_equipment]))
+                                     _generation_equipments,
+                                     _distribution_equipments,
+                                     _emission_equipments,
+                                     _storage_equipments))
     return _catalog_systems
 
   def _load_archetypes(self):