from tools.aoc import AOCDay
from tools.coordinate import Coordinate
from tools.grid import Grid
from typing import Any


def build_grid(max_num: int, p2: bool = False) -> (Grid, Coordinate, int):
    grid = Grid(0)
    grid.set(Coordinate(0, 0), 1)
    last_coord = Coordinate(0, 0)
    last_num = max_num
    dirs = [(-1, 0), (0, 1), (1, 0), (0, -1)]
    dir_index = 2
    cur_dir = dirs[dir_index]
    check_dir = dirs[dir_index + 1]

    for x in range(2, max_num + 1):
        last_coord += Coordinate(*cur_dir)
        if not p2:
            grid.set(last_coord, x)
        else:
            sum = grid.getNeighbourSum(last_coord)
            grid.set(last_coord, sum)
            if sum > max_num:
                return grid, last_coord, sum
        if not grid.get(last_coord + Coordinate(*check_dir)):
            dir_index = (dir_index + 1) % 4
            cur_dir = dirs[dir_index]
            check_dir = dirs[(dir_index + 1) % 4]

    return grid, last_coord, last_num


class Day(AOCDay):
    inputs = [
        [
            (31, "input3_test"),
            (438, "input3")
        ],
        [
            (266330, "input3")
        ]
    ]

    def part1(self) -> Any:
        g, c, _ = build_grid(int(self.getInput()))
        return abs(c.x) + abs(c.y)

    def part2(self) -> Any:
        g, c, i = build_grid(int(self.getInput()), True)
        return i


if __name__ == '__main__':
    day = Day(2017, 3)
    day.run(verbose=True)