from lib import Grid2D, add2
from collections import deque


DIRS = [Grid2D.N, Grid2D.W, Grid2D.E, Grid2D.S]


def find_move(unit, grid):
    assert grid[unit] in "GE"
    target_fields = set()
    enemies = grid.find("G") if grid[unit] == "E" else grid.find("E")
    for e in enemies:
        for dir in DIRS:
            nb = add2(e, dir)
            if grid[nb] == ".":
                target_fields.add(nb)

    # we are already next to an enemy
    if unit in target_fields:
        return None

    seen = set()
    states = deque([(unit, [])])
    while states:
        pos, hist = states.popleft()

        if pos in seen:
            continue
        else:
            seen.add(pos)

        for dir in DIRS:
            nb = add2(pos, dir)
            if nb in target_fields:
                if hist:
                    return hist[0]
                else:
                    return nb
            elif grid[nb] == ".":
                nhist = hist + [nb]
                states.append((nb, nhist))

    # we did not find a path to an enemy
    return None


def part_2(data):
    for elf_attack in range(3, 100):
        g = Grid2D(data)
        pos_to_health = {pos: 200 for pos in g.find("GE")}
        len_elfs_orig = len(g.find("E"))

        rounds = 0
        done = False
        while not done:
            for unit in sorted(list(pos_to_health.keys())):
                if not unit in pos_to_health:
                    continue  # unit died in the meantime
                enemy_type = "G" if g[unit] == "E" else "E"

                # move stage
                for dir in DIRS:
                    nb = add2(unit, dir)
                    if g[nb] == enemy_type:
                        break  # we are already next to an enemy
                else:
                    npos = find_move(unit, g)
                    if npos is not None:
                        assert npos not in pos_to_health and g[npos] == "."
                        pos_to_health[npos] = pos_to_health[unit]
                        del pos_to_health[unit]
                        g[npos] = g[unit]
                        g[unit] = "."
                        unit = npos

                # attack stage
                enemy_type = "G" if g[unit] == "E" else "E"
                enemy_hit = 999
                enemy = None
                for dir in DIRS:
                    nb = add2(unit, dir)
                    if g[nb] == enemy_type and pos_to_health[nb] < enemy_hit:
                        enemy_hit = pos_to_health[nb]
                        enemy = nb

                if enemy is not None:
                    if g[enemy] == "G":
                        pos_to_health[enemy] -= elf_attack
                    else:
                        pos_to_health[enemy] -= 3
                    if pos_to_health[enemy] <= 0:
                        del pos_to_health[enemy]
                        g[enemy] = "."

                if len(g.find("E")) == 0 or len(g.find("G")) == 0:
                    done = True
                    break
            else:
                rounds += 1

        if elf_attack == 3:
            print(sum(pos_to_health.values()) * rounds)

        if len(g.find("E")) == len_elfs_orig:
            print(sum(pos_to_health.values()) * (rounds))
            break


def main():
    input_file = __file__.replace(".py", ".txt")
    with open(input_file) as f:
        data = f.read()
    part_2(data)


if __name__ == "__main__":
    main()