# -*- 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
from aoc.mixins import AStarMixin


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, AStarMixin[Tuple[int, int]], 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


class AssignmentOne(Assignment):
    example_result = 31

    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 run(self, input: Iterator) -> Any:
        map = Map(map=list(input))

        path = self.a_star(
            map.start,
            lambda coordinate, _: coordinate == map.end,
            neighbours=lambda a: map.neighbours(*a),
            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))

        path = self.a_star(
            map.end,
            lambda coordinate, _: map.elevation(*coordinate) == 1,
            neighbours=lambda a: map.neighbours(*a),
            distance=lambda a, b: self.distance(map, b, a),
        )

        return len(path) - 1