############################################################
# graph and display helpers
############################################################


def get_neighbors(graph, node):
    return graph.get(node, [])


def path_to_string(path):
    """path is a list of nodes"""
    if path is None:
        return "None"
    node_strs = []
    for node in path:
        node_strs.append(str(node))
    return " --> ".join(node_strs)


def pathlist_to_string(queue):
    """path is a list of nodes"""
    path_strs = []
    for path in queue:
        path_strs.append(path_to_string(path))
    return "[" + ", ".join(path_strs) + "]"


############################################################
# checking paths in trees using depth-first search
############################################################


tiny_tree = {
    "home": ["school", "personal", "downloads"],
    "school": ["fall24", "spring25", "fall25"],
    "fall25": ["6.100", "8.01"],
    "personal": ["photos", "bills"],
    "downloads": ["slides.pdf", "ps1", "beach.jpg"],
    "ps1": ["pset.py", "test.py"],
}


def exists_in(current_node, goal):
    if current_node == goal:
        return True
    for next_node in get_neighbors(tiny_tree, current_node):
        result = exists_in(next_node,goal)
        if result:
            return result
    return False


# print(exists_in("home", "pset.py"))


############################################################
# finding paths in trees using breadth-first search
############################################################


def bfs_tree(graph, start, goal):
    current_frontier = [[start]]
    next_frontier = []

    while len(current_frontier) > 0:
        print("Current frontier:", pathlist_to_string(current_frontier))

        for path in current_frontier:
            print("  Current BFS path:", path_to_string(path))

            current_node = path[-1]
            if current_node == goal:
                return path

            for next_node in get_neighbors(graph, current_node):
                next_frontier.append(path + [next_node])

        current_frontier, next_frontier = next_frontier, []

    return None


def demo_bfs_tree():
    filesystem = {
        "home": ["school", "personal", "downloads"],
        "school": ["fall24", "spring25", "fall25"],
        "fall25": ["6.100", "8.01"],
        "personal": ["photos", "bills"],
        "downloads": ["slides.pdf", "ps1", "beach.jpg"],
        "ps1": ["pset.py", "test.py"],
    }

    print(bfs_tree(filesystem, "home", "pset.py"))
    # print()
    # print(bfs_tree(filesystem, "home", "beach.jpg"))
    # print()

    # # make "downloads" first child of "home"
    # folders = filesystem["home"]
    # filesystem["home"] = [folders[-1]] + folders[:-1]

    # # doesn't change number of frontiers explored,
    # # but explore less of last frontier
    # print(bfs_tree(filesystem, "home", "pset.py"))
    # print()
    # print(bfs_tree(filesystem, "home", "beach.jpg"))
    # print()


# demo_bfs_tree()


############################################################
# finding paths in general graphs using breadth-first search
############################################################


def bfs_graph(graph, start, goal):
    current_frontier = [[start]]
    next_frontier = []
    visited = {start}

    while len(current_frontier) > 0:
        print("Current frontier:", pathlist_to_string(current_frontier))

        for path in current_frontier:
            print("  Current BFS path:", path_to_string(path))

            current_node = path[-1]
            if current_node == goal:
                return path

            for next_node in get_neighbors(graph, current_node):

                # avoid nodes already seen in current or previous frontier
                if next_node in visited:
                    print(f"    AVOID revisiting {next_node}")
                    continue
                visited.add(next_node)
                next_frontier.append(path + [next_node])

        current_frontier, next_frontier = next_frontier, []

    return None


def demo_bfs_graph():
    flights = {
        "Boston": ["Providence", "New York"],
        "Providence": ["Boston", "New York"],
        "New York": ["Chicago"],
        "Chicago": ["Denver", "Phoenix"],
        "Denver": ["New York", "Phoenix"],
        "Los Angeles": ["Boston"],
    }

    print(bfs_graph(flights, "Boston", "Phoenix"))
    print()

    # change the ordering in which we explore Chicago's neighbors
    flights["Chicago"].reverse()

    # shortest path solution does not change
    print(bfs_graph(flights, "Boston", "Phoenix"))
    print()


# demo_bfs_graph()


############################################################
# weighted graph extension
############################################################


def add_edge(graph, u, v):
    if u not in graph:
        graph[u] = []
    graph[u].append(v)


def unroll_weighted_graph(graph):
    new_graph = {}
    for u, edges in graph.items():
        for v, weight in edges:
            if weight == 1:
                add_edge(new_graph, u, v)
            else:
                # create intermediate nodes
                prev = u
                for i in range(1, weight):
                    mid = f"{u}->{v}#{i}"
                    add_edge(new_graph, prev, mid)
                    prev = mid
                add_edge(new_graph, prev, v)
    return new_graph


def demo_bfs_weighted_graph():
    flights = {
        "Boston": [("Providence", 1), ("New York", 2)],
        "Providence": [("Boston", 1), ("New York", 2)],
        "New York": [("Chicago", 3)],
        "Chicago": [("Denver", 3), ("Phoenix", 5)],
        "Denver": [("New York", 4), ("Phoenix", 1)],
        "Los Angeles": [("Boston", 8)],
    }
    flights_expanded = unroll_weighted_graph(flights)

    result = bfs_graph(flights_expanded, "Boston", "Phoenix")
    filtered_path = []
    for node in result:
        if "#" not in node:
            filtered_path.append(node)
    print(path_to_string(filtered_path))


# demo_bfs_weighted_graph()