from lib import get_data, str_to_ints, Grid2D, add2
from collections import defaultdict
from d9 import Amp

DIRS = [
    (-1, 0),
    (0, 1),
    (1, 0),
    (0, -1),
]

DIRCHAR = list("^>v<")


def find_path(g):
    (pos,) = g.find("^><v")
    assert g[pos] == "^"

    to_visit = set(g.find("#"))
    seen = defaultdict(int)
    dir = DIRS[0]  # up
    seen[pos] = 1

    def best_neighbor(pos):
        best_nb = None
        seen_count = 20
        for nb in g.neighbors_ort(pos):
            if nb in to_visit:
                return nb
            if nb in seen and seen[nb] < seen_count:
                seen_count = seen[nb]
                best_nb = nb
        assert best_nb is not None
        return best_nb

    path = ""
    while to_visit:

        new_pos = add2(pos, dir)
        if new_pos in to_visit or (new_pos in seen and seen[new_pos] < 2):
            path += "F"
            g[pos] = "#"
            pos = new_pos
            if pos in to_visit:
                to_visit.remove(pos)
            g[pos] = DIRCHAR[DIRS.index(dir)]
            seen[pos] += 1
        else:
            best_nb = best_neighbor(pos)
            while add2(pos, dir) != best_nb:
                path += "R"
                dir = DIRS[(DIRS.index(dir) + 1) % len(DIRS)]
                g[pos] = DIRCHAR[DIRS.index(dir)]

        # For debugging:
        # g.print()
        # input()
        # print()
    path = path.replace("RRR", "L")
    return path


def part_1(data):
    xs = str_to_ints(data)
    a = Amp(xs)
    text = ""
    while not a.done:
        a.go()
        while a.outputs:
            o = a.pop()
            text += chr(o)

    result = 0
    g = Grid2D(text)
    for r in range(g.n_rows):
        for c in range(g.n_cols):
            nbs = g.neighbors_ort((r, c))
            if (
                g[(r, c)] == "#"
                and len(nbs) == 4
                and all([g[(nr, nc)] == "#" for nr, nc in nbs])
            ):
                # g[(r, c)] = 'o'
                result += r * c
    print(result)
    # g.print()

    # Merge F commands into counts
    path = find_path(g)
    path = list(path)
    new_path = []
    i = 0
    while i < len(path):
        if path[i] == "R":
            new_path.append("R")
            i += 1
        elif path[i] == "L":
            new_path.append("L")
            i += 1
        elif path[i] == "F":
            count = 0
            while i < len(path) and path[i] == "F":
                count += 1
                i += 1
            new_path.append(str(count))
        else:
            assert False
    path = new_path
    print("Manually translate into commands:", "".join(path))

    # manually created from above output
    inst = (
        "A,A,B,C,B,C,B,C,A,C\n"
        "R,6,L,8,R,8\n"
        "R,4,R,6,R,6,R,4,R,4\n"
        "L,8,R,6,L,10,L,10\n"
        "n\n"
    )

    xs = str_to_ints(data)
    xs[0] = 2
    a = Amp(xs)
    for c in inst:
        a.feed(ord(c))
    while not a.done:
        a.go()
    print(a.outputs[-1])


def main():
    data = get_data(__file__)
    part_1(data)


if __name__ == "__main__":
    main()