from __future__ import division

import os, sys
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
from matplotlib.pyplot import *
import matplotlib.patches as mpatches

sys.path.append('..') # so we can import modules from `code` directory
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 sim import GazeSimulation
from recording.util.tools import is_outlier
from recording.tracker import readCameraParams
from geom import getRotationMatrix
from parallax_analysis import Experiment

'''
Change lines 48 to 61 accordingly
'''


class Parallax2Dto3DMapping(Experiment):
    '''
    IMPORTANT!
    In all experiments, scene camera's rvec = (0, 0, 0) i.e. the corresponding rotation matrix is the identity matrix therefore
    I have not included the dot production with this rotation matrix to convert points in world coordinates
    into scene camera coordinates. however, one should know that if the scene camera is rotated differentl7y
    this transformation is essential. I would add the corresponding computations later on.
    '''

    def __run__(self):
        # Processing real world data
        ######################################################################################################
        print '> Processing real world data...'

        C12D2D = np.load("../results/MeansC1D2D2.npy")
        C12D3D = np.load("../results/MeansC1D2D3.npy")

        C22D2D = np.load("../results/MeansC2D2D2.npy")
        C22D3D = np.load("../results/MeansC2D2D3.npy")

        C32D2D = np.load("../results/MeansC3D2D2.npy")
        C32D3D = np.load("../results/MeansC3D2D3.npy")

        C42D2D = np.load("../results/MeansC4D2D2.npy")
        C42D3D = np.load("../results/MeansC4D2D3.npy")

        C52D2D = np.load("../results/MeansC5D2D2.npy")
        C52D3D = np.load("../results/MeansC5D2D3.npy")


        summeC12D2D = [] 
        summeC22D2D = [] 
        summeC32D2D = [] 
        summeC42D2D = [] 
        summeC52D2D = [] 

        summeC12D3D = [] 
        summeC22D3D = [] 
        summeC32D3D = [] 
        summeC42D3D = [] 
        summeC52D3D = [] 


        i = 0
        while i < len(C12D2D):
            j = 0
            while j < len(C12D2D[0]):
                summeC12D2D.append(C12D2D[i][j])
                j += 1
            i += 1
            
        i = 0
        while i < len(C22D2D):
            j = 0
            while j < len(C22D2D[0]):
                summeC22D2D.append(C22D2D[i][j])
                j += 1
            i += 1
            
        i = 0
        while i < len(C32D2D):
            j = 0
            while j < len(C32D2D[0]):
                summeC32D2D.append(C32D2D[i][j])
                j += 1
            i += 1
            
        i = 0
        while i < len(C42D2D):
            j = 0
            while j < len(C42D2D[0]):
                summeC42D2D.append(C42D2D[i][j])
                j += 1
            i += 1
            
        i = 0
        while i < len(C52D2D):
            j = 0
            while j < len(C52D2D[0]):
                summeC52D2D.append(C52D2D[i][j])
                j += 1
            i += 1
            
        i = 0
        while i < len(C12D3D):
            j = 0
            while j < len(C12D3D[0]):
                summeC12D3D.append(C12D3D[i][j])
                j += 1
            i += 1
            
        i = 0
        while i < len(C22D3D):
            j = 0
            while j < len(C22D3D[0]):
                summeC22D3D.append(C22D3D[i][j])
                j += 1
            i += 1
            
        i = 0
        while i < len(C32D3D):
            j = 0
            while j < len(C32D3D[0]):
                summeC32D3D.append(C32D3D[i][j])
                j += 1
            i += 1
            
        i = 0
        while i < len(C42D3D):
            j = 0
            while j < len(C42D3D[0]):
                summeC42D3D.append(C42D3D[i][j])
                j += 1
            i += 1
            
        i = 0
        while i < len(C52D3D):
            j = 0
            while j < len(C52D3D[0]):
                summeC52D3D.append(C52D3D[i][j])
                j += 1
            i += 1
            
            
        mean1 = np.mean(summeC12D2D)
        mean2 = np.mean(summeC22D2D) 
        mean3 = np.mean(summeC32D2D) 
        mean4 = np.mean(summeC42D2D) 
        mean5 = np.mean(summeC52D2D) 

        mean6 = np.mean(summeC12D3D) 
        mean7 = np.mean(summeC22D3D)
        mean8 = np.mean(summeC32D3D) 
        mean9 = np.mean(summeC42D3D) 
        mean10 = np.mean(summeC52D3D) 

        std1 = np.std(summeC12D2D)
        std2 = np.std(summeC22D2D) 
        std3 = np.std(summeC32D2D) 
        std4 = np.std(summeC42D2D) 
        std5 = np.std(summeC52D2D) 

        std6 = np.std(summeC12D3D) 
        std7 = np.std(summeC22D3D)
        std8 = np.std(summeC32D3D) 
        std9 = np.std(summeC42D3D) 
        std10 = np.std(summeC52D3D) 
            
        mean2D2D_real = [mean1,mean2,mean3,mean4,mean5]
        mean2D3D_real = [mean6,mean7,mean8,mean9,mean10]

        std2D2D_real = [std1,std2,std3,std4,std5]
        std2D3D_real = [std6,std7,std8,std9,std10]
        ######################################################################################################
        # Simulation
        print '> Processing simulation data...'
        ######################################################################################################
        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, results_std = [], []
        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() 
            aae_2ds_aae = []
            aae_2ds_phe = []
            aae_3ds_aae = []
            aae_3ds_phe = []
            aae_3D3Ds = [] # angular error

            for calibs in combinations(depths, num_of_calibration_depths):
                # Now calibs is a set of depths from each of which we need calibration data
                print '> Current calibration depths', 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() # last one is PHE std
                    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 only contains AAE
            results.append([np.mean(np.array(aae_2ds_aae)[:,0]), np.mean(np.array(aae_3ds_aae)[:,0]), np.mean(np.array(aae_3D3Ds)[:,0])])
            results_std.append([np.std(np.array(aae_2ds_aae)[:,0]), np.std(np.array(aae_3ds_aae)[:,0]), np.std(np.array(aae_3D3Ds)[:,0])])
        ######################################################################################################
        # Plotting
        print '> Plotting...'
        ######################################################################################################
        # New plot code based on EffectNumberofClusters.py
        mean2D2D = [res[0] for res in results]
        mean2D3D = [res[1] for res in results]
        mean3D3D = [res[2] for res in results]
        std2D2D = [res[0] for res in results_std]
        std2D3D = [res[1] for res in results_std]
        std3D3D = [res[2] for res in results_std]


        N = 5
        ind = np.asarray([0.25,1.25,2.25,3.25,4.25])

        width = 0.5       # the width of the bars

        #    x1 = [0.4,1.4,2.4,3.4,4.4]
        x2 = [0.45,1.45,2.45,3.45,4.45]
        #    x3 = [0.5,1.5,2.5,3.5,4.5]
        x4 = [0.55,1.55,2.55,3.55,4.55]
        #    x5 = [0.6,1.6,2.6,3.6,4.6]
        x6 = [0.50,1.50,2.50,3.50,4.50]

        fig = plt.figure(figsize=(14.0, 10.0))

        ax = fig.add_subplot(111)

        rrects1 = ax.errorbar(x2, mean2D2D_real,yerr=[std2D2D_real,std2D2D_real],fmt='o',color='red',ecolor='red',lw=3, capsize=8, capthick=3)
        plt.plot(x2, mean2D2D_real, marker="o", linestyle='-',lw=3,color='red',label = r'2D-to-2D')

        rrects2 =ax.errorbar(x4, mean2D3D_real,yerr=[std2D3D_real,std2D3D_real],fmt='o',color='blue',ecolor='blue',lw=3, capsize=8, capthick=3)
        plt.plot(x4, mean2D3D_real, marker="o", linestyle='-',lw=3,color='blue', label = r'2D-to-3D')


        rects1 = ax.errorbar(x2, mean2D2D,yerr=[std2D2D,std2D2D],fmt='o',color='red',ecolor='red',lw=3, capsize=5, capthick=2)
        plt.plot(x2, mean2D2D, marker="o", linestyle='--',lw=3,color='red',label = r'2D-to-2D Simulation')

        rects2 =ax.errorbar(x4, mean2D3D,yerr=[std2D3D,std2D3D],fmt='o',color='blue',ecolor='blue',lw=3, capsize=5, capthick=2)
        plt.plot(x4, mean2D3D, marker="o", linestyle='--',lw=3,color='blue', label = r'2D-to-3D Simulation')

        rects3 =ax.errorbar(x6, mean3D3D,yerr=[std3D3D,std3D3D],fmt='o',color='orange',ecolor='orange',lw=3, capsize=5, capthick=2)
        plt.plot(x6, mean3D3D, marker="o", linestyle='--',lw=3,color='orange', label = r'3D-to-3D Simulation')

        legend(fontsize=20,loc='upper right')

        #    rects3 = ax.errorbar(x3, meanC3,yerr=[stdC3,stdC3],fmt='o',color='black',ecolor='black',lw=3, capsize=5, capthick=2)
        #    plt.plot(x3, meanC3, marker="o", linestyle='-',lw=3,color='black')
        #    
        #    rects4 =ax.errorbar(x4, meanC4,yerr=[stdC4,stdC4],fmt='o',color='green',ecolor='green',lw=3, capsize=5, capthick=2)
        #    plt.plot(x4, meanC4, marker="o", linestyle='-',lw=3,color='green')
        #    
        #    rects5 =ax.errorbar(x5, meanC5,yerr=[stdC5,stdC5],fmt='o',color='orange',ecolor='orange',lw=3, capsize=5, capthick=2)
        #    plt.plot(x5, meanC5, marker="o", linestyle='-',lw=3,color='orange')


        ax.set_ylabel(r'Angular Error',fontsize=22)
        ax.set_xlabel(r'Number of Calibration Depths',fontsize=22)
        ax.set_xticks(ind+0.25)
        ax.set_xticklabels( ('D1', 'D2', 'D3','D4', 'D5') ,fontsize=18)

        TOPICs = [0.0,0.5,1.5,2.5,3.5,4.5,5.0]#,110]#,120]
        print TOPICs
        LABELs = ["",r'1',r'2', r'3', r'4', r'5', ""]#, ""]#, ""]

        #    fig.canvas.set_window_title('Distance Error Correlation')
        plt.xticks(TOPICs, LABELs,fontsize=18)

        #    legend([rects1,rects2], [r'\LARGE\textbf{2D2D}', r'\LARGE\textbf{2D3D}'], loc='lower right') 

        TOPICS = [0.0, 0.5,1,1.5,2,2.5,3,3.5,4,4.5,5]#,110]#,120]
        print TOPICS
        LABELS = [r'0.0', r'0.5',  r'1',r'1.5', r'2',r'2.5', r'3',r'3.5', r'4',r'4.5',r'5']#, ""]#, ""]

        #    fig.canvas.set_window_title('Accuracy - Activity Statistics')
        plt.yticks(TOPICS, LABELS,fontsize=18)

        def autolabel(rects):
            # attach some text labels
            for rect in rects:
                height = rect.get_height()
                ax.text(0.26+rect.get_x()+rect.get_width()/2., height +0.35, "%.2f"%float(height),
                ha='center', va='bottom',fontweight='bold',fontsize=13.5)

            #    autolabel(rects1)


        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()
        ######################################################################################################


def main():
    ex = Parallax2Dto3DMapping()
    ex.performExperiment()
        
if __name__ == "__main__":
    main()