import sys
import numpy as np 
import cv2
import matplotlib.pyplot as plt

sys.path.append('..') # so we can import from pupil
from pupil import player_methods

# from tracker import processFrame

DATA_DIR = '/home/mmbrian/HiWi/pupil_clone/pupil/recordings/2015_09_10/007/'
OUTPUT_DIR = DATA_DIR + 'pp.npy'


def capture(frame_number):
	cap = cv2.VideoCapture(DATA_DIR + "world.mp4")
	fc = int(cap.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT))
	assert frame_number<fc
	frame = 1
	status, img = cap.read() # extract the first frame
	while status:
		if frame == frame_number:
			save_dir = DATA_DIR + "frames/frame_%d.jpg" % frame_number
			cv2.imwrite(save_dir, img)	
			break
		frame+=1
		status, img = cap.read()

def main():
	p2d, t3d = [], [] # 2D-3D pupil position to target position correspondences

	cap = cv2.VideoCapture(DATA_DIR + "world.mp4")
	fc = int(cap.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT))
	
	wt = np.load(DATA_DIR + "world_timestamps.npy")

	# time is measured in seconds (floating point) since the epoch.
	# world timestamps correspond to each frame, we have to correlate timestamp information from
	# pupil positions and world timestamps to find a 1-to-1 mapping between pupil positions and 
	# marker positions in the video
	print fc, len(wt)
	assert fc == len(wt)

	########################################################################################################
	# Processing markers
	# print 'Processing markers...'

	## Processing frame by frame does not work as ArUco process opens a Gtk which cannot be terminated
	## automatically, therefore we better process all frames before and work with the data here

	# frame = 1
	# status, img = cap.read() # extract the first frame
	# while status:
	# 	# cv2.imwrite(DATA_DIR + "frames/frame-%d.jpg" % frame, img)
	# 	save_dir = DATA_DIR + "frames/current_frame.jpg"
	# 	cv2.imwrite(save_dir, img)

	# 	processFrame(save_dir)

	# 	frame+=1
	# 	status, img = cap.read()
	# print "Processed %d frames." % (frame-1)

	
	########################################################################################################
	# Processing pupil positions
	print 'Processing pupil positions...'
	pp = np.load(DATA_DIR + "pupil_positions.npy") # timestamp	confidence	id	pos_x	pos_y	diameter
	# pos_x and pos_y are normalized (Origin 0,0 at the bottom left and 1,1 at the top right)
	# converting each element to dictionary for correlation
	pp = map(lambda e: dict(zip(['timestamp', 'conf', 'id', 'x', 'y', 'diam'], e)), pp)
	pp_by_frame = player_methods.correlate_data(pp, wt)

	# Keeping only pupil positions with nonzero confidence
	pp_by_frame = map(lambda l: filter(lambda p: p['conf']>0, l), pp_by_frame)
	# Computing a single pupil position for the frame by taking mean of all detected pupil positions
	pp_by_frame = map(lambda data: 
		sum(np.array([pp['x'], pp['y']]) for pp in data)/len(data) if data else np.array([-1, -1]), pp_by_frame)
	# Now each nonempty value of pp_by_frame is a tuple of (x, y) for pupil position in that frame
	
	# Next we need to associate each frame to a detected marker and by taking mean pupil point and
	# mean 3D marker position over a series of frames corresponding to that marker find a 2D-3D 
	# mapping for calibration/test

	tdiff = map(lambda e: e-wt[0], wt)
	# This time correspondence to each marker was coordinated using the GazeHelper android application
	# for 005 > starting from 00:56, 3 seconds gaze, 1 second for saccade
	# These parameters are specific to the experiment
	# 005 > 56, 3, 1
	# 006 > 3, 3, 1
	# 007 > 7, 3, 1 (or 8)
	# 010 > 3, 3, 1
	starting_point, gaze_duration, saccade_duration = 56, 3, 1 # these are specific to the experiment
	# finding the starting frame
	ind = 0
	while tdiff[ind] < starting_point:
		ind+=1
	print ind

	data = []
	tstart = wt[ind]
	for i in xrange(9):
		print i
		while ind<len(wt) and wt[ind] - tstart < saccade_duration:
			ind+=1
		if ind<len(wt):
			tstart = wt[ind]
		starting_ind = ind

		while ind<len(wt) and wt[ind] - tstart < gaze_duration:
			ind+=1

		# all frames from starting_ind to ind-1 correspond to currently gazed marker
		c = 0
		cp = np.array([0, 0])
		all_corresponding_points = []
		for j in xrange(starting_ind, ind):
			if pp_by_frame[j][0] >= 0:
				c+=1
				cp = cp + pp_by_frame[j]
				all_corresponding_points.append(pp_by_frame[j])
		# print c
		if c>0:
			ret = cp/c
		else:
			ret = np.array([-1, -1])  # no detected pupil for this marker
		p2d.append(ret)
		data.append([
			np.array([starting_ind, ind-1]), # frame range
			ret, # mean pupil position
			all_corresponding_points]) # all pupil positions in range

		if ind<len(wt):
			tstart = wt[ind]

	# p2d is now the list of detected pupil positions (not always 9 points)
	print 'Saving data...'
	np.save(OUTPUT_DIR, data)

	# plt.plot([x[0] if x!=None else 0 for x in pp_by_frame], [y[1] if y!=None else 0 for y in pp_by_frame], 'ro')
	# plt.show()
	########################################################################################################

	print len(p2d), 'gaze points'
	print p2d
	print len(wt), 'frames...'

if __name__ == '__main__':
	main()