from emis_funky_funktions import *

from enum import Enum, unique
from math import isqrt
from typing import List, NamedTuple, Tuple

@unique
class Terrain(int, Enum):
	"""
	Various types of terrain understandable by the routing algorithm.

	Each element of the terrain is associated with an integer value indictating the cost
	that that terrain has on movement speed.  Units are (approximately) measured in
	seconds / kilometer.

	For reference, a typical jogging speed is roughly 500 s/km, and a brisk walking
	speed is about 700 s/km.
	"""
	OPEN_LAND = 510
	ROUGH_MEADOW = 700
	EASY_FOREST = 530
	MEDIUM_FOREST = 666
	WALK_FOREST = 800
	BRUSH = 24000
	WET = 6000
	ROAD = 500
	FOOTPATH = 505
	OOB = 2 ** 62

class Point(NamedTuple):
	x: int
	y: int

	def to_linear(self, array_width: int) -> int:
		"""
		Converts a point to an index in a linear 2D array

		In all cases, this is equivalent to `point.x + array_width * point.y`

		Arguments:
			array_width: The width of the array being indexed

		>>> my_array = [
		...     0, 1, 2,
		...     3, 4, 5,
		...     6, 7, 8
		... ]
		>>> my_array[Point(1, 1).to_linear(3)]
		4
		"""
		return self.x + array_width * self.y
	def __repr__(self):
		return f'({self.x}, {self.y})'

@dataclass(frozen=True)
class OobError:
	"""
	The distance cannot be computed because one point is out of bounds

	This means off the map, not the OOB terrain type.
	"""
	p: Point

@dataclass
class World:
	tiles: Sequence[Tuple[Terrain, int]]
	"""
	The tiles that make up the world

	Each tile should be a tuple containing the terrain type at that location, along with
	the elevation in millimeters at that point.
	"""

	width: int
	"The number of pixels wide the map is"

	lon_scale: int
	"Real world distance represented by each longitudinal 'pixel', in millimeters"

	lat_scale: int
	"Real world distance represented by each latitudinal 'pixel', in millimeters"

	_diag: int
	"Real world distance represented by a diagonal movement, in millimeters"

	_lon_scale2: int
	"Longitudinal scale squared"
	_lat_scale2: int
	"Latitudinal scale squared"
	_diag2: int
	"Diagonal scale squared"

	def __init__(self,
		  tiles: Sequence[Tuple[Terrain, int]],
		  width: int = 395,
		  lon_scale: int = 10_290,
		  lat_scale: int = 7_550
	):
		self.tiles = tiles
		self.width = width
		self.lon_scale = lon_scale
		self.lat_scale = lat_scale

		self._lon_scale2 = lon_scale * lon_scale
		self._lat_scale2 = lat_scale * lat_scale
		self._diag2 = self._lon_scale2 + self._lat_scale2
		self._diag = isqrt(self._diag2)

	def _adjacency(self, p: Point) -> List[Tuple[Point, int]]:
		"""
		Return a series of adjacent points, in any of the eight compass directions

		In addition to each point, a number is returned representing the distance as the
		crow flies between the center of the original point and the center of the adjacent
		point, in millimeters **AND SQUARED**.

		This does not take into account the presence, value, or elevation of tiles
		represented on these points.

		>>> world = World([], lon_scale=3, lat_scale=4)
		>>> world._adjacency(Point(13, 12)) #doctest: +NORMALIZE_WHITESPACE
		[((12, 11), 25),
		 ((13, 11), 16),
		 ((14, 11), 25),
		 ((12, 12),  9),
		 ((14, 12),  9),
		 ((12, 13), 25),
		 ((13, 13), 16),
		 ((14, 13), 25)]
		"""
		return [
			(Point(p.x - 1, p.y - 1), self._diag2     ),
			(Point(p.x    , p.y - 1), self._lat_scale2),
			(Point(p.x + 1, p.y - 1), self._diag2     ),
			(Point(p.x - 1, p.y    ), self._lon_scale2),
			(Point(p.x + 1, p.y    ), self._lon_scale2),
			(Point(p.x - 1, p.y + 1), self._diag2     ),
			(Point(p.x    , p.y + 1), self._lat_scale2),
			(Point(p.x + 1, p.y + 1), self._diag2     )
		]

	def __getitem__(self, loc: Point) -> Tuple[Terrain, int]:
		"""
		Look up terrain cost and elevation by point
		"""
		return self.tiles[loc.x + self.width * loc.y]

	def elevation_difference(self, p1: Point, p2: Point):
		"""
		Compute the elevation difference between two points

		Points do not need to be adjacent!

		>>> world = World([(Terrain.BRUSH, 1_000), (Terrain.BRUSH, 1_100)])
		>>> world.elevation_difference(Point(0, 0), Point(1, 0))
		100
		"""
		return self[p2][1] - self[p1][1]

	def neighbors(self, loc: Point) -> List[Tuple[Point, int]]:
		"""
		Produce a list of valid neighbors surrounding a given point.

		Neighbors in any of the eight cardinal directions will be considered, but only
		neighbors which are in bounds and correspond to a point on the map will be
		returned.

		In addition to each neighboring point, a time cost associated with moving from the
		center of this point to the center of the neighbor is included.  This time cost is
		measured in microseconds, and factors in three-dimensional distance as well as the
		movement speed allowed by the terrain.

		Example:
			In this example, we construct a simple world of only four tiles, and attempt
			to list the neighbors of the upper left tile.  Three other tiles can be
			considered.  Of these, the upper right tile is out of bounds terrain, leaving
			only the bottom two tiles.

			Each of these tiles is returned.  Notice, however, that the cost of moving to
			the bottom right tile is slightly larger than one might expect.  This is
			because diagonal travel is more expensive than latitudinal or longitudinal
			travel.

			>>> world = World(
			...    [ (Terrain.OPEN_LAND,    1_000), (Terrain.OOB,           950)
			...    , (Terrain.ROUGH_MEADOW, 1_080), (Terrain.EASY_FOREST, 1_010)
			...    ],
			...    width = 2, # Our simple world is only two tiles wide
			...    lon_scale = 500, # Pick an easy number for example
			...    lat_scale = 400  # and remember that these are in millimeters!
			... )
			>>> world.neighbors(Point(0, 0))
			[((0, 1), 246235), ((1, 1), 332800)]

			Let's look at how the time cost for the point at (1, 1) is computed.

			First, we find the distance as the crow flies between the center of the two
			points.  This is equal to the square root of the latitudinal distance squared
			plus the longitudinal distance squared.

			>>> isqrt(500 ** 2 + 400 ** 2)
			640

			So now we know that the distance between the centers of these two points is
			640 millimeters as the crow flies.

			Next, we factor in elevation.  The elevation difference between (0, 0) and
			(1, 1) is just `10` millimeters.  This should make an inconsiquential
			difference, but we compute it nonetheless.

			>>> isqrt(10 ** 2 + 640 ** 2)
			640

			Sure enough, the three-dimensional distance between these two points is still
			640 millimeters.

			Now, we need to know how difficult it will be to move over the terrain.  Half
			of the time, we expect to be moving over `OPEN_LAND`, which has a movement
			speed of 510 microseconds / millimeter.  Once we cross the border between the
			two tiles at the half-way point, we'll be travelling through `EASY_FOREST`,
			which has a movement speed of 530 microseconds / millimeter, just slightly
			slower.

			Since we'll our trip is perfectly balanced between these two terrain types, we
			can compute the average movement speed of the trip to be the average of these
			two values.

			>>> (510 + 530) // 2
			520

			So now we know we'll be travelling at an average of 520 microseconds /
			millimeter.

			Now that we know both the distance we need to travel and the speed we'll be
			travelling at, we can multiply them together to get our actual travel time.

			>>> 640 * 520
			332800

			And there it is!  The travel time returned by the `.neighbors` function!
		"""
		loc_terrain, loc_elevation = self[loc]
		return [
			(
				adj_point,
				(
					# 2 * Movement speed (seconds / km = microseconds / millimeter)
					(
						loc_terrain +
						adj_terrain
					)

					# * Distance travelled (millimeters) = 2 * Time cost (microseconds)
					* isqrt(
						# Crow Distance Squared
						crow_distance2
						# + Elevation Change Squared = 3D distance squared
						+ (loc_elevation - adj_elevation) ** 2
					)

					# / 2 = Time cost (microseconds)
					// 2
				)
			)
			for (adj_point, crow_distance2) in self._adjacency(loc)
			for (adj_terrain, adj_elevation) in (self[adj_point],)
			if adj_point.x >= 0
				and  adj_point.y >= 0
				and adj_point.x < self.width
				and adj_point.y < len(self.tiles) // self.width
				and adj_terrain < 4294967296
		]

	def heuristic(self, a: Point, b: Point) -> int:
		"""
		Estimate the time it will take to travel between two points

		The following assumptions are made to speed up the process, at the cost of
		accuracy:
		- All tiles have the exact same movement speed (500 milliseconds / meter)
		- All tiles are in-bounds
		- All tiles are at the same elevation

		This does NOT assume that we can travel at angles other than multiples of
		45 degrees, however, as this would not be beneficial to the computation nor the
		accuracy.

		Because this algorithm ignores both terrain types and elevations, it does not
		access world data at all, and the results of the computation depend exclusively on
		the start point, the finish point, and the longitudinal/latitudinal scales.

		>>> world = World([], lon_scale = 500, lat_scale = 400)
		>>> world.heuristic(Point(0, 0), Point(3, 5))
		1360000

		For a slower algorithm that also includes elevation, see `calculate_distance()`.
		"""

		# Taxicab distance in each direction
		lon_tiles_raw = abs(a.x - b.x)
		lat_tiles_raw = abs(a.y - b.y)

		# The number of moves necessary, allowing diagonal moves
		vh_moves = lon_tiles_raw - lat_tiles_raw
		if vh_moves > 0:
			lon_moves_real = vh_moves
			lat_moves_real = 0
			diag_moves_real = lat_tiles_raw
		else:
			lat_moves_real = -vh_moves
			lon_moves_real = 0
			diag_moves_real = lon_tiles_raw

		# Total distance necessary, not counting elevation
		total_flat_distance = (
			lon_moves_real * self.lon_scale +
			lat_moves_real * self.lat_scale +
			diag_moves_real * self._diag
		)

		# TODO: Test whether adding in elevation is beneficial

		estimated_speed = 500 # milliseconds / meter = microseconds / millimeters

		return estimated_speed * total_flat_distance

	def calculate_distance(self, a: Point, b:Point) -> Result[int, OobError]:
		"""
		Calculate the distance between the centers of two tiles, incl elevation difference

		Looks up the elevation at both points **a** and **b**, converting them into points
		in 3D space, then computes the integer distance between

		For a faster algorithm that does not include elevation, see `heuristic()`.

		>>> world = World( # We instantiate a simple, small world
		...    [ (Terrain.OPEN_LAND,    1_000), (Terrain.BRUSH,       850), (Terrain.WET, 500)
		...    , (Terrain.ROUGH_MEADOW,   950), (Terrain.EASY_FOREST, 750), (Terrain.WET, 500)
		...    ],
		...    width = 3, # Our simple world is only two tiles wide
		...    lon_scale = 600, # Pick an easy number for example
		...    lat_scale = 500
		... )
		>>> world.calculate_distance(Point(0, 0), Point(2, 1))
		Ok(1392)
		>>> world.calculate_distance(Point(0, 0), Point(2, 2))
		Err(OobError(p=(2, 2)))

		Notice that this is distance as-the-crow-flies between two points in 3D space.  If
		you're looking for actual distance, you must either use two adjacent points, or
		first pathfind between those two points and calculate the distance that way.
		"""
		lon_dist = abs(a.x - b.x) * self.lon_scale
		lat_dist = abs(a.y - b.y) * self.lat_scale
		try:
			a_elev = self[a][1]
		except IndexError:
			return Err(OobError(a))
		try:
			b_elev = self[b][1]
		except IndexError:
			return Err(OobError(b))
		elev_dist = abs(a_elev - b_elev)
		return Ok(isqrt(lon_dist * lon_dist + lat_dist * lat_dist + elev_dist * elev_dist))

	@tco_rec
	def calculate_path_length(self, points: Sequence[Point], addend: int = 0) -> Result[int, OobError]:
		"""
		Calculate the length of a path to the greatest degree of accuracy possible

		Points are expected to be a sequence of at least two adjacent points.  The
		returned distance will be the distance travelling between the centers of the tiles
		of each tile in sequence, including elevation changes.

		Asequential points will not result in an error, but will result in a drop in
		accuracy.  Diagonal moves to adjacent tiles, however, are valid and allowed.

		Any points which fall outside the bounds of the map will result in an `OobError`
		indicating the first point which was out of bounds.

		On a successful run, the returned units will be in the units the world was
		instantiated with, typically millimeters.

		>>> world = World( # We instantiate a simple, small world
		...    [ (Terrain.OPEN_LAND,    1_000), (Terrain.BRUSH,       850), (Terrain.WET, 500)
		...    , (Terrain.ROUGH_MEADOW,   950), (Terrain.EASY_FOREST, 750), (Terrain.WET, 500)
		...    ],
		...    width = 3, # Our simple world is only two tiles wide
		...    lon_scale = 600, # Pick an easy number for example
		...    lat_scale = 500
		... )
		>>> world.calculate_path_length([Point(0,0), Point(1, 0), Point(2,1)])
		Ok(1473)

		Calling this method with only two points is equivalent to calling
		`calculate_distance()`.

		>>> world.calculate_path_length([Point(0, 0), Point(2, 1)])
		Ok(1392)
		>>> world.calculate_distance(Point(0, 0), Point(2, 1))
		Ok(1392)
		"""
		match points:
			case [a, b, *rest]:
				match self.calculate_distance(a, b):
					case Ok(this_dist):
						return Recur(self, (b, *rest), addend + this_dist)
					case err:
						return Return(Ok(err))
			case _: # single or empty
				return Return(Ok(addend))


if __name__ == '__main__':
	import doctest
	doctest.testmod()