coreylowman / allyourturtles Goto Github PK

max_halite = max(gmap.values())

n = len(starts)
standard = (None,) * n

def actual_position(s, i):
    return s[i + n] if s[i + n] is not None else s[i]

def heuristic(s):
    return sum(2 * dist(actual_position(s, i), goals[i]) for i in range(n))

def neighbors(s, i):
    for position in position_neighbors(s[i], width, height):
        if position in s[n:]:
            continue
        neighbor = list(s)
        neighbor[i + n] = position
        if i == n - 1:
            yield tuple(neighbor[n:]) + standard, 1
        else:
            yield tuple(neighbor), 0

def _reconstruct_path(prev_by_node, current, n):
    total_path = [current[:n]]
    while current in prev_by_node:
        current = prev_by_node[current]
        if current[n:] == standard:
            total_path.append(current[:n])
    return list(reversed(total_path))

closed_set = set()
open_set = set()
g_score = defaultdict(lambda: math.inf)
h_score = defaultdict(lambda: math.inf)
f_score = defaultdict(lambda: math.inf)
c_score = defaultdict(int)
came_from = {}
halite_at = {}
extractions_at = {}
time_at = {}

start = tuple(starts) + standard

open_set.add(start)
g_score[start] = 0
h_score[start] = heuristic(start)
f_score[start] = g_score[start] + h_score[start]
c_score[start] = 0
halite_at[start] = list(halite)
extractions_at[start] = []
time_at[start] = 0

while len(open_set) > 0:
    # TODO break ties by heuristic and then conflicts
    current = min(open_set, key=f_score.get)

    open_set.remove(current)

    if current[n:] == standard:
        closed_set.add(current)
        if current[:n] == goals:
            return _reconstruct_path(came_from, current, n)

    actuals = tuple(actual_position(current, i) for i in range(n))
    t = time_at[current]
    halite_left = list(halite_at[current])
    halite_on_ground = list(gmap[actuals[i]] for i in range(n))
    for pos, _, amt in extractions_at[current]:
        for i in range(n):
            if pos == actuals[i]:
                halite_on_ground[i] -= amt

    agent_i = current.index(None) - n
    raw_move_cost = halite_on_ground[agent_i] / 4
    raw_extracted = halite_on_ground[agent_i] / 10
    move_cost = halite_on_ground[agent_i] / max_halite

    for neighbor, dt in neighbors(current, agent_i):
        if neighbor in closed_set:
            continue

        nt = t + dt
        pos = neighbor[agent_i + n]

        if pos in illegals_by_t[nt]:
            continue

        cost = 0 if current[agent_i] == neighbor[agent_i] else move_cost
        d = cost + 1
        g = g_score[current] + d

        if g > max_cost:
            continue

        if neighbor not in open_set:
            open_set.add(neighbor)
        elif g >= g_score[neighbor]:
            continue

        came_from[neighbor] = current
        g_score[neighbor] = g
        h_score[neighbor] = heuristic(neighbor)
        f_score[neighbor] = g_score[neighbor] + h_score[neighbor]
        time_at[neighbor] = nt
        c_score[neighbor] = c_score[current] + int(pos in conflicts_by_t[nt])

        if current[agent_i] == neighbor[agent_i]:
            extracted = raw_extracted
            halite_left[agent_i] += extracted
            halite_at[neighbor] = halite_left
            extractions_at[neighbor] = extractions_at[current] + [(current[agent_i], nt, extracted)]
        else:
            halite_left[agent_i] -= raw_move_cost
            halite_at[neighbor] = halite_left
            extractions_at[neighbor] = deepcopy(extractions_at[current])

connected_components = []
component_by_pos = {p: None for p in OBSTACLES}
unclaimed = positions - OBSTACLES
while len(unclaimed) > 0:
    open_set = {unclaimed.pop()}
    closed_set = set()
    component = len(connected_components)

    while len(open_set) > 0:
        current = open_set.pop()
        closed_set.add(current)
        component_by_pos[current] = component
        unclaimed.discard(current)

        for neighbor in cardinal_neighbors(current):
            if neighbor in OBSTACLES or neighbor in closed_set:
                continue
            open_set.add(neighbor)

    connected_components.append(closed_set)