"""Chapter 5.5, exercise 6 in Allen Downey's Think Complexity book.

2. Use the WS model to generate the largest graph you can in a reasonable amount
   of time. Plot P(k) versus k and see if you can characterize the tail
   behavior.

"""
import sys

from matplotlib import pyplot

from Graph import Vertex
from SmallWorldGraph import SmallWorldGraph



def main(script, n, k, p):

    # create a SmallWorldGraph with n vertices, k regular edges between
    # vertices, and rewiring probability p.

    n, k, p = int(n), int(k), float(p)
    vs = [Vertex() for i in xrange(n)]

    g = SmallWorldGraph(vs, k, p)

    # retrieve probabilities
    p = g.get_p()

    # plot P(k) versus k on a log-log scale

    vals = p.items()
    vals.sort(key=lambda t: t[0])
    x, y = zip(*vals)

    assert abs(sum(y) - 1.0) < 1e-6

    pyplot.clf()
    pyplot.xscale('log')
    pyplot.yscale('log')
    pyplot.title('P(k) versus k')
    pyplot.xlabel('k')
    pyplot.ylabel('P(k)')
    pyplot.plot(x, y, label='P(k) vs. k', color='green', linewidth=3)
    pyplot.legend(loc='upper right')
    pyplot.show()


if __name__ == '__main__':
    main(*sys.argv)