from __future__ import division
import os, sys, cv2
import numpy as np
import shutil

# PARTICIPANTS = ['p10', 'p16', 'p13', 'p24', 'p5', 'p14', 'p26', 'p12', 'p20', 'p7', 'p15', 'p11', 'p21', 'p25']
# PARTICIPANTS = ['p21']
# ROOT = '/BS/3D_Gaze_Tracking/archive00/participants/'
# DEST = '/BS/3D_Gaze_Tracking/work/Eye_Images/'
DEST = '/home/mmbrian/3D_Gaze_Tracking/work/Marker_Eye_Images_7/'
ROOT = '/home/mmbrian/HiWi/etra2016_mohsen/code/recording/data/participants'

def main():
	# distCoeffs = np.load("/BS/3D_Gaze_Tracking/work/dist.npy")
	# cameraMatrix = np.load("/BS/3D_Gaze_Tracking/work/cameraMatrix.npy")
	distCoeffs = np.load("dist.npy")
	cameraMatrix = np.load("cameraMatrix.npy")

	PARTICIPANTS = [sys.argv[1]]

	processed = 0
	for p in os.listdir(ROOT):
		if p in PARTICIPANTS:
			print '> Collecting images for participant', p
			d1 = os.path.join(ROOT, p)
			d1 = os.path.join(d1, os.listdir(d1)[0]) # ../p_i/../
			for d2 in os.listdir(d1):
				path = os.path.join(d1, d2)
				print '> Processing', path			
				processPath(path, p, d2, cameraMatrix, distCoeffs)
			processed+=1
			print '> Processed %s participants.' % processed

def processPath(path = None, participant = None, experiment = None, cameraMatrix = None, distCoeffs = None):
	if not path:
		path = sys.argv[1]
	raw_images_dir = os.path.join(DEST, 'ImagesRaw')
	raw_images_dir = os.path.join(raw_images_dir, participant)
	raw_images_dir = os.path.join(raw_images_dir, experiment)
	undist_images_dir = os.path.join(DEST, 'ImagesUndist')
	undist_images_dir = os.path.join(undist_images_dir, participant)
	undist_images_dir = os.path.join(undist_images_dir, experiment)
	if not os.path.exists(raw_images_dir):
		os.makedirs(raw_images_dir)
	else:
		print '> Already processed.'
		return
	if not os.path.exists(undist_images_dir):
		os.makedirs(undist_images_dir)
	else:
		print '> Already processed.'
		return
	# else:
	# 	print '> Removing old images...'
	# 	shutil.rmtree(raw_images_dir)
	# 	return
	# 	os.makedirs(raw_images_dir)

	p_frames = np.load(os.path.join(path, 'p_frames.npy'))
	frames = sorted(reduce(lambda l1,l2: list(l1)+list(l2), p_frames))

	cap = cv2.VideoCapture(os.path.join(path, 'eye0.mp4'))
	fc_eye = int(cap.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT))
	fps_eye = int(cap.get(cv2.cv.CV_CAP_PROP_FPS))

	eye_frame = 0
	world_frame = 0
	status, img = cap.read() # extract the first frame
	while status:
		try:
			if eye_frame == frames[world_frame]:
				save_dir = os.path.join(raw_images_dir, 'img_%s.png' %(frames[world_frame]))
				cv2.imwrite(save_dir, img)

				save_dir = os.path.join(undist_images_dir, 'img_%s.png' %(frames[world_frame]))
				undistImg = cv2.undistort(img, cameraMatrix, distCoeffs)
				cv2.imwrite(save_dir, undistImg)
				world_frame+=1
		except:
			break
		eye_frame+=1
		status, img = cap.read()
	cap.release()
	# print '> Removing world video...'
	# os.remove(world_video)
	print "> Processed %d frames." % (world_frame)

if __name__ == '__main__':
	main()