from lib import *

EXAMPLE = """
2413432311323
3215453535623
3255245654254
3446585845452
4546657867536
1438598798454
4457876987766
3637877979653
4654967986887
4564679986453
1224686865563
2546548887735
4322674655533
"""

def solve(i: Input, second=False):
    g = i.grid2()
    starts = [((0, 0), (0, None))]

    def is_goal(node):
        pos, _ = node
        return pos == (g.n_cols - 1, g.n_rows - 1)

    def neighbors(node):
        pos, dirs = node
        repeats, prev_dir = dirs
        nbs = []
        for dir in g.dirs_ort():
            if second:
                if repeats < 4 and prev_dir is not None and prev_dir != dir:
                    continue
                if repeats == 10 and prev_dir == dir:
                    continue
            else:
                if repeats == 3 and prev_dir == dir:
                    continue
            if prev_dir == g.flip_ort(dir):
                continue

            nb = add2(pos, dir)
            if nb not in g:
                continue
            nbs.append((nb, (repeats + 1 if dir == prev_dir else 1, dir)))
        return nbs

    def h(node):
        pos, _ = node
        return abs(g.n_rows - 1 - pos[0]) + abs(g.n_cols - 1 - pos[1])

    def d(_, b):
        pos, _ = b
        if pos == (0, 0):
            return 0
        return int(g[pos])

    a = A_Star(starts, is_goal, h, d, neighbors)
    return a.cost

def main():
    DAY_INPUT = "i17.txt"
    print("Example  1:", solve(Input(EXAMPLE)))
    print("Solution 1:", solve(Input(DAY_INPUT)))
    print("Example  2:", solve(Input(EXAMPLE), True))
    print("Solution 2:", solve(Input(DAY_INPUT), True))
    return

if __name__ == "__main__":
    main()