from z3 import *
import networkx as nx

# solver for Jane Street's bug byte, see:
# https://www.janestreet.com/bug-byte/

white_17, white_3, green_31, green_19_23, top_star, green_8 = [vertex for vertex in range(0, 6)]
green_6_9_16, white_54, white_49, white_60, white_79, white_75 = [vertex for vertex in range(6, 12)]
empty_left, white_29, empty_right, white_39, white_25, bot_star = [vertex for vertex in range(12, 18)]

vertex_names = dict([
    (white_17, "white_17"),
    (white_3, "white_3"),
    (green_31, "green_31"),
    (green_19_23, "green_19_23"),
    (top_star, "top_star"),
    (green_8, "green_8"),
    (green_6_9_16, "green_6_9_16"),
    (white_54, "white_54"),
    (white_49, "white_49"),
    (white_60, "white_60"),
    (white_79, "white_79"),
    (white_75, "white_75"),
    (empty_left, "empty_left"),
    (white_29, "white_29"),
    (empty_right, "empty_right"),
    (white_39, "white_39"),
    (white_25, "white_25"),
    (bot_star, "bot_star")
])

solver = Solver()
graph = nx.Graph()

graph.add_edge(white_17, green_19_23, weight=Int("(white_17, green_19_23)"))
graph.add_edge(white_17, white_3, weight=Int("(white_17, white_3)"))
graph.add_edge(white_3, top_star, weight=Int("(white_3, top_star)"))
graph.add_edge(green_31, white_54, weight=Int("(green_31, white_54)"))
graph.add_edge(green_19_23, white_54, weight=Int("(green_19_23, white_54)"))
graph.add_edge(green_19_23, top_star, weight=12)
graph.add_edge(top_star, white_49, weight=Int("top_star, white_49"))
graph.add_edge(green_8, white_49, weight=Int("(green_8, white_49)"))
graph.add_edge(green_6_9_16, white_60, weight=Int("(green_6_9_16, white_60)"))
graph.add_edge(white_54, white_79, weight=Int("(white_54, white_79)"))
graph.add_edge(white_54, white_60, weight=Int("(white_54, white_60)"))
graph.add_edge(white_49, white_60, weight=Int("(white_49, white_60)"))
graph.add_edge(white_49, white_75, weight=20)
graph.add_edge(white_60, white_79, weight=24)
graph.add_edge(white_60, white_75, weight=Int("(white_60, white_75)"))
graph.add_edge(white_79, empty_left, weight=Int("(white_79, empty_left)"))
graph.add_edge(white_79, white_39, weight=7)
graph.add_edge(white_79, white_29, weight=Int("(white_79, white_29)"))
graph.add_edge(white_75, white_25, weight=Int("(white_75, white_25)"))
graph.add_edge(white_75, empty_right, weight=Int("(white_75, empty_right)"))
graph.add_edge(white_29, white_25, weight=Int("(white_29, white_25)"))
graph.add_edge(white_29, white_39, weight=Int("white_29, white_39"))
graph.add_edge(white_39, bot_star, weight=Int("(white_39, bot_star)"))
graph.add_edge(white_25, bot_star, weight=Int("(white_25, bot_star)"))

vertices = list(graph.nodes)
weights = [weight for weight in map(lambda uvw: uvw[2], graph.edges.data("weight"))]

# all weights must be unique
solver.add(Distinct(weights))

# all weights fall in [1, 24]
for weight in weights:
    solver.add(And(weight >= 1, weight <= 24))

vertices_to_sums = [
    (white_17, 17), (white_3, 3), (white_54, 54), (white_49, 49), (white_60, 60),
    (white_79, 79), (white_75, 75), (white_29, 29), (white_39, 39), (white_25, 25)
]

# adjacent edge weight sum constraints for white vertices
for (vertex, sum) in vertices_to_sums:
    neighbour_edge_weights = [graph.adj[vertex][neighbour]["weight"] for neighbour in graph[vertex]]
    solver.add(Sum(neighbour_edge_weights) == sum)

vertices_to_path_sums = [
    (green_19_23, [19, 23]),
    (green_31, [31]),
    (green_6_9_16, [6, 9, 16]),
    (green_8, [8])
]

# existence of paths with particular weight sums for green vertices
for (vertex, path_sums) in vertices_to_path_sums:
    for path_sum in path_sums:
        constraints = []

        for destination in (set(vertices) - set([vertex])):
            for path in nx.all_simple_paths(graph, vertex, destination):
                constraints.append(nx.path_weight(graph, path, weight="weight") == path_sum)

        solver.add(Or(constraints))

# exhausts all possible solutions
while solver.check() == sat:
    model = solver.model()

    for vertex, neighbour in graph.edges():
        variable = graph.adj[vertex][neighbour]["weight"]

        if isinstance(variable, int):
            continue

        plugged_in = model.evaluate(variable).as_long()
        graph.adj[vertex][neighbour]["weight"] = plugged_in

    answer_vertices = list(nx.shortest_path(graph, top_star, bot_star, weight="weight"))
    answer_vertex_names = [vertex_names[vertex] for vertex in answer_vertices]
    print("the shortest path is: %s" % answer_vertex_names)

    answer_weights = [graph.adj[answer_vertices[index]][answer_vertices[index + 1]]["weight"] for index in range(0, len(answer_vertices) - 1)]
    chars = [c for c in map(lambda weight: chr(weight + 97 - 1), answer_weights)]
    answer = ""

    for char in chars:
        answer += char

    print("the answer is: \"%s\"" % answer)

    solver.add(Or([sym() != model[sym] for sym in model.decls()]))

print("no more solution found")