2018/Day15/Day15.py

from pprint import pprint
from collections import deque
import fileinput
import time


class Elf:
    power = 3
    locations = set()
    deaths = 0

    def __init__(self):
        self.health = 200
        self.targeted = False

    def __repr__(self):
        return "E"


class Goblin:
    enemy = Elf
    power = 3
    locations = set()
    deaths = 0

    def __init__(self):
        self.health = 200
        self.targeted = False

    def __repr__(self):
        return "G"


Elf.enemy = Goblin


class Wall:
    enemy = None

    def __init__(self):
        self.targeted = False

    def __repr__(self):
        return "#"


class Empty:
    enemy = None

    def __init__(self):
        self.targeted = False

    def __repr__(self):
        return "." if not self.targeted else "\033[33m?\033[0m"


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


def read():
    field = []
    Elf.deaths = 0
    Goblin.deaths = 0
    Elf.locations.clear()
    Goblin.locations.clear()
    for y, line in enumerate(fileinput.input()):
        row = []

        for x, char in enumerate(line.rstrip()):
            if char == "E":
                row.append(Elf())
                Elf.locations.add((x, y))
            elif char == "G":
                row.append(Goblin())
                Goblin.locations.add((x, y))
            elif char == "#":
                row.append(Wall())
            elif char == ".":
                row.append(Empty())
            else:
                raise ValueError(char)

        field.append(row)

    return field


def distance(a, b):
    (ax, ay), (bx, by) = a, b
    return abs(ax - bx) + abs(ay - by)


def find_closest(field, pos, targets):
    width = len(field[0])
    height = len(field)

    queue = deque([[pos]])
    seen = {pos}
    paths = []

    while queue:
        path = queue.popleft()
        x, y = path[-1]

        if (x, y) in targets:
            paths.append(path[1:])

        for dx, dy in NEIGHBORS:
            x_, y_ = x + dx, y + dy
            if 0 <= x_ < width and 0 <= y_ < height \
               and type(field[y_][x_]) is Empty \
               and (x_, y_) not in seen:
                queue.append(path + [(x_, y_)])
                seen.add((x_, y_))

    if not paths or [] in paths:
        return None

    return min(paths, key=lambda path: (len(path), path[-1][1], path[-1][0], path[0][1], path[0][0]))



def move(field, x, y, value):
    for row in field:
        for col in row:
            col.targeted = False

    targets = []
    for tx, ty in value.enemy.locations:
        targets.extend([(tx + dx, ty + dy) for dx, dy in NEIGHBORS])

    if not targets:
        raise StopIteration

    for tx, ty in targets:
        field[ty][tx].targeted = True

    closest = find_closest(field, (x, y), targets)
    # print_field(field, highlight=(x, y), path=closest)
    # time.sleep(0.3)

    if closest:
        tx, ty = closest[0]
        if type(field[ty][tx]) is Empty:
            # print(f"{field[y][x]}@{x},{y} moves to {tx},{ty}")
            type(value).locations.discard((x, y))
            type(value).locations.add((tx, ty))
            field[ty][tx] = field[y][x]
            field[y][x] = Empty()
            return tx, ty
    return x, y


def attack(field, x, y, value):
    width = len(field[0])
    height = len(field)

    neighbors = [(field[y + y_][x + x_], (x + x_, y + y_))
                    for x_, y_ in NEIGHBORS
                    if 0 <= y + y_ < width and 0 <= x + x_ < width and \
                    type(field[y + y_][x + x_]) is value.enemy]
    if not neighbors:
        return

    target, (tx, ty) = \
        min(neighbors, key=lambda x: (x[0].health, x[1][1], x[1][0]))
    # print(f"{value}@{x},{y}:{value.health} attacks "
    #         f"{target}@{tx},{ty}:{target.health}->{target.health - value.power}")
    target.health -= value.power

    if target.health <= 0:
        # print(f"{target}@{tx,ty} dies :(")
        type(target).locations.discard((tx, ty))
        type(target).deaths += 1
        field[ty][tx] = Empty()


def print_field(field, steps=None, highlight=None, path=None):
    if steps is not None:
        print(f"round={steps}")
    for y, row in enumerate(field):
        units = []
        for x, value in enumerate(row):
            if highlight == (x, y):
                print(end="\033[31;1m")
            if path and (x, y) in path:
                print(end="\033[44;1m")

            print(value, end="\033[0m")

            if type(value) in (Goblin, Elf):
                units.append(value)

        etc = sorted(units, key=lambda unit: (str(unit), unit.health))
        print(" ", ", ".join(f"{x}({x.health})" for x in etc))
    print()


def step(field, steps):
    # print_field(field, steps)
    seen_units = set()
    for y, row in enumerate(field):
        for x, value in enumerate(row):
            if type(value) not in (Elf, Goblin) or value in seen_units:
                continue
            seen_units.add(value)

            tx, ty = move(field, x, y, value)
            attack(field, tx, ty, value)


def main():
    field = read()

    steps = 0
    while True:
        try:
            step(field, steps)
        except StopIteration:
            total_health = 0
            for row in field:
                for value in row:
                    if type(value) in (Elf, Goblin) and value.health > 0:
                        total_health += value.health

            print("Rounds:", steps)
            print("Health:", total_health)
            print("Elf Deaths:", Elf.deaths)
            print(total_health * steps)
            break

        steps += 1


while True:
    main()
    if Elf.deaths == 0:
        print("yay")
        break

    Elf.power += 1
    print("oh no :(, trying", Elf.power)