#!/usr/bin/python

import sys, random, math, time

from OpenGL.GL import *
from OpenGL.GLE import *
from OpenGL.GLU import *
from OpenGL.GLUT import *
from PIL import Image

try:
    from _frustum import Frustum
except:
    print "not using C frustum"
    from frustum import Frustum

try:
    from _octree import Octree
except:
    print "not using C octree"
    from octree import Octree

from util import normal, normalise

def average_normal(*normals):
    x, y, z = 0, 0, 0
    for nx, ny, nz in normals:
        x = x + nx
        y = y + ny
        z = z + nz
    return normalise(x, y, z)

class Ground:
    def __init__(self, heightfield_file, texture_file):
        self.treelist = 0
        self.maplist = 0
        self.alltreeslist = 0

        # create the tree display list
        #      p3
        #      /|\
        #     / | \
        #    /  |  \
        # p2 ------- p4
        #      p1
        p1 = (0, 5, -10)
        p2 = (-10, 5, 0)
        p3 = (0, 20, 0)
        p4 = (10, 5, 0)
        self.treelist = glGenLists(1)
        glNewList(self.treelist, GL_COMPILE)
        glColor(.3, .8, .3)
        glDisable(GL_TEXTURE_2D)
        glEnable(GL_LIGHTING)
        glShadeModel(GL_FLAT)
        glBegin(GL_TRIANGLE_STRIP)
        # (normals go clockwise)
        glNormal(*normal(p2, p3, p1))
        glVertex(*p2)
        glVertex(*p1)
        glVertex(*p3)
        glNormal(*normal(p4, p1, p3))
        glVertex(*p4)
        glEnd()
        glEnable(GL_TEXTURE_2D)
        glEndList()

        print "initialising ground... "
        self.image = Image.open(heightfield_file).convert('L')
        l = self.limit = self.image.size[0]
        s = self.spacing = 300
        hm= height_multiplier = 10
        self.span = s*l
        tempmap = []
        h = self.image.getpixel
        for x in range(0, l):
            z_row = []
            tempmap.append(z_row)
            for z in range(0, l):
                # figure the x and x+1 points
                xm = x*s
                if x+1 < l: x1 = x+1
                else: x1 = 0
                xm1 = (x+1) * s

                # figure the z and z+1 points
                zm = z*s
                if z+1 < l: z1 = z+1
                else: z1 = 0
                zm1 = (z+1) * s

                # figure the four corners
                #  p2 --- p3
                #    |  /|
                #    | / |
                #    |/  |
                #  p1 --- p4
                p1 = (float(xm),   float(h((x, z))*hm),  float(zm))
                p2 = (float(xm), float(h((x, z1))*hm), float(zm1))
                p3 = (float(xm1),   float(h((x1, z1))*hm), float(zm1))
                p4 = (float(xm1), float(h((x1, z))*hm), float(zm))

                # save it
                z_row.append((normal(p1, p2, p3), normal(p1, p3, p4),
                    p1, p2, p3, p4))

        # figure the point normals
        #   --- --- ---
        #  |p /|q /|r /|
        #  | / | / | / |
        #  |/ a|/ b|/ c|
        #   ---2---3---
        #  |d /|e /|f /|
        #  | / | / | / |
        #  |/ g|/ h|/ i|
        #   ---1---4---
        #  |j /|k /|o /|     NOTE: o and l are swapped so I don't have
        #  | / | / | / |           to rebind l in the loop! slack, yes :)
        #  |/ m|/ n|/ l|
        #   --- --- ---
        #
        # so, 1 = j, g, e, h, k, m
        #     2 = a, q, b, e, g, d
        #     3 = b, r, c, f, h, e
        #     4 = k, h, f, i, o, n
        # pa = x-1, z+1
        # qb = x, z+1
        # rc = x+1, z+1
        # dg = x-1, z
        # eh = x, z
        # fi = x+1, z
        # jm = x-1, z-1
        # kn = x, z-1
        # ol = x+1, z-1
        # *phew*!
        self.octree = Octree((8192., 0, 8192.), 16384.)
        self.treeoctree = Octree((8192., 0, 8192.), 16384.)
        tex_scale = 4./256
        print '... %d map squares'%(l*l)

        self.maplist = glGenLists(1)
        glNewList(self.maplist, GL_COMPILE)
        glColor(.3, .3, .3)
        glBegin(GL_TRIANGLES)
        for x in range(0, l):
            sys.stdout.write('\rworking %3d%%'%((x*100)/l))
            sys.stdout.flush()
            for z in range(0, l):
                # handle the x plus/minus
                xm = x-1
                xp = x+1
                if x == 0:
                    xm = l-1
                elif x == l-1:
                    xp = 0
                # and the z plus/minus
                zm = z-1
                zp = z+1
                if z == 0:
                    zm = l-1
                elif z == l-1:
                    zp = 0

                # get the info from the tempmap
                # pa = x-1, z+1
                a = tempmap[xm][zp][1]
                # qb = x, z+1
                q, b = tempmap[x][zp][:2]
                # rc = x+1, z+1
                r, c = tempmap[xp][zp][:2]
                # dg = x-1, z
                d, g = tempmap[xm][z][:2]
                # eh = x, z
                e, h, p1, p2, p3, p4 = tempmap[x][z]
                # fi = x+1, z
                f, i = tempmap[xp][z][:2]
                # jm = x-1, z-1
                j, m = tempmap[xm][zm][:2]
                # kn = x, z-1
                k, n = tempmap[x][zm][:2]
                # ol = x+1, z-1
                o = tempmap[xp][zm][0]

                # figure the averaged normals
                # so, 1 = j, g, e, h, k, m
                n1 = average_normal(j, g, e, h, k, m)
                #     2 = a, q, b, e, g, d
                n2 = average_normal(a, q, b, e, g, d)
                #     3 = b, r, c, f, h, e
                n3 = average_normal(b, r, c, f, h, e)
                #     4 = k, h, f, i, o, n
                n4 = average_normal(k, h, f, i, o, n)

                # figure the texture coords
                # n2/t1 --- t4/n3
                #      |   |
                #      |   |
                #      |   |
                # n1/t2 --- t3/n4
                t1 = (x*tex_scale, 1. -z*tex_scale)
                t2 = (x*tex_scale, 1. -(z*tex_scale + tex_scale))
                t3 = (x*tex_scale + tex_scale, 1. -(z*tex_scale + tex_scale))
                t4 = (x*tex_scale + tex_scale, 1. -z*tex_scale)

                # right, done
                glTexCoord2d(*t1)
                glNormal3d(*n1)
                glVertex3d(*p1)
                glTexCoord2d(*t2)
                glNormal3d(*n2)
                glVertex3d(*p2)
                glTexCoord2d(*t3)
                glNormal3d(*n3)
                glVertex3d(*p3)
                glTexCoord2d(*t1)
                glNormal3d(*n1)
                glVertex3d(*p1)
                glTexCoord2d(*t3)
                glNormal3d(*n3)
                glVertex3d(*p3)
                glTexCoord2d(*t4)
                glNormal3d(*n4)
                glVertex3d(*p4)
        glEnd()
        glEndList()

        print "\rworking 100%..."

        print "... texture"
        self.texture = Image.open(texture_file)
        ix, iy = self.texture.size
        image = self.texture.tostring("raw", "RGB", 0, -1)
        self.maptex = glGenTextures(1)
        glBindTexture(GL_TEXTURE_2D, self.maptex)
        glPixelStorei(GL_UNPACK_ALIGNMENT, 1)
        glTexImage2D(GL_TEXTURE_2D, 0, 3, ix, iy, 0, GL_RGB, GL_UNSIGNED_BYTE,
            image)
        print "... done"

    def __del__(self):
        self.clean()
    def clean(self):
        print 'cleaning up display lists'
        if self.treelist:
            glDeleteLists(self.treelist, 1)
            self.treelist = 0
        if self.alltreeslist:
            glDeleteLists(self.alltreeslist, 1)
            self.alltreeslist = 0
        if self.maplist:
            glDeleteLists(self.maplist, 1)
            self.maplist = 0

    def animate(self, world):
        pass

    def draw(self, world):
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        glTexEnvf(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_MODULATE)
        glEnable(GL_TEXTURE_2D)
        glBindTexture(GL_TEXTURE_2D, self.maptex)
        glEnable(GL_LIGHTING)
        glShadeModel(GL_SMOOTH)

        glCallList(self.maplist)

        mapsize = self.limit * self.spacing
        if world.ufo.x < mapsize/2:
            glPushMatrix()
            glTranslate(-mapsize, 0, 0)
            glCallList(self.maplist)
            glPopMatrix()
            if world.ufo.z < 2000:
                glPushMatrix()
                glTranslate(-mapsize, 0, -mapsize)
                glCallList(self.maplist)
                glPopMatrix()
                glPushMatrix()
                glTranslate(0, 0, -mapsize)
                glCallList(self.maplist)
                glPopMatrix()
            else:
                glPushMatrix()
                glTranslate(-mapsize, 0, mapsize)
                glCallList(self.maplist)
                glPopMatrix()
                glPushMatrix()
                glTranslate(0, 0, mapsize)
                glCallList(self.maplist)
                glPopMatrix()
        else:
            glPushMatrix()
            glTranslate(mapsize, 0, 0)
            glCallList(self.maplist)
            glPopMatrix()
            if world.ufo.z < 2000:
                glPushMatrix()
                glTranslate(mapsize, 0, -mapsize)
                glCallList(self.maplist)
                glPopMatrix()
                glPushMatrix()
                glTranslate(0, 0, -mapsize)
                glCallList(self.maplist)
                glPopMatrix()
            else:
                glPushMatrix()
                glTranslate(mapsize, 0, mapsize)
                glCallList(self.maplist)
                glPopMatrix()
                glPushMatrix()
                glTranslate(0, 0, mapsize)
                glCallList(self.maplist)
                glPopMatrix()

        glDisable(GL_TEXTURE_2D)


class Exhaust:
    def __init__(self, point_file, num_points=100):
        self.num_points = num_points
        self.points = [(0, [], [], [0, 0, 0, 0])]*self.num_points
        self.index = 0

        # texture
        self.texture = Image.open(point_file)
        ix, iy = self.texture.size
        image = self.texture.tostring("raw", "RGBA", 0, -1)
        self.pointtex = glGenTextures(1)
        glBindTexture(GL_TEXTURE_2D, self.pointtex)
        glPixelStorei(GL_UNPACK_ALIGNMENT,1)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, ix, iy, 0, GL_RGBA,
            GL_UNSIGNED_BYTE, image)

    def addPoint(self, ufo):
        # reverse direction to thrust
        f = -3 - (random.random()/3 + 0.3)
        yaw = ufo.r_yaw + (random.random() - 0.5)/2
        pitch = ufo.r_pitch + (random.random() - 0.5)/2
        dx = f * math.sin(pitch) * math.sin(yaw) + ufo.dx
        dy = f * math.cos(pitch) + ufo.dy
        dz = f * math.sin(pitch) * math.cos(yaw) + ufo.dz

        # point is position, velocity, colour
        self.points[self.index] = (.05, [ufo.x, ufo.y, ufo.z],
            [dx, dy, dz], [1, random.random(), 0, 1])
        self.index = (self.index + 1)%self.num_points

    def animate(self, world):
        for fade, position, velocity, colour in self.points:
            if colour[3] <= 0: continue
            position[0] = position[0] + velocity[0]
            velocity[1] = velocity[1] - world.gravity 
            position[1] = position[1] + velocity[1]
            position[2] = position[2] + velocity[2]
            colour[3] = colour[3] - fade

    def draw(self, world):
        # set up the texture parameters
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        glTexEnvf(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_MODULATE)
        glBlendFunc(GL_SRC_ALPHA, GL_ONE);
        glEnable(GL_BLEND)

        # set up the polygon drawing
        glEnable(GL_TEXTURE_2D)
        glBindTexture(GL_TEXTURE_2D, self.pointtex)
        glDisable(GL_LIGHTING)
        glColor(1, 1, 1)

        # draw the quads for the points
        glBegin(GL_QUADS)
        for life, position, velocity, colour in self.points:
            if colour[3] <= 0: continue
            glColor4f(*colour)
            x, y, z = position
            glTexCoord2f(0, 0)
            glVertex3f(x, y, z)
            glTexCoord2f(1, 0)
            glVertex3f(x+4, y, z)
            glTexCoord2f(1, 1)
            glVertex3f(x+4, y+4, z)
            glTexCoord2f(0, 1)
            glVertex3f(x, y+4, z)
        glEnd()

        # re-enable everything
        glEnable(GL_LIGHTING)
        glDisable(GL_BLEND)
        glDisable(GL_TEXTURE_2D)

class UFO:
    def __init__(self):
        self.yaw = self.pitch = self.r_yaw = self.r_pitch = 0
        self.x, self.y, self.z = 0, 700, 0

        self.dx = self.dy = self.dz = 0
        self.ddx = self.ddy = self.ddz = 0
        self.exhaust = Exhaust('data/point_64.tga')

    def rotateAngles(self, x, y, pi2 = 2*math.pi):
        self.r_yaw = math.atan2(x, y) + math.pi
        self.pitch = -math.sqrt(x**2 + y**2)
        if self.pitch < -180: self.pitch = -180
        self.r_pitch = self.pitch * math.pi / 180
        self.yaw = self.r_yaw * 180 / math.pi

    def thrust(self, boosters_are_go):
        if boosters_are_go:
            self.exhaust.addPoint(self)
            f = 0.30
            self.ddx = f * math.sin(self.r_pitch) * math.sin(self.r_yaw)
            self.ddy = f * math.cos(self.r_pitch)
            self.ddz = f * math.sin(self.r_pitch) * math.cos(self.r_yaw)
        else:
            self.ddx = self.ddy = self.ddz = 0

    def animate(self, world):
        # TODO: introduce drag :)
        self.x = self.x + self.dx
        gs = world.ground.span
        if self.x >= gs: self.x = self.x - gs
        elif self.x < 0: self.x = self.x + gs
        self.dx = self.dx + self.ddx
        self.y = self.y + self.dy
        self.dy = self.dy + self.ddy - world.gravity
        self.z = self.z + self.dz
        if self.z >= gs: self.z = self.z - gs
        elif self.z < 0: self.z = self.z + gs
        self.dz = self.dz + self.ddz
        self.exhaust.animate(world)

    def draw(self, world):
        glShadeModel(GL_FLAT)
        glPushMatrix()
        glTranslate(self.x, self.y, self.z)
        glRotate(self.yaw, 0, 1, 0)
        glRotate(self.pitch, 1, 0, 0)

        glColor3f(0.6, 0.8, 0.3)
        # draw the UFO body
        glBegin(GL_TRIANGLE_STRIP)
        # normal((0, 5, 5), (-10, 0, 10), (10, 0, 10))
        glNormal(0.0, 0.70710678118654746, 0.70710678118654746)
        glVertex3f(-10, 0, 10)
        glVertex3f(10, 0, 10)
        glVertex3f(0, 5, 5)

        # normal((0, 5, 5), (10, 0, 10), (0, 0, -10))
        glNormal(0.53452248382484879, 0.80178372573727319, -0.2672612419124244)
        glVertex3f(0, 0, -10)

        # normal((0, 5, 5), (0, 0, -10), (-10, 0, 10))
        glNormal(-0.53452248382484879, 0.80178372573727319, -0.2672612419124244)
        glVertex3f(-10, 0, 10)

        # normal((0, 0, -10), (10, 0, 10), (-10, 0, 10))
        glNormal(0, -1, 0)
        glVertex3f(10, 0, 10)
        glEnd()

        glPopMatrix()

        # draw the exhaust
        self.exhaust.draw(world)


class World:
    def __init__(self, args):
        # initial mouse position
        self.button0_down = 0
        self.width = 640
        self.height = 480

        # pause flag
        self.pause = 0

        # world parameters
        self.gravity = 0

        # set the display mode and create the window
        glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH)
        glutInitWindowSize(self.width, self.height)
        glutCreateWindow("UFO!")

        # display all the time
        glutDisplayFunc(self.on_display)
        glutIdleFunc(self.on_display)

        # go with both normal and passive (mouse button down and not)
        glutPassiveMotionFunc(self.on_motion)
        glutMotionFunc(self.on_motion)
        # and detect mouse clicks
        glutMouseFunc(self.on_mouse)
        # and key clicks
        glutKeyboardFunc(self.on_key)
        glutReshapeFunc(self.on_reshape)
        if '-fs' in args:
            glutFullScreen()

        # set up the display specifics
        glClearDepth(1.0)
        glEnable(GL_DEPTH_TEST)
        glClearColor(0.0, 0.0, 0.0, 0.0)
        glMatrixMode(GL_MODELVIEW)

        # initialize lighting
        glLightfv(GL_LIGHT0, GL_POSITION,  (-40.0, 200, -100.0, 0.0))
        glLightfv(GL_LIGHT0, GL_DIFFUSE, (0.99, 0.99, 0.99, 1.0))
        glEnable(GL_LIGHT0)
        glLightfv(GL_LIGHT1, GL_POSITION, (-100.0, 400, -40.0, 0.0))
        glLightfv(GL_LIGHT1, GL_DIFFUSE, (0.99, 0.99, 0.99, 1.0))
        glEnable(GL_LIGHT1)
        glEnable(GL_LIGHTING)
        glEnable(GL_COLOR_MATERIAL)

        self.ufo = UFO()
        self.ground = Ground('data/heightfield_64.png',
            'data/texmap_2_256.png')
        self.angle = 0
        self.dump = 0

        self.frames = 0

        self.ts = None
        self.frustum = Frustum()

    def on_key(self, key, x, y):
        if key == 'g':
            if self.gravity:
                self.gravity = 0
            else:
                self.gravity = 0.05

        elif key == 'p':
            self.pause = not self.pause

        elif key == 'd':
            print '***** DUMP %s *****'%self.ts
            print (self.ufo.x, self.ufo.y, self.ufo.z)
#            print (int(self.ufo.x/self.ground.spacing), int(self.ufo.z/self.ground.spacing))
            self.dump = 1

        elif key == 't':
            print 1/((time.clock()-self.ts)/self.frames)

        else:
            print "timing:", 1/((time.clock()-self.ts)/self.frames)
            self.ground.clean()
            sys.exit()

    def on_mouse(self, button, state, x, y):
        if button == GLUT_RIGHT_BUTTON:
            self.button0_down = not state

    def on_motion(self, x, y):
        if self.pause: return
        x = self.width/2 - x
        y = self.height/2 - y
        self.ufo.rotateAngles(x, y)

    def on_reshape(self, width, height):
        self.width = width
        self.height = height

    def on_display(self):
        try:
            self._on_display()
        except:
            import traceback
            traceback.print_exc()
            sys.exit()

    def _on_display(self):          
        self.frames = self.frames + 1
        if not self.ts:
            self.ts = time.clock()
        # do world calculations
        if not self.pause:
            self.ufo.animate(self)
            self.ufo.thrust(self.button0_down)

        # setup the viewport
        glViewport(0, 0, self.width, self.height)

        # setup the projection matrix
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()

        # setup the projection
        # calculate left/right and top/bottom clipping planes based the
        # smallest square viewport
        a = 9.0/min(self.width, self.height)
        lr, tb = a*self.width, a*self.height
        glFrustum(-lr, lr, -tb, tb, 50.0, 20000.0)

#        aspect = self.width*1./self.height
#        gluPerspective(30, aspect, 50, 6000)

        # and place the camera
        x, y, z = self.ufo.x, self.ufo.y, self.ufo.z
        gluLookAt(x, y+400, z-600, x, y, z, 0, 1, 0)

        # clear the buffer
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

        self.frustum.calculate()

        # Set up the model view matrix
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()

#        glPushMatrix()
#        glTranslate(self.ufo.x, self.ufo.y, self.ufo.z)
#        glEnable(GL_FOG)
#        glFog(GL_FOG_MODE, GL_LINEAR)
#        glFog(GL_FOG_COLOR, (0, 0, 0, 1))
#        glFog(GL_FOG_START, 1)
#        glFog(GL_FOG_END, 3000)
#        glFog(GL_FOG_HINT, GL_NICEST)
#        glPopMatrix()

        self.ground.draw(self)
        self.ufo.draw(self)

        self.dump = 0

        # swap the buffer
        glutSwapBuffers()

if __name__ == '__main__':
    # initialize GLUT 
    glutInit(sys.argv)
    # create the demo window
    World(sys.argv)
    # enter the event loop
    glutMainLoop ()

