#!/usr/bin/python

## Massimiliano Patacchiola, Plymouth University 2016
#
# In this demo a NAO robot is used to learn the association between
# a teacher head pose and its own head pose. The setup requires a
# table with some objects. Stick the naomarks on the objects.
# Pressing the button on the keyboard you can teach the robot and
# control its actions. 
#
# Requirements: you need the pynaoqi (>=2.1.4) for python 2.7 (32 or 64 bit)
# You can download it from the aldebaran website. The folder containing
# pynaoqi must be copied in the ex_nao_head_imitation folder.
#
# The joints limit in degrees [radians] for the NAO's head are:
# PITCH: +29.5 [0.51] (down), -38.5 [-0.67] (up)
# YAW: -119.5 [-2.08] (right), +119.5 [2.08] (left)

import numpy as np
import cv2
import time

#The naoqi libraries
#Requires pynaoqi >= 2.1.4
#Change this path according to your pynaoqi folder location
import sys
sys.path.insert(1, "./pynaoqi-python2.7-2.1.4.13-linux64")
from naoqi import ALProxy

#from haar_cascade import haarCascade
from head_pose_estimation import CnnHeadPoseEstimator
import tensorflow as tf

#It requires the pyERA library
from pyERA.som import Som
from pyERA.hebbian import HebbianNetwork

def main():
    # The state machine has the following states:
    # [VOID, INIT, FIND, IMG, SHOW, KEY, COUNTER, NAO, WHICH, QUIT]
    #
    # VOID: Empty state, to use for test
    # INIT: It is called only once for the initialisation
    # FIND: use the hardware face and landmark detection libraries
    # IMG: Aquire an image from the camera and convert in in numpy array
    # SHOW: Print the image on screen using OpenCV
    # KEY: Check which key is pressed
    # COUNTER: Check if some counters are alive. It avoids the use of sleeps.
    # NAO: Pressing the (h) button the robot look in front of itself
    # WHICH: Pressing the (w) button is like asking to the robot to look to a object on the table
    # QUIT: Pressing (q) unsubscribe the proxy and close the script

    #Configuration Variables, adjust to taste
    DEBUG = False
    NAO_IP = "192.168.0.100"
    NAO_PORT = 9559
    VOICE_ENBLED = False #If True the robot speaks
    RECORD_VOICE = False #If True record all the sentences
    RECORD_VIDEO = False #If True record a video from the NAO camera
    STATE = "VOID" #The initial state
    #These two lists contains the landmarks ID
    #and the associated object name. Stick many naomarks
    #on different objects and be sure the robot can look at 
    #them when turning the head.
    landmark_id_list = [64, 68, 80, 85]
    landmark_name_list = ["sponge", "stapler", "book", "cup"]

    while(True):

        #Empty STATE, to be used for test
        if(STATE == "VOID"):            
            STATE = "INIT"

        # The zero state is an init phase
        # In this state all the proxies are
        # created and tNAO subscribe to the services
        #
        elif(STATE=="INIT"):
            #Init some generic variables
            #This counter allows continuing the program flow
            #without calling a sleep
            which_counter = 0 #Counter increased when WHICH is called
            which_counter_limit = 15 #Limit for the which counter
            learning_rate = 0.1

            #Init the deepgaze head pose estimator and 
            # launch the graph in a session.
            print("[STATE " + str(STATE) + "] " + "Deepgaze init" + "\n")
            sess = tf.Session()
            my_head_pose_estimator = CnnHeadPoseEstimator(sess)
            my_head_pose_estimator.allocate_yaw_variables()
            #my_head_pose_estimator.print_allocated_variables()
            my_head_pose_estimator.load_yaw_variables("./cnn_cccdd_30k")
            yaw_teacher = 0.0 #Init the teacher yaw variable
            yaw_robot = 0.0
            pitch_robot = 0.0

            #Init the Self-Organizing Maps
            print("[STATE " + str(STATE) + "] " + "Self-Organizing Maps init" + "\n")
            som_robot = Som(matrix_size=3, input_size=1) #the robot SOM is 8x8 matrix
            print("[STATE " + str(STATE) + "] " + "Loading the network from: som_nao_eight.npz" + "\n")
            #som_robot.load("som_nao_eight.npz")
            som_robot.set_unit_weights(-80.0, 0, 0); som_robot.set_unit_weights(-60.0, 0, 1); som_robot.set_unit_weights(-40.0, 0, 2); 
            som_robot.set_unit_weights(-20.0, 1, 0); som_robot.set_unit_weights(  0.0, 1, 1); som_robot.set_unit_weights(+20.0, 1, 2);
            som_robot.set_unit_weights(+40.0, 2, 0); som_robot.set_unit_weights(+60.0, 2, 1); som_robot.set_unit_weights(+80.0, 2, 2);
            #For simplicity we consider only the Yaw angle of the teacher
            #Instead of a SOM we store the values in a vector. We can have 7 discrete positions:
            #som_teacher = np.array([-80.0, -45.0, -25.0, 0.0, +25.0, +45.0, +80.0])
            som_teacher = Som(matrix_size=3, input_size=1)
            #We set the weights value manually
            som_teacher.set_unit_weights(-80.0, 0, 0); som_teacher.set_unit_weights(-60.0, 0, 1); som_teacher.set_unit_weights(-40.0, 0, 2); 
            som_teacher.set_unit_weights(-20.0, 1, 0); som_teacher.set_unit_weights(  0.0, 1, 1); som_teacher.set_unit_weights(+20.0, 1, 2);
            som_teacher.set_unit_weights(+40.0, 2, 0); som_teacher.set_unit_weights(+60.0, 2, 1); som_teacher.set_unit_weights(+80.0, 2, 2);

            #Init the Hebbian Network
            hebbian_network = HebbianNetwork("naonet")
            hebbian_network.add_node("som_robot", (3,3))
            hebbian_network.add_node("som_teacher", (3,3))
            hebbian_network.add_connection(0, 1) #connecting som_robot -> som_teacher
            hebbian_network.print_info()
          
            #Getting the nao proxies
            print("[STATE " + str(STATE) + "] " + "ALProxy init" + "\n")
            _al_motion_proxy = ALProxy("ALMotion", NAO_IP, int(NAO_PORT))
            _al_posture_proxy = ALProxy("ALRobotPosture", NAO_IP, int(NAO_PORT))
            _al_tts_proxy = ALProxy("ALTextToSpeech", NAO_IP, int(NAO_PORT))
            _al_landmark_proxy = ALProxy("ALLandMarkDetection", NAO_IP, int(NAO_PORT))
            _al_face_proxy = ALProxy("ALFaceDetection", NAO_IP, int(NAO_PORT))
            _al_memory_proxy = ALProxy("ALMemory", NAO_IP, int(NAO_PORT))
            _al_video_proxy = ALProxy("ALVideoDevice", NAO_IP, int(NAO_PORT))
            if(RECORD_VIDEO == True): _al_video_rec_proxy = ALProxy("ALVideoRecorder", NAO_IP, int(NAO_PORT))
            if(RECORD_VIDEO == True): _al_video_rec_proxy.stopRecording() #reset a dangling connection
            #_al_speechrecognition_proxy = ALProxy("ALSpeechRecognition", NAO_IP, int(NAO_PORT))

            # Subscribe to the proxies services
            print("[STATE " + str(STATE) + "] " + "ALProxy Landmark init" + "\n")
            period = 500 #in msec
            #_al_face_proxy.subscribe("Test_Face", period, 0.0 )
            _al_face_proxy.subscribe("Test_Face")
            #_al_landmark_proxy.subscribe("Test_LandMark", period, 0.0 )
            _al_landmark_proxy.subscribe("Test_LandMark")
            try:
                #"Test_Video", CameraIndex=1, Resolution=1, ColorSpace=0, Fps=5
                #CameraIndex= 0(Top), 1(Bottom)
                #Resolution= 0(160*120), 1(320*240), VGA=2(640*480), 3(1280*960)
                #ColorSpace= AL::kYuvColorSpace (index=0, channels=1), AL::kYUV422ColorSpace (index=9,channels=3),
                #AL::kRGBColorSpace RGB (index=11, channels=3), AL::kBGRColorSpace BGR (to use in OpenCV) (index=13, channels=3)
                #Fps= OV7670 VGA camera can only run at 30, 15, 10 and 5fps. The MT9M114 HD camera run from 1 to 30fps.
                #Settings for resolution 1 (320x240)
                resolution_type = 1
                fps=15
                cam_w = 320
                cam_h = 240
                #Settigns for resolution 2 (320x240)
                #resolution_type = 2
                #fps = 15
                #cam_w = 640
                #cam_h = 480
                camera_name_id = _al_video_proxy.subscribeCamera("Test_Video", 0, resolution_type, 13, fps)
                print("[STATE " + str(STATE) + "] " + "Connected to the camera with resolution: " + str(cam_w) + "x" + str(cam_h) + "\n")
            except BaseException, err: