terrain/terrain.py

import random
import numpy
import pyglet
from pyglet.gl import *
from trackball_camera import TrackballCamera


class Terrain():
    def __init__(self, details):
        self.size = 2 ** details + 1
        self.map = numpy.zeros((self.size, self.size))

    def _square(self, x, y, size, offset):
        average = sum([
            self.map[x - size][y - size],
            self.map[x + size][y - size],
            self.map[x + size][y + size],
            self.map[x - size][y + size]]) / 4
        self.map[x][y] = average + offset

    def _diamond(self, x, y, size, offset):
        average = sum([
            self.map[x][y - size],
            self.map[x + size][y],
            self.map[x][y + size],
            self.map[x - size][y]]) / 4
        self.map[x][y] = average + offset

    def generate(self, high):
        i, half = self.size - 1, (self.size - 1) / 2
        scale = high * self.size

        while i > 1:
            for y in numpy.arange(half, self.size - 1, i):
                for x in numpy.arange(half, self.size - 1, i):
                    self._square(x, y, half, scale / 2 * random.random())
            for y in numpy.arange(0, self.size - 1, half):
                for x in numpy.arange((y + half) % i, self.size - 1, i):
                    self._diamond(x, y, half, scale / 2 * random.random())
            i /= 2
            half = i / 2
            scale = high * half * 2


class Window(pyglet.window.Window):
    def __init__(self, *args, **kwargs):
        self.terrain = None
        self.camera = kwargs['camera']
        del kwargs['camera']
        super().__init__(*args, **kwargs)

        # Initialize OpenGL
        glEnable(GL_DEPTH_TEST)
        glDisable(GL_CULL_FACE)

    def on_mouse_drag(self, x, y, dx, dy, buttons, modifiers):
        """Move camera/zoom on mouse drag"""
        if buttons & pyglet.window.mouse.LEFT:
            self.camera.mouse_roll(
                self._norm(x, self.width),
                self._norm(y, self.height))
        elif buttons & pyglet.window.mouse.RIGHT:
            self.camera.mouse_zoom(
                self._norm(x * 2, self.width),
                self._norm(y * 2, self.height))

    def on_mouse_press(self, x, y, button, modifiers):
        """Move camera/zoom on mouse drag"""
        if button == pyglet.window.mouse.LEFT:
            self.camera.mouse_roll(
                self._norm(x, self.width),
                self._norm(y, self.height),
                False)
        elif button == pyglet.window.mouse.RIGHT:
            self.camera.mouse_zoom(
                self._norm(x * 2, self.width),
                self._norm(y * 2, self.height),
                False)

    def on_resize(self, width, height):
        """Adjust drawing after window is resized"""
        self.width = width
        self.height = height
        glViewport(0, 0, self.width, self.height)
        self.on_show()

    def on_show(self):
        """Set OpenGl config"""
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective(40, self.width / self.height, 1, 400)
        self.camera.update_modelview()

    def on_draw(self):
        """Draw the current frame"""
        if self.terrain is None:
            return
        self.clear()
        map, size = self.terrain.map, self.terrain.size

        glPushMatrix()
        glTranslatef(-size / 2, 0, -size / 2)

        # Draw terrain
        for x in numpy.arange(size - 1):
            for y in numpy.arange(size - 1):
                glBegin(GL_TRIANGLE_STRIP)
                glColor3f(map[x][y] / 20, map[x][y] / 20, map[x][y] / 20)
                glVertex3f(x, map[x][y], y)
                glVertex3f(x + 1, map[x + 1][y], y)
                glVertex3f(x, map[x][y + 1], y + 1)
                glVertex3f(x + 1, map[x + 1][y + 1], y + 1)
                glEnd()
        glPopMatrix()

        # Draw axis
        glBegin(GL_LINES)
        glColor3f(1, 0, 0)
        glVertex3f(-size, 0, 0)
        glVertex3f(size, 0, 0)
        glColor3f(0, 1, 0)
        glVertex3f(0, -size, 0)
        glVertex3f(0, size, 0)
        glColor3f(0, 0, 1)
        glVertex3f(0, 0, -size)
        glVertex3f(0, 0, size)
        glEnd()

    def _norm(self, x, max_x):
        """given x within [0,max_x], scale to a range [-1,1]"""
        return (2 * x - float(max_x)) / float(max_x)

    def draw(self, terrain):
        """Render the height map"""
        self.terrain = terrain


def main():
    terrain = Terrain(5)
    camera = TrackballCamera(150)
    window = Window(caption="Terrain", resizable=True,
            width=800, height=600, camera=camera)
    terrain.generate(0.7)
    window.draw(terrain)
    pyglet.app.run()

if __name__ == '__main__':
    main()