# -*- coding: utf-8 -*-
"""
Affichage vidéo (et autres) de la kinect pour expérimentations / debug.

$Id$
$URL$
"""

from openni import *
import numpy
import cv
import pygame

SCREEN_SIZE = 640, 480
SCREEN_TITLE = "Kinect debug"
FPS = 30
MOUSE_SCREEN_COORDINATES = 0
MOUSE_REAL_COORDINATES = 1
DEFAULT_HAND_DISTANCE_RANGE = 150 # distance de débatement, autour du centre, en mm

class RGB :
    def __init__(self, mouseMode=MOUSE_REAL_COORDINATES) :
        self.mouseMode = mouseMode
        self.context = Context()
        self.context.init()
        self.imgGene = ImageGenerator()
        self.imgGene.create(self.context)
        self.imgGene.set_resolution_preset(RES_VGA)
        self.imgGene.fps = FPS
        
        self.depthGene = DepthGenerator()
        self.depthGene.create(self.context)
        self.depthGene.set_resolution_preset(RES_VGA)
        self.depthGene.fps = FPS
        
        self.handsGene = HandsGenerator()
        self.handsGene.create(self.context)
        self.handsGene.register_hand_cb(self.handCreateCB, self.handUpdateCB, self.handDestroyCB)
        
        self.gestGene = GestureGenerator()
        self.gestGene.create(self.context)
        self.gestGene.add_gesture('Click')
        self.gestGene.register_gesture_cb(self.gestureRecognizedCB, self.gestureProgressCB)
        

        self.userGene = UserGenerator()
        self.userGene.create(self.context)
        self.userGene.register_user_cb(self.newUserCB, self.lostUserCB)
        
        self.context.start_generating_all()
        
        screen = pygame.display.get_surface()
        self.xratio = float(screen.get_size()[0]) / SCREEN_SIZE[0]
        self.yratio = float(screen.get_size()[1]) / SCREEN_SIZE[1]
        self.handCenteredPosition = None
        self.pygameScreenWidth, self.pygameScreenHeight = screen.get_size()
    
    def getProjPos(self, realPos) :
        return self.depthGene.to_projective([realPos])[0]
    
    def setMousePosition(self, realPos) :
        if self.mouseMode == MOUSE_SCREEN_COORDINATES :
            x, y, z = self.getProjPos(realPos)
            # mirror
            x = SCREEN_SIZE[0] - x
            # scale
            x, y = x * self.xratio, y * self.yratio
        elif self.mouseMode == MOUSE_REAL_COORDINATES :
            x, y = realPos[:2]
            x, y = x - self.handCenteredPosition[0], y - self.handCenteredPosition[1]
            x = -x # miroir
            x, y = x + DEFAULT_HAND_DISTANCE_RANGE, DEFAULT_HAND_DISTANCE_RANGE - y
            x, y = map(lambda i : numpy.clip(i, 0, DEFAULT_HAND_DISTANCE_RANGE * 2), [x, y])
            x = x * (self.pygameScreenWidth / (2. * DEFAULT_HAND_DISTANCE_RANGE))
            y = y * (self.pygameScreenHeight / (2. * DEFAULT_HAND_DISTANCE_RANGE))

        pygame.mouse.set_pos(int(x), int(y))

    def handCreateCB(self, src, id, pos, time):
        print 'Create ', id, pos

    def handUpdateCB(self, src, id, pos, time):
        self.setMousePosition(pos)

    def handDestroyCB(self, src, id, time):
        print 'Destroy ', id

    def gestureRecognizedCB(self, src, gesture, id, end_point) :
        print 'gestureDetected', src, gesture, id, end_point

    def gestureProgressCB(self, src, gesture, point, progress) :
        print 'gestureProgress', src, gesture, point, progress
        self.handsGene.start_generating()
        self.handsGene.start_tracking(point)
        self.handCenteredPosition = point

    def newUserCB(self, *args, **kw) :
        print 'newUserCB', args, kw

    def lostUserCB(self, *args, **kw) :
        print 'lostUserCB', args, kw
    
    
    def capture(self) :
        rgb_frame = numpy.fromstring(self.imgGene.get_raw_image_map_bgr(), dtype=numpy.uint8).reshape(480, 640, 3)
        image = cv.fromarray(rgb_frame)
        cv.CvtColor(cv.fromarray(rgb_frame), image, cv.CV_BGR2RGB)
        pyimage = pygame.image.frombuffer(image.tostring(), cv.GetSize(image), 'RGB')
        return pyimage
    
    def update(self) :
        self.context.wait_any_update_all()
        #self.context.wait_and_update_all()
        #return self.context.wait_one_update_all(self.imgGene)


class RGBSprite(pygame.sprite.DirtySprite, RGB) :
    
    def __init__(self, alpha=255, size=SCREEN_SIZE) :
        pygame.sprite.DirtySprite.__init__(self)
        RGB.__init__(self)
        self.dirty = 2 # toujours dirty !
        self.size = size
        self.image = pygame.Surface(size)
        self.workSur = pygame.Surface(SCREEN_SIZE)
        self.image.set_alpha(alpha)
        self.rect = pygame.Rect((0, 0), (0, 0))
    
    def update(self) :
        RGB.update(self)
        img = self.capture()
        self.workSur.blit(img, (0, 0))
        self.workSur = pygame.transform.flip(self.workSur, True, False) # miroir
        if self.size != SCREEN_SIZE :
            pygame.transform.scale(self.workSur, self.size, self.image) # étirement, blit implicite
        else :
            self.image.blit(self.workSur, (0, 0))

    
def main() :
    pygame.init()
    screen = pygame.display.set_mode(SCREEN_SIZE)
    pygame.display.set_caption(SCREEN_TITLE)
    
    rgb = RGB()
    
    sur = pygame.Surface((640, 480))
    sur.fill((255, 255, 255))
    
    while True :
        for event in pygame.event.get():
            pass
        
        rgb.update()

        rgbImg = rgb.capture()
        sur.blit(rgbImg, (0, 0))
        screen.blit(pygame.transform.flip(sur, True, False), (0, 0))
        pygame.display.flip()
    

if __name__ == "__main__" :
    main()