class Vertex(object):
    """Describes a vertex object in terms of a "key" and a
    dictionary that indicates edges to neighboring vertices with
    a specified weight.
    """
    def __init__(self, key):
        """Constructs a vertex with a key value and an empty dictionary 
        in which we'll store other vertices to which this vertex is connected.
        """
        self.key = key
        self.neighbors = {}   # empty dictionary for neighboring vertices. key is
                              # the Vertex, value is the weight
        self.color = 'white'
        self.distance = 0
        self.previous = None

    def set_neighbor(self, other, weight=0):
        """Adds a reference to a neighboring Vertex object `other` to the dictionary, 
        to which this vertex is connected by an edge. If a weight is not indicated, 
        default weight is 0.
        """
        self.neighbors[other] = weight

    def get_color(self):
        return self.color
    
    def set_color(self, color):
        self.color = color

    def get_distance(self):
        return self.distance
    
    def set_distance(self, new_distance):
        self.distance = new_distance

    def get_previous(self):
        return self.previous
    
    def set_previous(self, new_prev):
        self.previous = new_prev

    def __repr__(self):
        """Returns a representation of the vertex and its neighbors, suitable for 
        printing. Check out the example of 'list comprehension' here!
        """
        return f"Vertex({self.key})"
        
    def __str__(self):
        return ( f"{self.key} (color={self.color}), connected to: "
        + f"{[x.key for x in self.neighbors]}")

    def get_neighbors(self):
        return self.neighbors.keys()    # returns Vertex objects

    def get_key(self):
        return self.key

    def get_neighbor(self, other):
        """Returns the weight of an edge connecting this vertex with another,
        or None if the neighbor doesn't exist
        """
        return self.neighbors.get(other, None)
    
class Graph(object):
    """Describes the Graph class, which is primarily a dictionary
    mapping vertex names to Vertex objects, along with a few methods
    that can be used to manipulate them.
    """
    def __init__(self):
        self.vertices = {}

    def set_vertex(self, key):
        self.vertices[key] = Vertex(key)

    def get_vertex(self, key):
        """Looks for the key in the dictionary of Vertex objects, and
        returns the Vertex if found. Otherwise, returns None.
        """
        '''
        # This is the classic way of doing that
        if key in self.graph.keys():
            return self.graph[key]
        else:
            return None
        '''
        # Single-line alternative
        return self.vertices.get(key, None)

    def __contains__(self, key):
        """This 'dunder' expression is written so we can use Python's "in"
        operation: If the parameter 'key' is in the dictionary of vertices,
        the value of "key in my_graph" will be True, otherwise False.
        """
        return key in self.vertices

    def add_edge(self, from_vert, to_vert, weight=0):
        """Adds an edge connecting two vertices (specified by key
        parameters) by modifying those vertex objects. Note that
        the weight can be specified as well, but if one isn't
        specified, the value of weight will be the default value
        of 0.
        """
        # if the from_key doesn't yet have a vertex, create it
        if from_vert not in self.vertices:
            self.set_vertex(from_vert)
        # if the to_key doesn't yet have a vertex, create it
        if to_vert not in self.vertices:
            self.set_vertex(to_vert)
        # now we can create the edge between the two
        self.vertices[from_vert].set_neighbor(self.vertices[to_vert], weight)

    def get_vertices(self):
        return self.vertices.keys()

    def __iter__(self):
        """Another 'dunder' expression that allows us to iterate through
        the list of vertices.
        Example use:
        for vertex in graph:  # Python understands this now!
            print(vertex)
        """
        return iter(self.vertices.values())