__author__ = 'Martin Felder, felder@in.tum.de'

from OpenGL.GL import * #@UnusedWildImport
from OpenGL.GLU import * #@UnusedWildImport
from OpenGL.GLUT import * #@UnusedWildImport

from math import acos, pi, sqrt
from tools.mathhelpers import crossproduct, norm, dotproduct

import time
import Image #@UnresolvedImport

from pybrain.tools.networking.udpconnection import UDPClient


class ODEViewer(object):
    def __init__(self, servIP="127.0.0.1", ownIP="127.0.0.1", port="21590", buf="16384"):

        # initialize the viewport size
        self.width = 800
        self.height = 600

        # initialize object which the camera follows
        self.centerObj = None
        self.mouseView = True
        self.viewDistance = 30
        self.lastx = -0.5
        self.lasty = 1
        self.lastz = -1

        self.dt = 1
        self.fps = 50
        self.lasttime = time.time()
        self.starttime = time.time()
        self.captureScreen = False
        self.isCapturing = False
        self.isFloorGreen = True

        self.message = None
        self.keyboardCallback = None

        # capture only every frameT. frame
        self.counter = 0
        self.frameT = 1

        self.init_GL()

        # set own callback functions
        glutMotionFunc (self._motionfunc)
        glutPassiveMotionFunc(self._passivemotionfunc)
        glutDisplayFunc (self._drawfunc)
        glutIdleFunc (self._idlefunc)
        glutKeyboardFunc (self._keyfunc)


        self.dt = 1.0 / self.fps
        self.lasttime = time.time()
        self.starttime = self.lasttime

        # initialize udp client
        self.client = UDPClient(servIP, ownIP, port, buf)


    def start(self):
        # start the OpenGL main loop
        while True:
            glutMainLoop()

    def setFrameRate(self, fps):
        self.fps = fps
        self.dt = 1.0 / self.fps

    def setCaptureScreen(self, capture):
        self.captureScreen = capture

    def getCaptureScreen(self):
        return self.captureScreen

    def waitScreenCapturing(self):
        self.isCapturing = True

    def isScreenCapturing(self):
        return self.isCapturing

    def setCenterObj(self, obj):
        self.centerObj = obj


    def updateData(self):
        try:
            self.message = self.client.listen()
        except:
            pass


    def init_GL(self, width=800, height=600):
        """ initialize OpenGL. This function has to be called only once before drawing. """
        glutInit([])

        # Open a window
        glutInitDisplayMode (GLUT_RGB | GLUT_DOUBLE | GLUT_DEPTH)
        self.width = width
        self.height = height
        glutInitWindowPosition (500, 0)
        glutInitWindowSize (self.width, self.height)
        self._myWindow = glutCreateWindow ("ODE Viewer")

        # Initialize Viewport and Shading
        glViewport(0, 0, self.width, self.height)
        glShadeModel(GL_SMOOTH)
        glHint(GL_PERSPECTIVE_CORRECTION_HINT, GL_NICEST)
        glClearColor(1.0, 1.0, 1.0, 0.0)

        # Initialize Depth Buffer
        glClearDepth(1.0)
        glEnable(GL_DEPTH_TEST)
        glDepthFunc(GL_LESS)

        # Initialize Lighting
        glEnable(GL_LIGHTING)
        glLightfv(GL_LIGHT1, GL_AMBIENT, [0.5, 0.5, 0.5, 1.0])
        glLightfv(GL_LIGHT1, GL_DIFFUSE, [1.0, 1.0, 1.0, 1.0])
        glLightfv(GL_LIGHT1, GL_POSITION, [0.0, 5.0, 5.0, 1.0])
        glEnable(GL_LIGHT1)

        # enable material coloring
        glEnable(GL_COLOR_MATERIAL)
        glColorMaterial(GL_FRONT, GL_AMBIENT_AND_DIFFUSE)

        glEnable(GL_NORMALIZE)

    def prepare_GL(self):
        """Prepare drawing. This function is called in every step. It clears the screen and sets the new camera position"""
        # Clear the screen
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

        # Projection mode
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective (45, 1.3333, 0.2, 500)

        # Initialize ModelView matrix
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()

        # View transformation (if "centerOn(...)" is set, keep camera to specific object)
        if self.centerObj is not None:
            (centerX, centerY, centerZ) = self.centerObj.getPosition()
        else:
            centerX = centerY = centerZ = 0
        # use the mouse to shift eye sensor on a hemisphere
        eyeX = self.viewDistance * self.lastx
        eyeY = self.viewDistance * self.lasty + centerY
        eyeZ = self.viewDistance * self.lastz
        gluLookAt (eyeX, eyeY, eyeZ, centerX, centerY, centerZ, 0, 1, 0)

    def draw_item(self, item):
        """ draws an object (spere, cube, plane, ...) """
        glDisable(GL_TEXTURE_2D)

        glPushMatrix()

        if item['type'] in ['GeomBox', 'GeomSphere', 'GeomCylinder', 'GeomCCylinder']:
            # set color of object (currently dark gray)
            if item.has_key('color'):
                glEnable (GL_BLEND)
                glBlendFunc (GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
                glColor4f(*(item['color']))
            else: glColor3f(0.1, 0.1, 0.1)

            # transform (rotate, translate) body accordingly
            (x, y, z) = item['position']
            R = item['rotation']

            rot = [R[0], R[3], R[6], 0.0,
                   R[1], R[4], R[7], 0.0,
                   R[2], R[5], R[8], 0.0,
                      x, y, z, 1.0]

            glMultMatrixd(rot)

            # switch different geom objects
            if item['type'] == 'GeomBox':
                # cube
                (sx, sy, sz) = item['scale']
                glScaled(sx, sy, sz)
                glutSolidCube(1)
            elif item['type'] == 'GeomSphere':
                # sphere
                glutSolidSphere(item['radius'], 20, 20)

            elif item['type'] == 'GeomCCylinder':
                quad = gluNewQuadric()
                # draw cylinder and two spheres, one at each end
                glTranslate(0.0, 0.0, -item['length'] / 2)
                gluCylinder(quad, item['radius'], item['radius'], item['length'], 32, 32)
                glutSolidSphere(item['radius'], 20, 20)
                glTranslate(0.0, 0.0, item['length'])
                glutSolidSphere(item['radius'], 20, 20)

            elif item['type'] == 'GeomCylinder':
                glTranslate(0.0, 0.0, -item['length'] / 2)
                quad = gluNewQuadric()
                gluDisk(quad, 0, item['radius'], 32, 1)
                quad = gluNewQuadric()
                gluCylinder(quad, item['radius'], item['radius'], item['length'], 32, 32)
                glTranslate(0.0, 0.0, item['length'])
                quad = gluNewQuadric()
                gluDisk(quad, 0, item['radius'], 32, 1)
            else:
                # TODO: add other geoms here
                pass

        elif item['type'] == 'GeomPlane':
            # set color of plane (currently green)
            if self.isFloorGreen:
                glColor3f(0.2, 0.6, 0.3)
            else:
                glColor3f(0.2, 0.3, 0.8)

            # for planes, we need a Quadric object
            quad = gluNewQuadric()
            gluQuadricTexture(quad, GL_TRUE)

            p = item['normal']      # the normal vector to the plane
            d = item['distance']    # the distance to the origin
            q = (0.0, 0.0, 1.0)     # the normal vector of default gluDisks (z=0 plane)

            # calculate the cross product to get the rotation axis
            c = crossproduct(p, q)
            # calculate the angle between default normal q and plane normal p
            theta = acos(dotproduct(p, q) / (norm(p) * norm(q))) / pi * 180

            # rotate the plane
            glPushMatrix()
            glTranslate(d * p[0], d * p[1], d * p[2])
            glRotate(-theta, c[0], c[1], c[2])
            gluDisk(quad, 0, 20, 20, 1)
            glPopMatrix()

        glPopMatrix()

    @staticmethod
    def _loadTexture(textureFile):
        image = open(textureFile)
        ix = image.size[0]
        iy = image.size[1]

        image = image.tostring("raw", "RGBX", 0, -1)

        # Create Texture
        textures = glGenTextures(3)
        glBindTexture(GL_TEXTURE_2D, textures[0])       # 2d texture (x and y size)

        glPixelStorei(GL_UNPACK_ALIGNMENT, 1)
        glTexImage2D(GL_TEXTURE_2D, 0, 3, ix, iy, 0, GL_RGBA, GL_BYTE, image)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        glTexEnvf(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_DECAL)

        # Create Linear Filtered Texture
        glBindTexture(GL_TEXTURE_2D, textures[1])
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
        glTexImage2D(GL_TEXTURE_2D, 0, 3, ix, iy, 0, GL_RGBA, GL_UNSIGNED_BYTE, image)

        # Create MipMapped Texture
        glBindTexture(GL_TEXTURE_2D, textures[2])
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_NEAREST)
        gluBuild2DMipmaps(GL_TEXTURE_2D, 3, ix, iy, GL_RGBA, GL_UNSIGNED_BYTE, image)
        return textures



    def _drawfunc (self):
        """ draw callback function """
        # Draw the scene
        self.prepare_GL()

        if self.message:
            for item in self.message:
                self.draw_item(item)

        glutSwapBuffers()
        if self.captureScreen:
            self._screenshot()

    def _idlefunc(self):
        self.updateData()
        t = self.dt - (time.time() - self.lasttime)
        if (t > 0):
            time.sleep(t)
        self.lasttime = time.time()
        glutPostRedisplay ()

    def _keyfunc (self, c, x, y):
        """ keyboard call-back function. """
        if c == 's':
            self.setCaptureScreen(not self.getCaptureScreen())
            print "Screen Capture: " + (self.getCaptureScreen() and "on" or "off")
        if c in ['x', 'q']:
            sys.exit()
        if c == 'v':
            self.mouseView = not self.mouseView

    def _motionfunc(self, x, z):
        """Control the zoom factor"""
        if not self.mouseView: return
        zn = 2.75 * float(z) / self.height + 0.25   # [0.25,3]
        self.viewDistance = 3.0 * zn * zn
        self._passivemotionfunc(x, z)

    def _passivemotionfunc(self, x, z):
        """ Store the mouse coordinates (relative to center and normalized)
         the eye does not exactly move on a unit hemisphere; we fudge the projection
         a little by shifting the hemisphere into the ground by 0.1 units,
         such that approaching the perimeter dows not cause a huge change in the
         viewing direction. The limit for l is thus cos(arcsin(0.1))."""
        if not self.mouseView: return
        x1 = 3 * float(x) / self.width - 1.5
        z1 = -3 * float(z) / self.height + 1.5
        lsq = x1 * x1 + z1 * z1
        l = sqrt(lsq)
        if l > 0.994987:
            # for mouse outside window, project onto the unit circle
            x1 = x1 / l
            z1 = z1 / l
            y1 = 0
        else:
            y1 = max(0.0, sqrt(1.0 - x1 * x1 - z1 * z1) - 0.1)
        self.lasty = y1
        self.lastx = x1
        self.lastz = z1

    def _screenshot(self, path_prefix='.', format='PNG'):
        """Saves a screenshot of the current frame buffer.
        The save path is <path_prefix>/.screenshots/shot<num>.png
        The path is automatically created if it does not exist.
        Shots are automatically numerated based on how many files
        are already in the directory."""

        if self.counter == self.frameT:
            self.counter = 1
            dir = os.path.join(path_prefix, 'screenshots')
            if not os.path.exists(dir):
                os.makedirs(dir)

            num_present = len(os.listdir(dir))
            num_digits = len(str(num_present))
            index = '0' * (5 - num_digits) + str(num_present)

            path = os.path.join(dir, 'shot' + index + '.' + format.lower())
            glPixelStorei(GL_PACK_ALIGNMENT, 1)
            data = glReadPixels(0, 0, self.width, self.height, GL_RGB, GL_UNSIGNED_BYTE)
            image = Image.fromstring("RGB", (self.width, self.height), data)
            image = image.transpose(Image.FLIP_TOP_BOTTOM)
            image.save(path, format)
            print 'Image saved to %s' % (os.path.basename(path))
        else:
            self.counter += 1

        self.isCapturing = False


if __name__ == '__main__':
    s = sys.argv[1:]
    odeview = ODEViewer(*s)
    odeview.start()

