#!/usr/bin/env python
from __future__ import print_function
from tf.transformations import quaternion_from_euler

import sys
import rospy
import dlib
import cv2
import numpy as np
import tf2_geometry_msgs
import tf2_ros
import math
import pickle
import socket
import os
from std_msgs.msg import String
#import matplotlib.pyplot as plt
from sensor_msgs.msg import Image
from geometry_msgs.msg import PointStamped, Vector3, Pose, Twist
from cv_bridge import CvBridge, CvBridgeError
from visualization_msgs.msg import Marker, MarkerArray
from std_msgs.msg import ColorRGBA
from exercise4.msg import faces_info

class face_localizer:
    def __init__(self):
        rospy.init_node('face_localizer', anonymous=True)

        # An object we use for converting images between ROS format and OpenCV format
        self.bridge = CvBridge()

        # The function for performin HOG face detection
        #self.face_detector = dlib.get_frontal_face_detector()
        self.face_cascade = cv2.CascadeClassifier('/opt/ros/kinetic/share/OpenCV-3.3.1-dev/haarcascades/haarcascade_frontalface_default.xml')

        self.haar_scale = 1.5
        self.min_neighbors = 3
        self.haar_flags = 0

        # A help variable for holding the dimensions of the image       
        self.dims = (0, 0, 0)

        # Marker array object used for showing markers in Rviz
        self.marker_array = MarkerArray()
        self.marker_num = 1
        self.marker_array_2 = MarkerArray()
        self.marker_num_2 = 1

        # Subscribe to the image and/or depth topic
        # self.image_sub = rospy.Subscriber("/camera/rgb/image_raw", Image, self.image_callback)
       # self.depth_sub = rospy.Subscriber("/camera/depth/image_raw", Image, self.depth_callback)

        # Publiser for the visualization markers
        self.markers_pub = rospy.Publisher('face_markers', MarkerArray, queue_size=1000)
        self.markers_pub_2 = rospy.Publisher('face_goals', MarkerArray, queue_size=1000)
        self.pose_pub = rospy.Publisher('new_faces', Pose, queue_size=10)

        # Object we use for transforming between coordinate frames
        self.tf_buf = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buf)

        #for publishing faces information
        self.faces = faces_info()
        self.faces_pub = rospy.Publisher('faces_info', faces_info, queue_size=10)

        path = os.environ['XDG_RUNTIME_DIR']
        self.server_address = path + '/uds_socket'

        # Create a UDS socket
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)

        # try to connect
        print('connecting to', self.server_address)
        try:
            self.sock.connect(self.server_address)
        except socket.error, msg:
            print(msg)
            sys.exit(1)

    def call_face_recognition(self, img_face, img_features):
        # send 2 images for face / hair + length recognition
        print('Sending face')
        img_face = cv2.cvtColor(img_face, cv2.COLOR_BGR2RGB)
        img_features = cv2.cvtColor(img_features, cv2.COLOR_BGR2RGB)
        face_and_features = (img_face, img_features)
        data = pickle.dumps(face_and_features, protocol=2)
        self.sock.sendall(data)
        self.sock.send('End')

        # recieve recognition result
        result = self.sock.recv(4096)

        # extract info from result
        result = str.decode(result)

        print("RESULT", result)
        if result == '':
            info = [0 , 1 , 0]
            print("----> Empty string exception!!! BIG OOF <----")
        else: 
            info = [int(x) for x in result.split(' ')]

        print('face info :', info)

        return info

    def get_pose(self,coords,dist,stamp):
        # Calculate the position of the detected face

        k_f = 554 # kinect focal length in pixels

        x1, x2, y1, y2 = coords

        face_x = self.dims[1] / 2 - (x1+x2)/2.
        face_y = self.dims[0] / 2 - (y1+y2)/2.

        angle_to_target = np.arctan2(face_x,k_f)

        # Get the angles in the base_link relative coordinate system
        x, y = dist*np.cos(angle_to_target), dist*np.sin(angle_to_target)

        ### Define a stamped message for transformation - directly in "base_link"
        #point_s = PointStamped()
        #point_s.point.x = x
        #point_s.point.y = y
        #point_s.point.z = 0.3
        #point_s.header.frame_id = "base_link"
        #point_s.header.stamp = rospy.Time(0)

        # Define a stamped message for transformation - in the "camera rgb frame"
        point_s = PointStamped()
        point_s.point.x = -y
        point_s.point.y = 0
        point_s.point.z = x
        point_s.header.frame_id = "camera_rgb_optical_frame"
        point_s.header.stamp = stamp
       
        # Get the point in the "map" coordinate system
        try:
            point_world = self.tf_buf.transform(point_s, "map")
            add = True
            margin = 0.5
            margin1 = 0.1

            for i in range(len(detected_faces)):
                 if detected_faces[i][0] + margin > point_world.point.x and detected_faces[i][0] - margin < point_world.point.x and detected_faces[i][1] + margin > point_world.point.y and detected_faces[i][1] - margin < point_world.point.y and detected_faces[i][2] + margin > point_world.point.z and detected_faces[i][2] - margin < point_world.point.z:
                    # print("face already detected!")
                    add = False
                    break

            # Create a Pose object with the same position
            if add:
                pose = Pose()
        
                if self.faceMargin(last_detected_face,margin1,point_world) and self.faceMargin(last_detected_face2,margin1,point_world):
                   
                   pose.position.x = (point_world.point.x + last_detected_face[0] + last_detected_face2[0])/3
                   pose.position.y = (point_world.point.y + last_detected_face[1] + last_detected_face2[1])/3
                   pose.position.z = (point_world.point.z + last_detected_face[2] + last_detected_face2[2])/3
                   print("Face number ", len(detected_faces)+1," detected at: X:",pose.position.x, " Y:", pose.position.y)
                
                   detected_faces.append([pose.position.x,pose.position.y, pose.position.z])
                else:
                    pose = None
                last_detected_face2[0] = last_detected_face[0]
                last_detected_face2[1] = last_detected_face[1]
                last_detected_face2[2] = last_detected_face[2]   
                last_detected_face[0] = point_world.point.x
                last_detected_face[1] = point_world.point.y
                last_detected_face[2] = point_world.point.z
            else:
                pose = None
        except Exception as e:
            print(e)
            pose = None

        return pose
    
    def faceMargin(self,lastFace, margin, point):
        if lastFace[0] + margin > point.point.x and lastFace[0] - margin < point.point.x and lastFace[1] + margin > point.point.y and lastFace[1] - margin < point.point.y and lastFace[2] + margin> point.point.z and lastFace[2] - margin < point.point.z:
            return True
            
        return False    

    def find_faces(self):

       # print('I got a new image!')
        # Get the next rgb and depth images that are posted from the camera
        try:
            rgb_image_message = rospy.wait_for_message("/camera/rgb/image_raw", Image)
        except Exception as e:
            print(e)
            return 0
        try:
            depth_image_message = rospy.wait_for_message("/camera/depth/image_raw", Image)
        except Exception as e:
            print(e)
            return 0
        #get current velocity
        try:
            velocity = rospy.wait_for_message("/mobile_base/commands/velocity", Twist, timeout=0.05)
            if abs(velocity.angular.z) > 0.8 or abs(velocity.linear.x) > 0.3:
                print("Cannot triangulate face positions when spinning/moving this fast! [", velocity.angular.z,", ", velocity.linear.x, "]")
                return 0
        except Exception as e:
            None
        #if angular velocity is too big do not proceed with detecting faces
       
        

        # Convert the images into a OpenCV (numpy) format
        try:
            rgb_image = self.bridge.imgmsg_to_cv2(rgb_image_message, "bgr8")
        except CvBridgeError as e:
            print(e)

        try:
            depth_image = self.bridge.imgmsg_to_cv2(depth_image_message, "32FC1")
        except CvBridgeError as e:
            print(e)

        # Set the dimensions of the image
        self.dims = rgb_image.shape

        # Preprocessing
        gray = cv2.cvtColor(rgb_image, cv2.COLOR_BGR2GRAY)
        img = gray 
    
        image_scale = rgb_image.shape[1]

        #smallImage = cv2.resize(gray, (int(rgb_image.shape[1] / image_scale), int(rgb_image.shape[0] / image_scale)), interpolation=cv2.INTER_LINEAR)
        #smallImage = cv2.equalizeHist(smallImage)

        faces = self.face_cascade.detectMultiScale(img, self.haar_scale, self.min_neighbors, self.haar_flags)
        # Detect the faces in the image
        # face_rectangles = self.face_detector(rgb_image, 0)

        # For each detected face, extract the depth from the depth image
        self.markers_pub.publish(self.marker_array)
        self.markers_pub_2.publish(self.marker_array_2)
        self.faces_pub.publish(self.faces)
        for (x,y,w,h) in faces:
            #print('Faces were detected\nNumber of faces: ',len(face_rectangles))

            # The coordinates of the rectanle
            x1 = x
            x2 = x+w
            y1 = y
            y2 = y+h

            dx=int((x2-x1)/4)
            dy=int((y2-y1)/4)

            x1d= x1 + dx 
            x2d= x2 - dx
            y1d= y1 + dy
            y2d= y2 - dy
            
            # Extract region containing face
            #print(x1,x2,y1,y2)
            #print(x1d,x2d,y1d,y2d)
            face_region = rgb_image[y1:y2,x1:x2]

            # Extract region for hair 
            dh = int(h/3.5)
            dw = int(w/3.5)
            hair_region = rgb_image[y-dh:y+h+dh, x-dw:x+w+dw]
        
            # Visualize the extracted face
            #cv2.imshow("Depth window", face_region)
            #cv2.waitKey(1)

            # Find the distance to the detected face
            face_distance = float(np.nanmean(depth_image[y1d:y2d,x1d:x2d]))
            

            #print('Distance to face', face_distance)

            # Get the time that the depth image was recieved
            depth_time = depth_image_message.header.stamp
           
            # Find the location of the detected face
            pose = self.get_pose((x1,x2,y1,y2), face_distance, depth_time)
            
           
            if pose is not None:

                # Create a marker used for visualization
                self.marker_num += 1
                marker = Marker()
                marker.header.stamp = rospy.Time(0)
                marker.header.frame_id = 'map'
                marker.pose = pose
                marker.type = Marker.CUBE
                marker.action = Marker.ADD
                marker.frame_locked = False
                marker.lifetime = rospy.Duration.from_sec(10)
                marker.id = self.marker_num
                marker.scale = Vector3(0.1, 0.1, 0.1)
                marker.color = ColorRGBA(0, 1, 0, 1)
                self.marker_array.markers.append(marker)
                self.markers_pub.publish(self.marker_array)

               
                pose_2 = self.pub_front_face(pose)

                #cv2.imshow("Face", face_region)
                #cv2.waitKey(1)

                try:
                    rgb_msg = self.bridge.cv2_to_imgmsg(face_region, "bgr8")
                    
                except CvBridgeError as e:
                    print(e)

                #sending detected faces 
                self.faces.poses.append(pose)
                self.faces.robotPoses.append(pose_2)
                #TO DO: add other features
                info = self.call_face_recognition(face_region, hair_region)
                self.faces.ids.append(info[0]) 
                self.faces.hairColor.append(str(info[1]))
                self.faces.hairLength.append(str(info[2]))
                self.faces.images.append(rgb_msg)
                self.faces_pub.publish(self.faces)
                # print(self.faces)
                self.marker_num_2 += 1
                marker_2 = Marker()
                marker_2.header.stamp = rospy.Time(0)
                marker_2.header.frame_id = 'map'
                marker_2.pose = pose_2
                marker_2.type = Marker.ARROW
                marker_2.action = Marker.ADD
                marker_2.frame_locked = False
                marker_2.lifetime = rospy.Duration.from_sec(10)
                marker_2.id = self.marker_num_2
                marker_2.scale = Vector3(0.2, 0.08, 0.1)
                marker_2.color = ColorRGBA(0, 0, 1, 1)
                self.marker_array_2.markers.append(marker_2)
                self.markers_pub_2.publish(self.marker_array_2)
            
    
    def pub_front_face(self, face_pose):
        xF = face_pose.position.x
        yF = face_pose.position.y
        point_s = PointStamped()
        point_s.point.x = 0
        point_s.point.y = 0
        point_s.point.z = 0
        point_s.header.frame_id = "base_link"
        point_s.header.stamp = rospy.Time(0)

        try:
            point_world = self.tf_buf.transform(point_s, "map")
        except Exception as e:
            print(e)

        x, y = point_world.point.x, point_world.point.y

        
        dx = x - xF
        dy = y - yF
        xN = dx / math.sqrt(float(dx*dx + dy *dy))  # normalize
        yN = dy / math.sqrt(float(dx*dx + dy *dy))

        dist = 0.5

        pose = Pose()
        pose.position.x = xF + dist*xN
        pose.position.y = yF + dist*yN
        pose.position.z = 0

        angle = np.arctan2((yF-y),(xF-x))
        
        q = quaternion_from_euler(0,0,angle)

        pose.orientation.x = q[0]
        pose.orientation.y = q[1]
        pose.orientation.z = q[2]
        pose.orientation.w = q[3]
           
        # post pose
        self.pose_pub.publish(pose)
        return pose

    def depth_callback(self,data):

        try:
            depth_image = self.bridge.imgmsg_to_cv2(data, "32FC1")
        except CvBridgeError as e:
            print(e)

        # Do the necessairy conversion so we can visuzalize it in OpenCV
        
        image_1 = depth_image / np.nanmax(depth_image)
        image_1 = image_1*255
        
        image_viz = np.array(image_1, dtype=np.uint8)

        cv2.imshow("Depth window", image_viz)
        cv2.waitKey(1)

        #plt.imshow(depth_image)
        #plt.show()

def main():

        global detected_faces   
        detected_faces = []
        global last_detected_face
        last_detected_face = [0, 0, 0]

        global last_detected_face2
        last_detected_face2 = [0, 0, 0]

        face_finder = face_localizer()
           
        rate = rospy.Rate(10)
        print("face localization started!")
        while not rospy.is_shutdown():
            face_finder.find_faces()
            rate.sleep()

        cv2.destroyAllWindows()
        face_finder.sock.close()

if __name__ == '__main__':
    main()
