From 9ef415cb8cbafe7e07cd40360e1837be6639be62 Mon Sep 17 00:00:00 2001
From: spjuhel <pro@sjuhel.org>
Date: Fri, 12 Jul 2024 17:46:27 +0200
Subject: [PATCH] Possibly working EventArbitraryProd (remains to be tested)

---
 boario/event.py      | 12 ++++++++++--
 boario/simulation.py | 14 ++++++++++----
 tests/test_events.py |  6 +++---
 3 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/boario/event.py b/boario/event.py
index 18a5e3e..be314b4 100644
--- a/boario/event.py
+++ b/boario/event.py
@@ -187,7 +187,14 @@ def from_series(
             recovery_function=recovery_function,
         )
     elif event_type == "arbitrary":
-        raise NotImplementedError("This type of event is not implemented yet.")
+        return EventArbitraryProd._from_series(
+            impact=impact,
+            occurrence=occurrence,
+            duration=duration,
+            name=name,
+            recovery_tau=recovery_tau,
+            recovery_function=recovery_function,
+        )
 
     else:
         raise ValueError(f"Wrong event type: {event_type}")
@@ -530,7 +537,7 @@ def _from_scalar_regions_sectors(
         **kwargs,
     ) -> Event:
         affected_industries = cls._build_industries_idx(
-            affected_regions, affected_sectors
+            regions=affected_regions, sectors=affected_sectors
         )
         regional_distrib = cls._level_distrib(
             affected_industries.levels[0], impact_regional_distrib
@@ -615,6 +622,7 @@ def from_scalar_regions_sectors(
         if impact <= 0:
             raise ValueError("Impact is null.")
 
+    @classmethod
     def _build_industries_idx(cls, regions: RegionsList, sectors: SectorsList):
         # TODO: Move this in utils?
         if isinstance(regions, str):
diff --git a/boario/simulation.py b/boario/simulation.py
index 4381b8e..2214ac9 100644
--- a/boario/simulation.py
+++ b/boario/simulation.py
@@ -1361,6 +1361,12 @@ def __init__(
                 source_event.impact.copy(), self.sim.model.industries
             )
             self._indus_dmg = self._indus_dmg_0.copy()
+            if source_event.impact_households is not None:
+                self._house_dmg_0 = _thin_to_wide(
+                    source_event.impact_households.copy(), self.sim.model.regions
+                )
+                self._house_dmg = self._house_dmg_0.copy()
+
         if isinstance(source_event, EventKapitalRebuild):
             self._init_distrib("indus", indus_rebuild_distribution, source_event)
             assert self._indus_dmg_0 is not None
@@ -1379,10 +1385,6 @@ def __init__(
             self.rebuild_dem_indus_distribution = self._reb_dem_indus_distribution
 
             if source_event.impact_households is not None:
-                self._house_dmg_0 = _thin_to_wide(
-                    source_event.impact_households.copy(), self.sim.model.regions
-                )
-                self._house_dmg = self._house_dmg_0.copy()
                 self._init_distrib("house", house_rebuild_distribution, source_event)
                 self._rebuild_demand_house_0 = pd.DataFrame(
                     source_event.rebuilding_sectors.values[:, None]
@@ -1411,6 +1413,10 @@ def __init__(
                 )
 
         if isinstance(self.event, EventArbitraryProd):
+            self._prod_delta_from_arb_0 = _thin_to_wide(
+                source_event.impact.copy(), self.sim.model.industries
+            )
+            self._prod_delta_from_arb = self._prod_delta_from_arb_0.copy()
             self._recovery_function_arb_delta = partial(
                 self.event.recovery_function,
                 init_impact_stock=self._prod_delta_from_arb_0,
diff --git a/tests/test_events.py b/tests/test_events.py
index 3a6c5b6..dc9090c 100644
--- a/tests/test_events.py
+++ b/tests/test_events.py
@@ -140,7 +140,7 @@ def test_event_from_series_rebuilding(self, sample_series):
         )
 
     def test_event_from_series_arbitrary(self, sample_series):
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(ValueError):
             event.from_series(sample_series, event_type="arbitrary")
 
     def test_event_from_series_wrongtype(self, sample_series):
@@ -626,7 +626,7 @@ def r_fun2(init_impact_stock, recovery_time):
     )
     def test_event_from_series_arbitrary(self, impact, expected):
         if isinstance(expected, str) and expected == "error":
-            with pytest.raises(ValueError):
+            with pytest.raises((ValueError, AttributeError)):
                 event.from_series(
                     impact,
                     event_type="arbitrary",
@@ -639,5 +639,5 @@ def test_event_from_series_arbitrary(self, impact, expected):
                 recovery_tau=1,
             )
             pd.testing.assert_series_equal(
-                event_instance.prod_cap_delta_arbitrary, expected
+                event_instance.prod_cap_delta_arbitrary, expected, check_names=False
             )