from __future__ import division

import os
import seaborn
from pylab import rcParams
import cv2

import numpy as np
from numpy import linalg as LA
from time import time
from itertools import combinations

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

from minimize import findInitialW, _q, g, minimizeEnergy, g3D3D, findW3D3D
from minimize import g as gaze_ray
from geom import getSphericalCoords, getAngularDiff
from vector import Vector as v
# from experiment import Experiment
from sim import GazeSimulation
from recording.util.tools import is_outlier
from recording.tracker import readCameraParams
from geom import getRotationMatrix


class Experiment:
	def performExperiment(self):
		print 'Running Experiment...'
		start = time()
		self.__run__()
		self.runningTime = time() - start
		print 'Running Time:', self.runningTime
	def __run__(self):
		pass

class Parallax2Dto3DMapping(Experiment):
	def __run__(self):
		sim = GazeSimulation(log = False)

		sim.place_eyeball_on_scene_camera = False
		sim.setEyeRelativeToSceneCamera(v(-65, -33, -73))
		# sim.setEyeRelativeToSceneCamera(v(-65, -33, 0)) # assuming eyeball and scene camera are coplanar i.e. e = (e.x, e.y, 0)
		sim.setCalibrationDepth(1 * 1000) # mm, wrt scene camera
		sim.setTestDepth(1.5 * 1000)
		sim.calibration_grid = True
		sim.calibration_random_depth = False
		sim.test_grid = True
		sim.test_random_depth = False
		sim.test_random_fixed_depth = False

		depths = map(lambda d:d*1000, [1, 1.25, 1.5, 1.75, 2.0])
		print '> Computing results for multiple calibration depths...'
		results = []
		for num_of_calibration_depths in xrange(1, 6): # from 1 calibration depths to 5
			print '> Considering only %s calibration depth(s)...' %num_of_calibration_depths
			sim.reset()	
			if num_of_calibration_depths != 3: continue

			for calibs in combinations(depths, num_of_calibration_depths):
				aae_2ds_aae = []
				aae_2ds_phe = []
				aae_3ds_aae = []
				aae_3ds_phe = []
				aae_3D3Ds = [] # angular error

				# Now calibs is a set of depths, from each of those we need calibration data
				# print calibs
				if not calibs in [(1000., 1250., 1500.), (1250., 1500., 1750.), (1500., 1750., 2000.)]:
					continue
				print 'curr calibs', calibs
				calibs = list(calibs)
				cp, ct = [], []
				sim.reset()	
				sim.setCalibrationDepth(calibs)

				# Perform calibration
				sim.runCalibration()
				cp, ct, p3d = sim.tr_pupil_locations, sim.calibration_points, sim.tr_3d_pupil_locations

				# target positions are computed relative to the scene CCS
				ti = map(lambda target: v(target) - v(sim.scene_camera.t), ct)
				# Computing pupil pose for each gaze
				ni = map(lambda p: (v(p)-v(sim.sclera_pos)).norm(), p3d) # ground truth gaze vectors
				
				w, e, w0 = minimizeEnergy(cp, ti)
				e = v(e)

				# transforming pupil pose to eye camera CS
				eyeR = np.array(sim.eye_camera.R[:3])
				ni = map(lambda pose: eyeR.dot(np.array(pose)), ni) 
				
				R, e3d3d = minimizeEnergy(ni, ti, pose_given=True)
				e3d3d = v(e3d3d)
				# R = LA.inv(R)

				# Now we have calibration data from multiple depths, we can test on all depths
				for test_depth in depths:
					sim.setTestDepth(test_depth)
					aae_2d_aae, aae_2d_phe, aae_2d_std, _ = sim.runTest() # 2nd one is PHE
					aae_2ds_aae.append((aae_2d_aae, aae_2d_std))
					aae_2ds_phe.append(aae_2d_phe)

					# Fetching test points
					t, p, p3d = sim.test_points, sim.te_pupil_locations, sim.te_3d_pupil_locations
					t = map(lambda target: v(target) - v(sim.scene_camera.t), t) # target coords in scene CCS

					# 3D3D
					t_3d3d = t[:]

					ni = map(lambda p: v(v(p)-v(sim.sclera_pos)).norm(), p3d) # ground truth gaze vectors
					# transforming pupil pose to eye camera CS
					ni = map(lambda r: v(eyeR.dot(np.array(r))), ni)

					# applying estimated rotation to pose vector in eye camera coordinates (Rn)
					# R is estimated rotation between scene camera and eye coordinate system (not eye camera!)
					# in other words, R is the rotation part of e
					Rni = map(lambda n: v(R.dot(np.array(n))), ni) # now ready to compare Rn with t-e
					# Intersecting gaze rays originating from the eye with the planes defined by each
					# target. then we can simply compute angular error between each intersection and 
					# the corresponding 3D target
					gis = map(lambda vec: v(vec), Rni) # gaze rays originating from eyeball
					# we multiply g such that it hits t's z-plane i.e. multiply all coordinates by factor (t.z-e.z)/g.z
					# then we add e to the final g so that it originates from scene camera. now both g and t are in the
					# same coordinate system and originate from the same point, so we can compare them
					gprimes = map(lambda tg: v(((tg[0].z - e3d3d.z)/tg[1].z)*tg[1] + e3d3d), zip(t_3d3d, gis))
					AE = list(np.degrees(np.arctan((v(p[0]).cross(p[1])/(v(p[0]).dot(p[1]))).mag)) for p in zip(gprimes, t_3d3d))

					N = len(t)
					AAE = np.mean(AE)
					STD = np.std(AE)
					m, M = min(AE), max(AE)
					aae_3D3Ds.append((AAE, STD))


					qi = map(_q, p) # computing feature vectors from raw pupil coordinates in 2D
					# computing unit gaze vectors corresponding to pupil positions
					# here we use the computed mapping matrix w
					gis = map(lambda q: g(q, w), qi) 
					
					# Intersecting gaze rays originating from the eye with the planes defined by each
					# target. then we can simply compute angular error between each intersection and 
					# the corresponding 3D target
					t = map(lambda vec: v(vec), t)
					gis = map(lambda vec: v(vec), gis)
					gprimes = map(lambda tg: v(((tg[0].z - e.z)/tg[1].z)*tg[1] + e), zip(t, gis))

					AE = list(np.degrees(np.arctan((v(p[0]).cross(p[1])/(v(p[0]).dot(p[1]))).mag)) for p in zip(gprimes, t))
					N = len(t)
					AAE = np.mean(AE)
					STD = np.std(AE)
					m, M = min(AE), max(AE)

					# Computing physical distance error (in meters)
					PHE = list((u-v).mag/1000 for u,v in zip(t, gprimes))
					N = len(t)
					APHE = np.mean(PHE)
					PHE_STD = np.std(PHE)
					PHE_m, PHE_M = min(PHE), max(PHE)

					aae_3ds_aae.append((AAE, STD))
					aae_3ds_phe.append((PHE, PHE_STD))
				

				results.append([calibs, aae_2ds_aae, aae_3ds_aae, aae_3D3Ds])

		######################################################################################################
		## Plotting part
		plt.ylabel('Angular Error')
		plt.xlabel('Depth')

		fig = plt.figure(figsize=(14.0, 10.0))
		ax = fig.add_subplot(111)

		clrs = ['blue', 'red', 'orange']
		depth_map = {1000.:1, 1250.:2, 1500.:3, 1750.:4, 2000.:5}
		cdlabel = lambda c: ' + '.join(map(str, map(lambda d: depth_map[d], c)))

		x1 = [0.375,1.375,2.375,3.375,4.375]
		x2 = [0.425,1.425,2.425,3.425,4.425]
		x3 = [0.475,1.475,2.475,3.475,4.475]
		x4 = [0.525,1.525,2.525,3.525,4.525]
		x5 = [0.575,1.575,2.575,3.575,4.575]
		x6 = [0.625,1.625,2.625,3.625,4.625]
		xrange_2d2d = [x1, x3, x5]
		xrange_2d3d = [x2, x4, x6]

		for i in [0, 1, 2]:
			res = results[i]
			calibs = res[0]

			_xrange = xrange_2d2d[i]
			aae_2d2d = np.array(res[1])[:,0]
			std_2d2d = np.array(res[1])[:,1]
			rects1 = ax.errorbar(_xrange, aae_2d2d,yerr=[std_2d2d, std_2d2d],fmt='o',color=clrs[i],ecolor=clrs[i],lw=3, capsize=5, capthick=2)
			plt.plot(_xrange, aae_2d2d, marker="o", linestyle='-',lw=3,color=clrs[i],label = '2D-to-2D Calibration Depth '+cdlabel(calibs))

		for i in [0, 1, 2]:
			res = results[i]
			calibs = res[0]

			_xrange = xrange_2d3d[i]
			aae_2d3d = np.array(res[2])[:,0]
			std_2d3d = np.array(res[2])[:,1]
			rects2 = ax.errorbar(_xrange, aae_2d3d,yerr=[std_2d3d, std_2d3d],fmt='o',color=clrs[i],ecolor=clrs[i],lw=3, capsize=5, capthick=2)
			plt.plot(_xrange, aae_2d3d, marker="o", linestyle='--',lw=3,color=clrs[i],label ='2D-to-3D Calibration Depth '+cdlabel(calibs))

		for i in [0, 1, 2]:
			res = results[i]
			calibs = res[0]

			_xrange = xrange_2d2d[i]
			aae_3d3d = np.array(res[3])[:,0]
			std_3d3d = np.array(res[3])[:,1]
			rects3 = ax.errorbar(_xrange, aae_3d3d,yerr=[std_3d3d, std_3d3d],fmt='o',color=clrs[i],ecolor=clrs[i],lw=3, capsize=5, capthick=2)
			plt.plot(_xrange, aae_3d3d, marker="o", linestyle='-.',lw=3,color=clrs[i],label ='3D-to-3D Calibration Depth '+cdlabel(calibs))

		ax.set_ylim(0, 1.5)

		ax.set_ylabel(r'Angular Error',fontsize=22, fontweight='bold')
		ax.set_xlabel(r'Depth',fontsize=22, fontweight='bold')
		
		TOPICS = [-0.2, 0, 0.2, 0.4,0.6,0.8,1.0,1.2,1.4,1.6,1.8,2.0,2.2,2.4]#,110]#,120]
		LABELS = [r'', r'0', r'0.2',r'0.4',r'0.6',  r'0.8', r'1.0', r'1.2', r'1.4',  r'1.6',  r'1.8', r'2.0',  r'2.2', r'2.4']#, ""]#, ""]

		plt.yticks(TOPICS, LABELS,fontsize=18)

		plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
           ncol=3, mode="expand", borderaxespad=0., fontsize=15)
		plt.yticks(fontsize='18')

		TOPICs = [0.0,0.5,1.5,2.5,3.5,4.5,5.0]
		LABELs = ['', 'D1 - 1m', 'D2 - 1.25m', 'D3 - 1.5m', 'D4 - 1.75m', 'D5 - 2.0m', '']
		# ax.set_xticklabels(LABELs,fontsize=18)
		plt.xticks(TOPICs, LABELs,fontsize=18)
		
		left  = 0.1  # the left side of the subplots of the figure
		right = 0.975    # the right side of the subplots of the figure
		bottom = 0.075   # the bottom of the subplots of the figure
		top = 0.925      # the top of the subplots of the figure
		wspace = 0.2   # the amount of width reserved for blank space between subplots
		hspace = 0.4   # the amount of height reserved for white space between subplots

		plt.subplots_adjust(left=left, bottom=bottom, right=right, top=top, wspace=wspace, hspace=hspace)
		plt.show()
		######################################################################################################

if __name__ == '__main__':
	# This also does 3D gaze estimation and plots estimation results for both 3D and 2D estimation
	ex = Parallax2Dto3DMapping()
	ex.performExperiment()