#!/usr/bin/python

import cv
import sys
import math
import pickle
import signal

sk_file = open(sys.argv[1])
pic_dir = sys.argv[2].rstrip('/') + '/' + sys.argv[1].split('/')[-2] + '/'
out_dir = sys.argv[1][0:sys.argv[1].rfind('/')+1]

j1 = 'Head'
j2 = 'ShoulderCenter'

def display(lines):
    global j1,j2
    limbs = {}
    pframe = 0
    while cv.WaitKey(5) != -1: pass
    for line in lines:
        run, frame, nid, joint, state, x, y, z, X, Y = line.strip().split(',')
        frame = int(frame)
        if frame != pframe:
            if j1 in limbs and j2 in limbs:
                dx = limbs[j1][0] - limbs[j2][0]
                dy = limbs[j1][1] - limbs[j2][1]

                img = cv.LoadImage(pic_dir+str(pframe)+".jpg")
                r = math.sqrt(dx*dx+dy*dy)
                x1 = max(int(limbs[j1][0] - r),0)
                y1 = max(int(limbs[j1][1] - r),0)
                x2 = min(int(x1 + 2*r), img.width)
                y2 = min(int(y1 + 3*r), img.height)

                cv.Rectangle(img,(x1,y1),(x2,y2),(0,255,0),3,8,0)
                #cv.SaveImage(out_dir+str(run)+"-"+str(frame)+'.jpg',img)
                cv.ShowImage('video', img)
                k = cv.WaitKey(5)
                if k != -1:
                    return chr(k&255)
            pframe = frame 
            limbs = {}
        id = int(nid)

        limbs[joint] = [int(X)*640.0/400.0,int(Y)*480.0/300.0]
    return chr(cv.WaitKey()&255)

def quit(save):
    global people,labels
    print save
    if save:
        pickle.dump(people,open(out_dir+'people.dump','w'))
        pickle.dump(labels,open(out_dir+'labels.dump','w'))
    sys.exit(0)

prun = 0
lines = []
people = []
#chars = '123456789qwerasdfzxcvQWERASDFZXCV'
chars = 'abcdefghijklmnopqrstuvwxyz'
try:
    people = pickle.load(open(out_dir+'people.dump'))
    people.sort()
except:
    pass
labels = {}
try:
    labels = pickle.load(open(out_dir+'labels.dump'))
except:
    pass
signal.signal(signal.SIGINT, lambda *args: quit(True))
signal.signal(signal.SIGQUIT, lambda *args: quit(False))
for line in sk_file:
    run, frame, nid, joint, state, x, y, z, X, Y = line.strip().split(',')
    if frame != 'Frame':
        run = int(run)
        if run != prun:
            if prun != 0 and prun not in labels:
                k = '\n'
                while k == '\n':
                    print "Run #",prun
                    print "0: None"
                    for i in range(len(people)):
                        print chars[i]+": "+people[i]
                    print "Enter: Replay"
                    #print "N: New person"
                    print "?: Unknown"
                    k = display(lines)
                    # Key pressed is SHIFT, grab next character and make it uppercase
                    if k == '\xe1':
                        k = chr(cv.WaitKey()&255).upper()

                    # Key pressed is CTRL, grab next character and make quit if 'c','d', or '\' follows
                    if k == '\xe3':
                        k = chr(cv.WaitKey()&255)
                        if k == 'c':
                            quit(True)
                        elif k == 'd' or k == '\\':
                            quit(False)

                    if k == '0':
                        labels[prun] = 'None'
                    elif k == '?':
                        labels[prun] = '?'
                    elif k in chars:
                        labels[prun] = people[chars.find(k)]
                    elif k == 'N':
                        name = raw_input("Type name: ")
                        if people != "":
                            people += name.split(',')
                            people.sort()
                        #labels[prun] = name
                        k = '\n'
                    elif k != '~':
                        k = '\n'
                    if cv.WaitKey(1000) != -1:
                        k = '\n'
            prun = run
            lines = []
        lines += [line]
quit(True)