import networkx as nx

from semistaticsim.rendering.floor import ReachabilityGraph
from semistaticsim.rendering.utils.dict2xyztup import dict2xyztuple


def cleanup_request(requested_node):
    if isinstance(requested_node, dict) and all([x in requested_node for x in ['x', 'y', 'z']]):
        requested_node = dict2xyztuple(requested_node)
    if isinstance(requested_node, tuple):
        requested_node = list(requested_node)
    return requested_node

def astar(requested_start_node, requested_end_node, filtered_rg: ReachabilityGraph, controller):
    requested_start_node = cleanup_request(requested_start_node)
    requested_end_node = cleanup_request(requested_end_node)

    # filtered_rg = full_reachability_graph.filter_(controller.last_event.metadata["objects"], mindist_to_obj=0.5)

    # filtered_rg.parent_floor.plot_self(reachability_graph=filtered_rg)
    # filtered_rg.parent_floor.plot_self(reachability_graph=filtered_rg)
    # filtered_rg.parent_floor.plot_self(samples=jnp.array(path))

    start_node = filtered_rg.shunt(requested_start_node)
    end_node = filtered_rg.shunt(requested_end_node)

    assert start_node in filtered_rg.full_graph
    assert end_node in filtered_rg.full_graph

    path = nx.astar_path(filtered_rg.full_graph, start_node, end_node)[1:]


    for node in path:
        yield node
