from tools.aoc import AOCDay
from typing import Any
from tools.coordinate import Coordinate, DistanceAlgorithm

DIR = {
    'R': Coordinate(1, 0),
    'L': Coordinate(-1, 0),
    'U': Coordinate(0, -1),
    'D': Coordinate(0, 1),
}


class Day(AOCDay):
    inputs = [
        [
            (13, "input9_test"),
            (6498, "input9"),
        ],
        [
            (1, "input9_test"),
            (36, "input9_test2"),
            (2531, "input9"),
        ]
    ]

    def get_visited(self, rope_length: int = 2) -> int:
        rope = [Coordinate(0, 0) for _ in range(rope_length)]
        visited = set()
        visited.add(rope[-1])

        for line in self.getInput():
            direction, count = line.split(" ")
            for _ in range(int(count)):
                rope[0] += DIR[direction]
                for knot in range(1, rope_length):
                    dist = rope[knot].getDistanceTo(
                        rope[knot - 1],
                        algorithm=DistanceAlgorithm.MANHATTAN,
                        includeDiagonals=True
                    )
                    step = rope[knot - 1] - rope[knot]
                    if dist == 2 or abs(step.x) == 2 and abs(step.y) == 2:
                        rope[knot] += Coordinate(step.x // 2, step.y // 2)
                    elif dist > 2:
                        if abs(step.x) == 2:
                            rope[knot] += Coordinate(step.x // 2, step.y)
                        else:
                            rope[knot] += Coordinate(step.x, step.y // 2)

                visited.add(rope[-1])

        return len(visited)

    def part1(self) -> Any:
        return self.get_visited()

    def part2(self) -> Any:
        return self.get_visited(10)


if __name__ == '__main__':
    day = Day(2022, 9)
    day.run(verbose=True)