[Feature] Scenario Creation (#78)

Co-authored-by: Tim Lorsbach <tim@lorsba.ch>
Reviewed-on: enviPath/enviPy#78

epdb/models.py

@@ -20,6 +20,7 @@ from django.db import models, transaction
 from django.db.models import JSONField, Count, Q, QuerySet
 from django.utils import timezone
 from django.utils.functional import cached_property
+from envipy_additional_information import EnviPyModel
 from model_utils.models import TimeStampedModel
 from polymorphic.models import PolymorphicModel
 from sklearn.metrics import precision_score, recall_score, jaccard_score
@@ -2292,7 +2293,8 @@ class ApplicabilityDomain(EnviPathModel):
                 transformation = {
                     'rule': rule_data,
                     'reliability': rule_reliabilities[rule_idx],
-                    # TODO
+                    # We're setting it here to False, as we don't know whether "assess" is called during pathway
+                    # prediction or from Model Page. For persisted Nodes this field will be overwritten at runtime
                     'is_predicted': False,
                     'local_compatibility': local_compatibilities[rule_idx],
                     'probability': preds[rule_idx].probability,
@@ -2407,27 +2409,88 @@ class Scenario(EnviPathModel):
 
     @staticmethod
     @transaction.atomic
-    def create(package, name, description, date, type, additional_information):
+    def create(package: 'Package', name:str, description:str, scenario_date:str, scenario_type:str, additional_information: List['EnviPyModel']):
         s = Scenario()
         s.package = package
+
+        if name is None or name.strip() == '':
+            name = f"Scenario {Scenario.objects.filter(package=package).count() + 1}"
+
         s.name = name
-        s.description = description
-        s.date = date
-        s.type = type
-        s.additional_information = additional_information
+
+        if description is not None and description.strip() != '':
+            s.description = description
+
+        if scenario_date is not None and scenario_date.strip() != '':
+            s.scenario_date = scenario_date
+
+        if scenario_type is not None and scenario_type.strip() != '':
+            s.scenario_type = scenario_type
+
+        add_inf = defaultdict(list)
+
+        for info in additional_information:
+            cls_name = info.__class__.__name__
+            ai_data = json.loads(info.model_dump_json())
+            ai_data['uuid'] = f"{uuid4()}"
+            add_inf[cls_name].append(ai_data)
+
+
+        s.additional_information = add_inf
 
         s.save()
 
         return s
 
-    def add_additional_information(self, data):
-        pass
+    @transaction.atomic
+    def add_additional_information(self, data: 'EnviPyModel'):
+        cls_name = data.__class__.__name__
+        ai_data = json.loads(data.model_dump_json())
+        ai_data['uuid'] = f"{uuid4()}"
 
-    def remove_additional_information(self, data):
-        pass
+        if cls_name not in self.additional_information:
+            self.additional_information[cls_name] = []
 
-    def set_additional_information(self, data):
-        pass
+        self.additional_information[cls_name].append(ai_data)
+        self.save()
+
+
+    @transaction.atomic
+    def remove_additional_information(self, ai_uuid):
+        found_type = None
+        found_idx = -1
+
+        for k, vals in self.additional_information.items():
+            for i, v in enumerate(vals):
+                if v['uuid'] == ai_uuid:
+                    found_type = k
+                    found_idx = i
+                    break
+
+        if found_type is not None and found_idx >= 0:
+            if len(self.additional_information[found_type]) == 1:
+                del self.additional_information[k]
+            else:
+                self.additional_information[k].pop(found_idx)
+            self.save()
+        else:
+            raise ValueError(f"Could not find additional information with uuid {ai_uuid}")
+
+    @transaction.atomic
+    def set_additional_information(self, data: Dict[str, 'EnviPyModel']):
+        new_ais = defaultdict(list)
+        for k, vals in data.items():
+            for v in vals:
+                ai_data = json.loads(v.model_dump_json())
+                if hasattr(v, 'uuid'):
+                    ai_data['uuid'] = str(v.uuid)
+                else:
+                    ai_data['uuid'] = str(uuid4())
+
+                new_ais[k].append(ai_data)
+
+        self.additional_information = new_ais
+        self.save()
 
     def get_additional_information(self):
         from envipy_additional_information import NAME_MAPPING
@@ -2437,7 +2500,14 @@ class Scenario(EnviPathModel):
                 continue
 
             for v in vals:
-                yield NAME_MAPPING[k](**json.loads(v))
+                # Per default additional fields are ignored
+                MAPPING = {c.__name__: c for c in NAME_MAPPING.values()}
+                inst = MAPPING[k](**v)
+                # Add uuid to uniquely identify objects for manipulation
+                if 'uuid' in v:
+                    inst.__dict__['uuid'] = v['uuid']
+
+                yield inst
 
 
 class UserSettingPermission(Permission):