py-tools/coordinate.py
2021-12-05 07:36:38 +01:00

160 lines
6.4 KiB
Python

from __future__ import annotations
from dataclasses import dataclass
from enum import Enum
from math import gcd, sqrt, inf, atan2, degrees
from typing import Union, List, Optional
from tools import compare
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: 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 in 2D, 1.7 in 3D)
:return: Distance to Target
"""
assert isinstance(target, Coordinate)
assert isinstance(mode, DistanceAlgorithm)
assert isinstance(includeDiagonals, bool)
if mode == DistanceAlgorithm.PYTHAGORAS:
if self.z is None:
return sqrt(abs(self.x - target.x) ** 2 + abs(self.y - target.y) ** 2)
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:
return abs(self.x - target.x) + abs(self.y - target.y) + abs(self.z - target.z)
else:
dist = [abs(self.x - target.x), abs(self.y - target.y)]
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, 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:
nb_list = [(-1, 0), (1, 0), (0, -1), (0, 1)]
for dx, dy in nb_list:
tx = self.x + dx
ty = self.y + dy
if minX <= tx <= maxX and minY <= ty <= maxY:
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() # which angle?!?!
dx = target.x - self.x
dy = target.y - self.y
if not normalized:
return degrees(atan2(dy, dx))
else:
angle = degrees(atan2(dx, dy))
if dx >= 0:
return 180.0 - angle
else:
return 180.0 + abs(angle)
def getLineTo(self, target: Coordinate) -> List[Coordinate]:
diff = target - self
if self.z is None:
steps = gcd(diff.x, diff.y)
step_x = diff.x // steps
step_y = diff.y // steps
return [Coordinate(self.x + step_x * i, self.y + step_y * i) for i in range(steps + 1)]
else:
steps = gcd(diff.x, diff.y, diff.z)
step_x = diff.x // steps
step_y = diff.y // steps
step_z = diff.z // steps
return [Coordinate(self.x + step_x * i, self.y + step_y * i, self.z + step_z * i) for i in range(steps + 1)]
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,
from_z: int = None, to_z: int = None) -> List[Coordinate]:
if from_z is None and to_z is None:
return [Coordinate(x, y) for x in range(from_x, to_x + 1) for y in range(from_y, to_y + 1)]
else:
return [
Coordinate(x, y, z)
for x in range(from_x, to_x + 1)
for y in range(from_y, to_y + 1)
for z in range(from_z, to_z + 1)
]