from lib import *
from collections import deque

EXAMPLE = """#.#####################
#.......#########...###
#######.#########.#.###
###.....#.>.>.###.#.###
###v#####.#v#.###.#.###
###.>...#.#.#.....#...#
###v###.#.#.#########.#
###...#.#.#.......#...#
#####.#.#.#######.#.###
#.....#.#.#.......#...#
#.#####.#.#.#########v#
#.#...#...#...###...>.#
#.#.#v#######v###.###v#
#...#.>.#...>.>.#.###.#
#####v#.#.###v#.#.###.#
#.....#...#...#.#.#...#
#.#########.###.#.#.###
#...###...#...#...#.###
###.###.#.###v#####v###
#...#...#.#.>.>.#.>.###
#.###.###.#.###.#.#v###
#.....###...###...#...#
#####################.#
"""

SLOPES = {
    "^": (-1, 0),
    ">": (0, 1),
    "v": (1, 0),
    "<": (0, -1),
}

def first(input):
    g = input.grid2()
    start = (0, 1)
    end = (g.n_rows - 1, g.n_cols - 2)
    longest = 0
    paths = [(set([start]), start)]
    while True:
        new_paths = []
        for p in paths:
            hist, pos = p
            for d in g.COORDS_ORTH:
                nb = add2(pos, d)
                if nb[0] < 0 or nb[0] >= g.n_rows or nb[1] < 0 or nb[1] >= g.n_cols:
                    continue
                c = g[nb]
                if c in SLOPES.keys() and d != SLOPES[c]:
                    continue
                if c == "#" or nb in hist:
                    continue
                if nb == end:
                    l = len(hist)
                    if l > longest:
                        longest = l
                    continue
                nhist = hist.copy()
                nhist.add(nb)
                new_paths.append((nhist, nb))
        paths = new_paths
        if len(paths) == 0:
            break
    return longest

def solve(input: Input, second=False):
    if not second:
        return first(input)

    g = input.grid2()
    start = (0, 1)
    end = (g.n_rows - 1, g.n_cols - 2)

    seen = set()
    q = deque([[start, (1, 1)]])

    # The intuition is that we can brute force much quicker if we have a pure
    # graph instead of following the maze along the whole time. So, we create
    # a graph from the maze and then brute force on the maze.
    sg = {start: set()} # {node: {(node, dist), ...}}

    while q:
        trail = q.popleft()
        pos = trail[-1]
        while True:
            nbs = []
            for d in g.COORDS_ORTH:
                nb = add2(pos, d)
                if nb[0] < 0 or nb[0] >= g.n_rows or nb[1] < 0 or nb[1] >= g.n_cols:
                    continue
                if g[nb] == "#" or nb == trail[-2]:
                    continue
                nbs.append(nb)
            if len(nbs) == 1:
                pos = nbs[0]
                trail.append(pos)
            else:
                break

        if not pos in sg:
            sg[pos] = set()

        dist = len(trail) - 1
        sg[trail[0]].add((pos, dist))
        sg[pos].add((trail[0], dist))

        seen.add(pos)
        for nb in nbs:
            if not nb in seen:
                seen.add(nb)
                q.append([pos, nb])

    # for key, value in sg.items():
    #     print(key, value)

    # Brute force in bf order.
    longest = 0
    q = deque([(set(), start, 0)])
    while q:
        hist, pos, dist = q.popleft()
        if pos == end:
            if dist > longest:
                longest = dist
            continue

        for nb, d in sg[pos]:
            if nb in hist:
                continue
            nhist = hist.copy()
            nhist.add(nb)
            q.append((nhist, nb, dist + d))

    return longest


def main():
    DAY_INPUT = "i23.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))

if __name__ == "__main__":
    main()