From 7fd44f283448ac268551ad12fcb270242dd39e24 Mon Sep 17 00:00:00 2001
From: Kristofers Solo <kristians.cagulis@pm.me>
Date: Fri, 5 Jan 2024 19:06:48 +0200
Subject: [PATCH] feat(ai): add `get_wells`

---
 src/ai/fitness/fitness.py |  4 ----
 src/ai/fitness/wells.py   | 37 +++++++++++++++++++++++++++++++++++++
 tests/ai/test_fitness.py  | 10 ++++++++++
 3 files changed, 47 insertions(+), 4 deletions(-)
 create mode 100644 src/ai/fitness/wells.py

diff --git a/src/ai/fitness/fitness.py b/src/ai/fitness/fitness.py
index 2f769e9..7d51d16 100644
--- a/src/ai/fitness/fitness.py
+++ b/src/ai/fitness/fitness.py
@@ -13,10 +13,6 @@ def calculate_fitness(game: Game) -> float:
     return fitness
 
 
-def get_wells(field: np.ndarray) -> int:
-    pass
-
-
 def _calc_holes(field: np.ndarray) -> float:
     height, width = field.shape
     penalty = 0
diff --git a/src/ai/fitness/wells.py b/src/ai/fitness/wells.py
new file mode 100644
index 0000000..b1a82c1
--- /dev/null
+++ b/src/ai/fitness/wells.py
@@ -0,0 +1,37 @@
+from typing import Optional
+
+import numpy as np
+
+from ai.log import log
+
+from .peaks import get_peaks
+
+
+def get_wells(
+    *, peaks: Optional[np.ndarray] = None, field: Optional[np.ndarray] = None
+) -> np.ndarray:
+    if peaks is None and field is None:
+        raise ValueError("peaks and field cannot both be None")
+    elif peaks is None:
+        peaks = get_peaks(field)
+
+    wells = np.zeros_like(peaks)
+    log.debug(f"{peaks=}")
+
+    first_well = peaks[1] - peaks[0]
+    wells[0] = first_well if first_well > 0 else 0
+
+    last_well = peaks[-2] - peaks[-1]
+    wells[-1] = last_well if last_well > 0 else 0
+
+    for idx in range(1, len(peaks) - 1):
+        well_l = peaks[idx - 1] - peaks[idx]
+        well_l = well_l if well_l > 0 else 0
+
+        well_r = peaks[idx + 1] - peaks[idx]
+        well_r = well_r if well_r > 0 else 0
+
+        wells[idx] = well_l if well_l >= well_r else well_r
+
+    log.debug(f"{wells=}")
+    return wells
diff --git a/tests/ai/test_fitness.py b/tests/ai/test_fitness.py
index c274942..ec03d8f 100644
--- a/tests/ai/test_fitness.py
+++ b/tests/ai/test_fitness.py
@@ -8,6 +8,7 @@ from ai.fitness.transitions import (
     get_col_transition,
     get_row_transition,
 )
+from ai.fitness.wells import get_wells
 
 
 class TestFitness(unittest.TestCase):
@@ -58,3 +59,12 @@ class TestFitness(unittest.TestCase):
         )
         for field, answer in zip(self.fields, answers):
             self.assertTrue(np.array_equal(get_holes(field), answer))
+
+    def test_get_wells(self):
+        answers = (
+            np.array([1, 0, 2, 1, 0]),
+            np.array([0, 0, 0, 0, 0]),
+            np.array([2, 0, 2, 0, 0]),
+        )
+        for field, answer in zip(self.fields, answers):
+            self.assertTrue(np.array_equal(get_wells(field=field), answer))