From 83678252b1af3a208de6fab1f55512524cf49244 Mon Sep 17 00:00:00 2001
From: Lukas Jiriste <ljiriste@student.42prague.com>
Date: Fri, 13 Sep 2024 01:56:03 +0200
Subject: [PATCH 1/1] Create a helper function for basin recognition

This function accepts an array representing the height map of an area
and accumulates pixel count through steepest path downward.
This is a naive implementation that takes quite some time to complete
for larger maps.
---
 povodi.py | 44 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)
 create mode 100644 povodi.py

diff --git a/povodi.py b/povodi.py
new file mode 100644
index 0000000..86477ef
--- /dev/null
+++ b/povodi.py
@@ -0,0 +1,44 @@
+import numpy as np
+
+def	is_noncalced_max(height_map, calced, multi_ind):
+	low_x = max(0, multi_ind[0] - 1)
+	high_x = min(height_map.shape[0] - 1, multi_ind[0] + 1)
+	low_y = max(0, multi_ind[1] - 1)
+	high_y = min(height_map.shape[1] - 1, multi_ind[1] + 1)
+	masked_subarr = np.ma.masked_array(height_map[low_x:high_x + 1, low_y:high_y + 1],
+			mask=calced[low_x:high_x + 1, low_y:high_y + 1])
+	return (masked_subarr.max() == height_map[multi_ind])
+
+def	get_recipient(height_map, multi_ind):
+	sub_x = 1
+	sub_y = 1
+	high_x = min(height_map.shape[0] - 1, multi_ind[0] + 1)
+	low_x = multi_ind[0] - 1
+	if (low_x < 0):
+		sub_x = 0
+		low_x = 0
+	high_y = min(height_map.shape[1] - 1, multi_ind[1] + 1)
+	low_y = multi_ind[1] - 1
+	if (low_y < 0):
+		sub_y = 0
+		low_y = 0
+	sub_arr = height_map[low_x:high_x + 1, low_y:high_y + 1]
+	min_ind_small = np.unravel_index(sub_arr.argmin(), sub_arr.shape)
+	return ((min_ind_small[0] + multi_ind[0] - sub_x, min_ind_small[1] + multi_ind[1] - sub_y))
+
+def povodi(height_map):
+	height_map = np.array(height_map)
+	calced = height_map.astype('bool')
+	calced.fill(0)
+	result = height_map.astype("uint32")
+	result.fill(1)
+	while (not(calced.all())):
+		print(calced.sum() / calced.size)
+		it = np.nditer(height_map, flags=['multi_index'])
+		for el in it:
+			if (not(calced[it.multi_index]) and is_noncalced_max(height_map, calced, it.multi_index)):
+				calced[it.multi_index] = 1
+				recipient_ind = get_recipient(height_map, it.multi_index)
+				if (recipient_ind != it.multi_index):
+					result[recipient_ind] += result[it.multi_index]
+	return (result)
-- 
2.30.2