import pylab as P

def split_line(pt1, pt2):
    """Given two complex numbers p1 and p2 return a list of the
    necessary points pairwise.
    """
    # The points we generate are p1, p2, p3, p4 and p5
    p1, p5 =  pt1, pt2
    diff = pt2 - pt1
    segment = diff/3.0
    p2 = p1 + segment
    p4 = p5 - segment
    # Now rotate a line given by 60degrees.
    # Recall that complex multiplication does this easily.
    # Complex numbers are defined by complex(a, b) also.
    p3 = p2 + segment*complex(P.cos(P.pi/3), P.sin(P.pi/3))
    return [(p1, p2), (p2, p3), (p3, p4), (p4, p5)]

def koch(n):
    start = complex(0, 0)
    end = complex(1, 0)
    points = [(start, end)]
    for level in range(n):
        new_points = []
        for (p1, p2) in points:
            new_points.extend(split_line(p1, p2))
        points = new_points
    points.insert(0, (start, points[0][0]) )
    points.append((points[-1][1], end))
    return points

def plot_koch(n):
    points = koch(n)
    x = [t[0].real for t in points]
    y = [t[0].imag for t in points]

    P.plot(x, y, 'r-')
    P.ylim(0, 1)
    P.xlim(0, 1)