"""This module contains the SmallWorldGraph class for 4.4, exercise 4.

"""
import random

from Graph import Edge
from RandomGraph import RandomGraph


class SmallWorldGraph(RandomGraph):
    """A small world graph for 4.4, exercise 4 in Think Complexity."""

    def __init__(self, vs, k, p):
        """Create a small world graph. Parameters:
            vs - a list of vertices
            k - regular graph parameter k: the number of edges between vertices
            p - probability for rewiring.

        First a regular graph is created from the list of vertices and parameter
        k, which specifies how many edges each vertex should have. Then the
        rewire() method is called with parameter p to rewire the graph according
        to the algorithm by Watts and Strogatz.

        """
        super(SmallWorldGraph, self).__init__(vs=vs)
        self.add_regular_edges(k)
        self.rewire(p)
        self.set_edge_length(1)

    def rewire(self, p):
        """Assuming the graph is initially a regular graph, rewires the graph
        using the Watts and Strogatz algorithm.

        """
        # We iterate over the edges in a random order:
        edges = self.edges()
        random.shuffle(edges)
        vertices = self.vertices()

        for e in edges:
            if random.random() < p:
                v, w = e    # remember the 2 vertices this edge connected
                self.remove_edge(e)     # remove from graph

                # pick another vertex to connect to v; duplicate edges are
                # forbidden
                while True:
                    x = random.choice(vertices)
                    if v is not x and not self.get_edge(v, x):
                        self.add_edge(Edge(v, x))
                        break