from __future__ import division

import visual as vs
import wx 

import numpy as np
import pylab as plt

import math 
from random import random as rnd

from util import *
from geom import *

from pupil.calibrate import get_map_from_cloud
mp = None # mapping obtained from training set

display_width = 800
display_height = 600

panel_width = 224
win = vs.window(width=display_width + panel_width, height=display_height, title='Gaze Simulation')
scene = vs.display(window=win, width=display_width, height=display_height)

################################################################################
## World Axes
drawAxes(None, vs.color.white, 13, (0, 0, 0))
################################################################################
################################################################################
## Eye Representation
eye_camera_pos = (-15, 20, 0)
scene_camera_pos = (0, 20, 0)

# This is gonna be the offset distance between the two spheres which are modelling an eye
eye_ball_offset = 4

eye_outer_pos = vs.vector((-15, 20, 15))

##_diff = vs.vector(scene_camera_pos) - eye_outer_pos
##eye_outer_pos = eye_outer_pos + _diff
##eye_camera_pos = vs.vector(eye_camera_pos) + _diff
##eye_camera_pos = (eye_camera_pos.x, eye_camera_pos.y, eye_camera_pos.z)


eye_outer_radius = 5.5
eye_outer = vs.sphere(pos=eye_outer_pos,
                      radius=eye_outer_radius, 
                      color=vs.color.red)

# This direction should always be towards the target 3D point
# this is just the default value
pupil_direction = vs.vector(1, 0, 0)
pupil_position = None # computed whenever target is changed
pupil_positions = [] # 3D positions of the pupil to be projected onto the eye camera

eye_inner_radius = 3
eye_inner_pos = eye_outer_pos + pupil_direction * eye_ball_offset
eye_inner = vs.sphere(pos=eye_inner_pos,
                      radius = eye_inner_radius,
                      color=vs.color.white)


target_position = vs.vector(0, 0, 0)
def get_target_dir():
  '''
  Returns the direction of target from the eye
  '''
  return (target_position - eye_outer_pos).norm()

def adjust_eye():
  '''
  Adjusts eye in a way that it gazes at the target
  '''
  global pupil_direction, pupil_position
  pupil_direction = get_target_dir()
  eye_inner.pos = eye_outer_pos + pupil_direction * eye_ball_offset
  pupil_position = eye_inner.pos + pupil_direction * eye_inner_radius

adjust_eye()

def drawPupilPoint():
  vs.points(pos=[pupil_position], size=10, color=vs.color.blue)
################################################################################
################################################################################
## 3D Points for Training and Test
num_training = 25
num_test = 25

min_xyz = -25, -5, -110
max_xyz = 45, 65, -30

training_points = generatePoints(num_training, min_xyz, max_xyz, True, False)
test_points = generatePoints(num_test, min_xyz, max_xyz, False, True)


def drawTrainingPoints():
  vs.points(pos=training_points, size=10, color=vs.color.yellow)
def drawTestPoints():
  vs.points(pos=test_points, size=10, color=vs.color.blue)

drawTrainingPoints()
# drawTestPoints()
################################################################################

################################################################################
## Camera Object
class Camera(PinholeCamera):
  default_radius = 2.0

  def drawCameraFrame(self):  # create frame and draw its contents
    c = vs.vector(self.t)
    self.center = vs.sphere(pos=c,
                            radius=Camera.default_radius,
                            color=vs.color.green)
    self.dir = vs.arrow(pos=c,
                        axis=vs.vector(self.direction) * self.f,
                        shaftwidth=1.0)
    self.img_plane = vs.box(pos=c + self.dir.axis,
                            length=self.image_width, 
                            width=0.5,
                            height=self.image_width,
                            color=vs.color.white,
                            opacity=0.5)

  def updateShape(self):
    c = vs.vector(self.t)
    # Update camera center
    self.center.pos = c
    # Update the arrow
    self.dir.pos = c
    self.dir.axis = vs.vector(self.direction) * self.f
    # Update image plane
    self.img_plane.pos = c + self.dir.axis
    self.img_plane.length = self.image_width
    self.img_plane.height = self.image_width

# Eye Camera
eye_camera = Camera(label='Eye Camera', f = 5)
eye_camera.setR((0, 0, 0))
eye_camera.setT(v(eye_camera_pos))# + v(30, 0, 0))
# eye_camera.setDir((0, 0, 1))
eye_camera.drawCameraFrame()
# Scene Camera
scene_camera = Camera(label='Scene Camera', f = 25, fov=math.pi/3)
scene_camera.setT(v(scene_camera_pos) + v(10, 5, -5))
scene_camera.setR((0, np.pi, 0))
# scene_camera.setDir((0, 0, -1))
scene_camera.drawCameraFrame()

# Cast rays from scene camera towards training points
for point in training_points:
  diff = vs.vector(point) - vs.vector(scene_camera.t)
  drawLine(None, vs.vector(scene_camera.t), diff.mag, diff.norm())

################################################################################
## Widgets
pan = win.panel   # addr of wx window object 
pan.SetSize((display_width + panel_width, display_height))

pan_minx = display_width + 7
wx.StaticText(pan, pos=(pan_minx, 7),
    label = "Select an item..." )

def hReset(evt):
  pass

c = 0
def gazeAtNextPoint(evt):
  global c, target_position, pupil_positions
  drawTrainingPoints()
  if c >= num_training:
    print 'Processed all training points...'
    return

  print 'Shifted gaze to point %d = %s...' % (c, training_points[c])

  vs.points(pos=[training_points[c]], size=10, color=vs.color.red)

  target_position = vs.vector(training_points[c])
  c+=1
  adjust_eye()
  # drawPupilPoint()
  pupil_positions.append(pupil_position)

def computeMap(evt):
  global mp, c, pupil_positions
  # Processing Projections
  targets = map(lambda p:vs.vector(p), training_points)
  targets = map(lambda p: scene_camera.project((p.x, p.y, p.z, 1)), targets)
  targets = map(lambda p: (-p[0], p[1]), targets) # x values are in a different CS, hence there are inverted
  # TODO: instead of inverting which only holds for this specific setting, rotate using scene_camera.R
  
  pupil_locations = map(lambda p: eye_camera.project((p.x, p.y, p.z, 1)), pupil_positions)
  # eye_camera is using the same world CS, so no reversing is required

  # filtering out points outside of the image plane
  old = len(targets)
  invalid = []
  for i in xrange(len(targets)):
    p = targets[i]
    if abs(p[0])>scene_camera.image_width/2. or abs(p[1])>scene_camera.image_width/2.:
      invalid.append(i)
  for i in invalid[::-1]:
    targets = targets[:i] + targets[i+1:]
    pupil_locations = pupil_locations[:i] + pupil_locations[i+1:]
  print 'Removed %s invalid target points...' % (old-len(targets))
  # Displaying target point projections
  base = scene_camera.center.pos + scene_camera.dir.axis
  pts = map(lambda p: np.array((p[0], p[1], 0)), targets)
  vs.points(pos=map(lambda p:p+np.array([base.x, base.y, base.z]), pts), size=7, color=vs.color.red)
  print 'Eye Camera Image Width:', eye_camera.image_width
  # Displaying pupil point projections
  base = eye_camera.center.pos + eye_camera.dir.axis
  # drawLine(None, eye_camera.t, eye_camera.f * 3, eye_camera.direction, vs.color.yellow)
  pts = map(lambda p: np.array((p[0], p[1], 0)), pupil_locations)
  vs.points(pos=map(lambda p:p+np.array([base.x, base.y, base.z]), pts), size=3, color=vs.color.red)
  # normalizing points
  print 'Normalizing points...'
  targets = scene_camera.getNormalizedPts(targets)
  pupil_locations = eye_camera.getNormalizedPts(pupil_locations)
  # computing polynomial map
  print 'Computing polynomial map...'
  mp = get_map_from_cloud(np.array([(pupil_locations[i][0], pupil_locations[i][1],
                            targets[i][0], targets[i][1]) for i in xrange(len(targets))]))
  # from sklearn.neighbors import KNeighborsRegressor as knn
  # ngh = knn(n_neighbors = 1)
  # ngh.fit(pupil_locations, targets)
  # mp = ngh.predict
  print 'Successfully computed map...'

  # Converting these to numpy arrays to support nice operands
  targets = np.array(targets)
  pupil_locations = np.array(pupil_locations)

  print 'Scene Camera Image Width:', scene_camera.image_width
  plt.subplot(211)
  plt.title('Target Points')
  plt.plot(targets[:, 0], targets[:, 1], 'ro')
  plt.axis([0, 1, 0, 1])
  plt.xlabel('X')
  plt.ylabel('Y')

  labels = ['p{0}'.format(i) for i in xrange(len(targets))]
  for label, x, y in zip(labels, targets[:, 0], targets[:, 1]):
    plt.annotate(
        label, 
        xy = (x, y), xytext = (-13, 13),
        textcoords = 'offset points', ha = 'right', va = 'top',
        bbox = dict(boxstyle = 'round,pad=0.25', fc = 'yellow', alpha = 0.25),
        arrowprops = dict(arrowstyle = '->', connectionstyle = 'arc3,rad=0'))

  plt.subplot(212)
  plt.title('Pupil Points')
  plt.plot(pupil_locations[:, 0], pupil_locations[:, 1], 'bo')
  plt.axis([0, 1, 0, 1])
  plt.xlabel('X')
  plt.ylabel('Y')

  labels = ['p{0}'.format(i) for i in xrange(len(pupil_locations))]
  for label, x, y in zip(labels, pupil_locations[:, 0], pupil_locations[:, 1]):
    plt.annotate(
        label, 
        xy = (x, y), xytext = (-13, 13),
        textcoords = 'offset points', ha = 'right', va = 'top',
        bbox = dict(boxstyle = 'round,pad=0.25', fc = 'yellow', alpha = 0.25),
        arrowprops = dict(arrowstyle = '->', connectionstyle = 'arc3,rad=0'))

  c=0
  pupil_positions = []
  plt.show()

def showTestPts(evt):
  global c, target_position, pupil_positions
  drawTestPoints()
  if c >= num_test:
    print 'Processed all test points...'
    return

  vs.points(pos=[test_points[c]], size=10, color=vs.color.red)  

  print 'Shifted gaze to point %d = %s...' % (c, test_points[c])
  target_position = vs.vector(test_points[c])
  adjust_eye()
  pupil_positions.append(pupil_position)
  c+=1

def projectAndMap(evt):
  # Processing Projections
  targets = map(lambda p:vs.vector(p), test_points)
  targets = map(lambda p: scene_camera.project((p.x, p.y, p.z, 1)), targets)
  targets = map(lambda p: (-p[0], p[1]), targets) # due to orientation of SC
  pupil_locations = map(lambda p: eye_camera.project((p.x, p.y, p.z, 1)), pupil_positions)
  # filtering out points outside of the image plane
  old = len(targets)
  invalid = []
  for i in xrange(len(targets)):
    p = targets[i]
    if abs(p[0])>scene_camera.image_width/2. or abs(p[1])>scene_camera.image_width/2.:
      invalid.append(i)
  for i in invalid[::-1]:
    targets = targets[:i] + targets[i+1:]
    pupil_locations = pupil_locations[:i] + pupil_locations[i+1:]
  print 'Removed %s invalid target points...' % (old-len(targets))
  # normalizing points
  print 'Normalizing points...'
  targets = np.array(scene_camera.getNormalizedPts(targets))
  pupil_locations = np.array(eye_camera.getNormalizedPts(pupil_locations))
  # Compute mapped points
  print pupil_locations
  print 'Computing %s mapped points...' % len(pupil_locations)
  mapped_targets = np.array(map(mp, np.array(pupil_locations)))
  print 'Computed %s mapped points...' % len(mapped_targets)
  print mapped_targets
  
  C = scene_camera.center.pos
  # ssd = sum((vs.vector(mapped_targets[i]) - vs.vector(targets[i])).mag**2 for i in xrange(len(mapped_targets)))
  # Computing Angular Error (Sum of Angular Differences)

  # denormalize points and add camera plane's z to each 
  base = scene_camera.center.pos + scene_camera.dir.axis

  targets = scene_camera.getDenormalizedPts(targets)
  targets = np.array(map(lambda p: np.concatenate((p + np.array([base.x, base.y]), np.array([base.z])), axis=0), targets))

  mapped_targets = scene_camera.getDenormalizedPts(mapped_targets)
  mapped_targets = np.array(map(lambda p: np.concatenate((p + np.array([base.x, base.y]), np.array([base.z])), axis=0), mapped_targets))

  AE = list(getAngularDiff(vs.vector(targets[i]), vs.vector(mapped_targets[i]), C) for i in xrange(len(mapped_targets)))
  N = len(mapped_targets)
  AAE = sum(AE)/N
  VAR = sum((ae - AAE)**2 for ae in AE)/N
  print 'AAE = %s, Variance = %s, STD = %s' % (AAE, VAR, np.sqrt(VAR))


  vs.points(pos=mapped_targets, size=2, color=vs.color.yellow)

  plt.subplot(111)
  plt.plot(targets[:, 0], targets[:, 1], 'bo', label='Projections')
  plt.plot(mapped_targets[:, 0], mapped_targets[:, 1], 'ro', label='Mapped Estimations')

  for _mp, tp in zip(mapped_targets, targets):
    coords = zip(*(_mp, tp))
    plt.plot(coords[0], coords[1], 'k-')

  plt.legend()
  plt.show()

y = 25
Button('Reset', pan_minx, y, hReset, (panel_width - 28, 40), pan)
y += 40 + 7
Button('Gaze at Next Point', pan_minx, y, gazeAtNextPoint, (panel_width - 28, 40), pan)
y += 40 + 7
Button('Compute Map', pan_minx, y, computeMap, (panel_width - 28, 40), pan)
y += 40 + 7
Button('Generate Test Points', pan_minx, y, showTestPts, (panel_width - 28, 40), pan)
y += 40 + 7
Button('Project, Map, Plot!', pan_minx, y, projectAndMap, (panel_width - 28, 40), pan)
################################################################################