# -*- coding: utf-8 -*-
import sys
from abc import ABC
from math import inf, sqrt, floor
from typing import List, Tuple, Iterator, Callable, Any

from aoc import BaseAssignment


class Map:
    def __init__(self, map: List[str]):
        self.map = map

        self.height = len(map)
        self.width = len(map[0])

        try:
            self.start = next(self.find_position("S"))
        except StopIteration:
            self.start = (0, 0)

        try:
            self.end = next(self.find_position("E"))
        except StopIteration:
            self.start = (0, 0)

    def find_position(self, value: str) -> Iterator[Tuple[int, int]]:
        for y in range(self.height):
            for x in range(self.width):
                if self.map[y][x] == value:
                    yield x, y

    def neighbours(self, x: int, y: int) -> Iterator[Tuple[int, int]]:
        for dy in range(max(0, y - 1), min(self.height, y + 2)):
            for dx in range(max(0, x - 1), min(self.width, x + 2)):
                if (
                    (dy == y and dx == x)
                    or (dy == y - 1 and dx == x - 1)
                    or (dy == y - 1 and dx == x + 1)
                    or (dy == y + 1 and dx == x - 1)
                    or (dy == y + 1 and dx == x + 1)
                ):
                    continue
                yield dx, dy

    def elevation(self, x: int, y: int):
        character = self.map[y][x]

        match character:
            case "S":
                elevation = ord("a")
            case "E":
                elevation = ord("z")
            case _:
                elevation = ord(character)

        return elevation - 96


class Assignment(BaseAssignment, ABC):
    def distance(self, map: Map, a: Tuple[int, int], b: Tuple[int, int]):
        distance = map.elevation(*b) - map.elevation(*a)

        if distance > 1:
            return inf

        return 1

    def heuristic(self, map: Map, a: Tuple[int, int]):
        dx = abs(map.start[0] - a[0])
        dy = abs(map.start[1] - a[1])
        return floor(sqrt(dx**2 + dy**2))

    def a_star(
        self,
        map: Map,
        start: Tuple[int, int],
        end: Tuple[int, int],
        distance: Callable[[Tuple[int, int], Tuple[int, int]], int],
        heuristic: Callable[[Tuple[int, int]], int],
    ) -> List[Tuple[int, int]]:
        open = {start}
        came_from = {}

        g_scores = {start: 0}

        f_scores = {start: heuristic(start)}

        while True:
            try:
                current = sorted(
                    [item for item in open], key=lambda item: f_scores.get(item, inf)
                )[0]
            except IndexError:
                raise RuntimeError("No path found")

            if current == end:
                path = [current]

                while current in came_from:
                    current = came_from[current]
                    path = [current, *path]

                return path

            open.remove(current)

            for neighbour in map.neighbours(*current):
                g_score = g_scores.get(current, inf) + distance(current, neighbour)

                if g_score < g_scores.get(neighbour, inf):
                    came_from[neighbour] = current
                    g_scores[neighbour] = g_score
                    f_scores[neighbour] = g_score + heuristic(neighbour)

                    if neighbour not in open:
                        open.add(neighbour)


class AssignmentOne(Assignment):
    example_result = 31

    def run(self, input: Iterator) -> Any:
        map = Map(map=list(input))

        path = self.a_star(
            map,
            map.start,
            map.end,
            distance=lambda a, b: self.distance(map, a, b),
            heuristic=lambda a: self.heuristic(map, a),
        )

        return len(path) - 1


class AssignmentTwo(Assignment):
    example_result = 29

    def run(self, input: Iterator) -> Any:
        map = Map(list(input))
        lengths = []

        for a in map.find_position("a"):
            try:
                path = self.a_star(
                    map,
                    a,
                    map.end,
                    distance=lambda a, b: self.distance(map, a, b),
                    heuristic=lambda a: self.heuristic(map, a),
                )

                lengths.append(len(path) - 1)
            except RuntimeError:
                continue

        return sorted(lengths)[0]