import pygame, sys, random

from pygame.locals import *

def draw_screen(height_map, square):

    render_map(height_map, square)

    pygame.display.update()

def convert(grid):

    # initialize an empty MAP_SIZE x MAP_SIZE array
    out = [[0 for y in xrange(MAP_SIZE)] for x in xrange(MAP_SIZE)]

    largest = 0
    smallest = MAP_SIZE * 3.0

    for x in xrange(MAP_SIZE):

        for y in xrange(MAP_SIZE):

            # average surrounding points to get square's value
            val = (grid[x][y] + grid[x+1][y] + grid[x][y+1] + grid[x+1][y+1]) \
                / 4.0

            out[x][y] = val

            # find largest and smallest while we're at it
            if val > largest:

                largest = val

            if val < smallest:

                smallest = val

    smallest = float(smallest)

    #print largest
    #print smallest

    # if any point below 0, bump everything up by that amount
    if smallest < 0:

        bump = 0 - smallest

    else:

        bump = 0

    #print bump

    largest = float(largest + bump)

    # compress to 8-bits or whatever
    for x in xrange(MAP_SIZE):

        for y in xrange(MAP_SIZE):

            out[x][y] = int(((out[x][y] + bump) / largest) * 255)

    return out

def render_map(height_map, square):

    for x in xrange(MAP_SIZE):

        for y in xrange(MAP_SIZE):

            val = height_map[x][y]

            color = (val, val, val)

            pygame.draw.rect(DISPLAYSURF, color, (x * square, y * square,
                x * square + square, y * square + square))

def generate(size):

    # the size of the first square, or, the size of the height map itself
    extent = size

    # half the size of the heightmap
    # coincidentally, the middle index number as well
    # ex: size = 4, [0, 1, 2, 3, 4], 4 / 2 = 2 = middle index
    half = size / 2

    # initialize an empty grid with dimensions size + 1
    # example: a 16x16 heightmap calls for a 17x17 grid
    points = [[0 for y in xrange(size + 1)] for x in xrange(size + 1)]

    # "noise" happens here
    scale = half * 1.5

    # randomize corners
    points[0][0] = half + random.uniform(-scale, scale)
    points[size][0] = half + random.uniform(-scale, scale)
    points[0][size] = half + random.uniform(-scale, scale)
    points[size][size] = half + random.uniform(-scale, scale)

    # when size is 1, there are no more points to fill in
    # I don't know why some people's algorithms say "while size > 0",
    # since there is nothing more to do to a square like (0, 0), (1, 0),
    # (0, 1), (1, 1)
    while size > 1:

        #print "half:", half
        #print "scale:", scale

        # square step

        # here we step through both dimensions, with a stride of the current
        # square size
        #
        # example: a 4 x 4 heightmap
        # our range on the first iteration is (0, 4, 4)
        # first square found lies at (0, 0)
        # then we take 1 stride of length 4, which puts us at (4, 4),
        # which is out of range (ours is 0-4 exclusive)
        # so only 1 square is found
        #
        # next iteration size is 2, so we work on 4 squares, located at:
        # (0, 0), (0, 2), (2, 0), (2, 2)
        # and so on
        for x in xrange(0, extent, size):

            for y in xrange(0, extent, size):

                # the 4 corners of this square are averaged
                sq_avg = (points[x][y] + points[x+size][y] + \
                    points[x][y+size] + points[x+size][y+size]) / 4.0

                # set center point to avg + "noise"
                points[x+half][y+half] = sq_avg + random.uniform(-scale, scale)

        # diamond step

        # for each square, we also need to set the values of the "diamond"
        # points, that is, the ones to the left, right, above, and below our
        # newly created square center points
        #
        # to avoid errors, we only use 3 points for averaging points along the
        # edges, since the 4th is off the map/non-existent
        for y in xrange(0, extent, size):

            for x in xrange(0, extent, size):

                # the top point of the diamond
                if y > 0:

                    # 4 points to average,
                    # the upper left point, upper right, middle, and
                    # the middle of the square above it
                    avg = (points[x][y] + points[x+size][y] + \
                        points[x+half][y+half] + points[x+half][y-half]) / 4.0

                else:

                    # otherwise, we use all but the middle of the above
                    avg = (points[x][y] + points[x+size][y] + \
                        points[x+half][y+half]) / 3.0

                # set top point to that average + noise
                points[x+half][y] = avg + random.uniform(-scale, scale)

                # the bottom point of the diamond
                if y < extent - size:

                    # 4 points to average
                    avg = (points[x+size][y+size] + points[x+half][y+half] + \
                        points[x+size][y+size] + points[x+half][y+size+half]) \
                        / 4.0

                else:

                    # only 3 points
                    avg = (points[x+size][y+size] + points[x+half][y+half] + \
                        points[x+size][y+size]) / 3.0

                points[x+half][y+size] = avg + random.uniform(-scale, scale)

                # the left point of the diamond
                if x > 0:

                    # 4 points to average
                    avg = (points[x][y] + points[x+half][y+half] + \
                        points[x][y+size] + points[x-half][y+half]) / 4.0

                else:

                    # only 3 points
                    avg = (points[x][y] + points[x+half][y+half] + \
                        points[x][y+size]) / 3.0

                points[x][y+half] = avg + random.uniform(-scale, scale)

                # the right point of the diamond
                if x < extent - size:

                    # 4 points to average
                    avg = (points[x+size][y+size] + points[x+half][y+half] + \
                        points[x+size][y+size] + points[x+size+half][y+half]) \
                        / 4.0

                else:

                    # only 3 points
                    avg = (points[x+size][y+size] + points[x+half][y+half] + \
                        points[x+size][y+size]) / 3.0

                points[x+size][y+half] = avg + random.uniform(-scale, scale)

        # end of this iteration
        # halve size of squares and size of half squares
        size /= 2
        half /= 2

        # the scale, or "noise", is also reduced for each iteration
        scale = half * 1.5

    return points

#################################### setup #####################################

pygame.init()

# size of visible window
WIN = 900

# size of height map to be generated, needs to be a power of 2
MAP_SIZE = 256

# size of squares for display
SQUARE = WIN / float(MAP_SIZE)

# initialize surface and window
DISPLAYSURF = pygame.display.set_mode((WIN, WIN))
pygame.display.set_caption('Diamond-Square Terrain')

def main():

    # generate a grid based on the diamond-square algo
    grid = generate(MAP_SIZE)

    # convert that to a height map of size MAP_SIZE, with values 0-255
    height_map = convert(grid)

    draw_screen(height_map, SQUARE)

    while True:

        for event in pygame.event.get():

            if event.type == QUIT:
                pygame.quit()
                sys.exit()

        #draw_screen(height_map, SQUARE)

if __name__ == '__main__':

    main()