welcome to the 3D world

coordinate.py

@@ -1,51 +1,77 @@
 from __future__ import annotations
 from dataclasses import dataclass
+from enum import Enum
 from math import sqrt, inf, atan2, degrees
-from typing import Union, List
+from typing import Union, List, Optional
+
+
+class DistanceAlgorithm(Enum):
+    MANHATTAN = 0
+    PYTHAGORAS = 1
 
 
 @dataclass(frozen=True, order=True)
 class Coordinate:
     x: int
     y: int
+    z: Optional[int] = None
 
-    def getDistanceTo(self, target: Coordinate, mode: int = 1, includeDiagonals: bool = False) -> Union[int, float]:
+    def getDistanceTo(self, target: Coordinate, mode: DistanceAlgorithm = DistanceAlgorithm.PYTHAGORAS,
+                      includeDiagonals: bool = False) -> Union[int, float]:
         """
         Get distance to target Coordinate
 
         :param target:
         :param mode: Calculation Mode (0 = Manhattan, 1 = Pythagoras)
-        :param includeDiagonals: in Manhattan Mode specify if diagonal movements are allowed (counts as 1.4)
+        :param includeDiagonals: in Manhattan Mode specify if diagonal
+                                 movements are allowed (counts as 1.4 in 2D, 1.7 in 3D)
         :return: Distance to Target
         """
         assert isinstance(target, Coordinate)
-        assert mode in [0, 1]
+        assert isinstance(mode, DistanceAlgorithm)
         assert isinstance(includeDiagonals, bool)
 
-        if mode == 1:
+        if mode == DistanceAlgorithm.PYTHAGORAS:
+            if self.z is None:
                 return sqrt(abs(self.x - target.x) ** 2 + abs(self.y - target.y) ** 2)
-        elif mode == 0:
+            else:
+                ab = sqrt(abs(self.x - target.x) ** 2 + abs(self.y - target.y) ** 2)
+                return sqrt(ab ** 2 + abs(self.z - target.z) ** 2)
+        elif mode == DistanceAlgorithm.MANHATTAN:
             if not includeDiagonals:
+                if self.z is None:
                     return abs(self.x - target.x) + abs(self.y - target.y)
                 else:
-                x_dist = abs(self.x - target.x)
+                    return abs(self.x - target.x) + abs(self.y - target.y) + abs(self.z - target.z)
-                y_dist = abs(self.y - target.y)
+            else:
-                o_dist = max(x_dist, y_dist) - min(x_dist, y_dist)
+                dist = [abs(self.x - target.x), abs(self.y - target.y)]
-                return o_dist + 1.4 * min(x_dist, y_dist)
+                if self.z is None:
+                    o_dist = max(dist) - min(dist)
+                    return o_dist + 1.4 * min(dist)
+                else:
+                    dist.append(abs(self.z - target.z))
+                    d_steps = min(dist)
+                    dist.remove(min(dist))
+                    dist = [x - d_steps for x in dist]
+                    o_dist = max(dist) - min(dist)
+                    return 1.7 * d_steps + o_dist + 1.4 * min(dist)
 
     def getNeighbours(self, includeDiagonal: bool = True, minX: int = -inf, minY: int = -inf,
-                      maxX: int = inf, maxY: int = inf) -> list[Coordinate]:
+                      maxX: int = inf, maxY: int = inf, minZ: int = -inf, maxZ: int = inf) -> list[Coordinate]:
         """
         Get a list of neighbouring coordinates
 
         :param includeDiagonal: include diagonal neighbours
         :param minX: ignore all neighbours that would have an X value below this
         :param minY: ignore all neighbours that would have an Y value below this
+        :param minZ: ignore all neighbours that would have an Z value below this
         :param maxX: ignore all neighbours that would have an X value above this
         :param maxY: ignore all neighbours that would have an Y value above this
+        :param maxZ: ignore all neighbours that would have an Z value above this
         :return: list of Coordinate
         """
         neighbourList = []
+        if self.z is None:
             if includeDiagonal:
                 nb_list = [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]
             else:
@@ -54,15 +80,31 @@ class Coordinate:
             for dx, dy in nb_list:
                 tx = self.x + dx
                 ty = self.y + dy
-            if tx < minX or tx > maxX or ty < minY or ty > maxY:
+                if minX <= tx <= maxX and minY <= ty <= maxY:
-                continue
-
                     neighbourList.append(Coordinate(tx, ty))
+        else:
+            if includeDiagonal:
+                nb_list = [(x, y, z) for x in [-1, 0, 1] for y in [-1, 0, 1] for z in [-1, 0, 1]]
+                nb_list.remove((0, 0, 0))
+            else:
+                nb_list = [
+                    (-1, 0, 0), (0, -1, 0), (1, 0, 0), (0, 1, 0), (0, 0, 1), (0, 0, -1)
+                ]
+
+            for dx, dy, dz in nb_list:
+                tx = self.x + dx
+                ty = self.y + dy
+                tz = self.z + dz
+                if minX <= tx <= maxX and minY <= ty <= maxY and minZ <= tz <= maxZ:
+                    neighbourList.append(Coordinate(tx, ty, tz))
 
         return neighbourList
 
     def getAngleTo(self, target: Coordinate, normalized: bool = False) -> float:
         """normalized returns an angle going clockwise with 0 starting in the 'north'"""
+        if self.z is not None:
+            raise NotImplementedError()
+
         dx = target.x - self.x
         dy = target.y - self.y
         if not normalized:
@@ -75,13 +117,17 @@ class Coordinate:
                 return 180.0 + abs(angle)
 
     def __add__(self, other: Coordinate) -> Coordinate:
+        if self.z is None:
             return Coordinate(self.x + other.x, self.y + other.y)
+        else:
+            return Coordinate(self.x + other.x, self.y + other.y, self.z + other.z)
 
     def __sub__(self, other: Coordinate) -> Coordinate:
+        if self.z is None:
             return Coordinate(self.x - other.x, self.y - other.y)
+        else:
+            return Coordinate(self.x - other.x, self.y - other.y, self.z - other.z)
 
     @staticmethod
     def generate(from_x: int, to_x: int, from_y: int, to_y: int) -> List[Coordinate]:
         return [Coordinate(x, y) for x in range(from_x, to_x + 1) for y in range(from_y, to_y + 1)]
-
-