#!/usr/bin/python3

import sys
import copy
import rospy
import cv2
import math
import numpy as np
import tf2_geometry_msgs
import tf2_ros
from sensor_msgs.msg import Image
from geometry_msgs.msg import PointStamped, Vector3, Pose
from cv_bridge import CvBridge, CvBridgeError
from visualization_msgs.msg import Marker, MarkerArray
from std_msgs.msg import ColorRGBA
import message_filters
from object_detection.msg import RingDetected

from sklearn import neural_network
import joblib
import pathlib


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

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

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

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

        # Publish coordinates of rings
        self.detection_publisher = rospy.Publisher('ring_detection', RingDetected);

        self.mlpClf = joblib.load("./src/color_model/Models/MLPRGB.pkl")

        self.rgb_subscriber = message_filters.Subscriber("/camera/rgb/image_raw", Image)
        self.depth_subscriber = message_filters.Subscriber("/camera/depth/image_raw", Image)

        # Sync rbg and depth message
        self.time_sync = message_filters.TimeSynchronizer([self.rgb_subscriber, self.depth_subscriber], 10)
        self.time_sync.registerCallback(self.register_images)

        self.rgb_message = None
        self.depth_message = None


    def register_images(self, rgb_message, depth_message):   
        self.rgb_message = rgb_message
        self.depth_message = depth_message


    def publish_detection(self, e, dist, color, stamp):
        k_f = 554 # kinect focal length in pixels // TODO: in face detection it is 554

        elipse_x = self.dims[1] / 2 - e[0][0]
        elipse_y = self.dims[0] / 2 - e[0][1]

        angle_to_target = np.arctan2(elipse_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 - 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 = rospy.Time(0)

        # Get the point in the "map" coordinate system
        point_world = self.tf_buf.transform(point_s, "map")

        self.detection_publisher.publish(point_world.point.x, point_world.point.y, point_world.point.z, color)
        

    def find_ring(self):
        rospy.loginfo('New image processing started')

        if((self.rgb_message is not None) or (self.depth_message is not None)):
            # Convert the images into a OpenCV (numpy) format
            try:
                rgb_image = self.bridge.imgmsg_to_cv2(self.rgb_message, "bgr8")
            except CvBridgeError as e:
                print(e)

            try:
                
                depth_image = self.bridge.imgmsg_to_cv2(self.depth_message, "32FC1")
            except CvBridgeError as e:
                print(e)
        else:
            return 0

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

        # Color pipeline
        # gray = cv2.cvtColor(rgb_image, cv2.COLOR_BGR2GRAY)
        # img = cv2.equalizeHist(gray)

        # Depth pipeline
        img = self.depthToGrayScale(depth_image)

        # Binarize the image
        ret, thresh = cv2.threshold(img, 20, 255, cv2.THRESH_BINARY) # Changed treshlod to 150
        thresh = thresh.astype(np.uint8)
        # cv2.imshow("Image window",thresh)
        # cv2.waitKey(1)

        # Extract contours
        contours, hierarchy = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) 

        # Fit elipses to all extracted contours
        # elps = (x_centre,y_centre),(minor_axis,major_axis), angle
        elps = []
        for cnt in contours:
            #     print cnt
            #     print cnt.shape
            if cnt.shape[0] >= 20: # TODO: test parameter
                ellipse = cv2.fitEllipse(cnt)
                elps.append(ellipse)


        # Find two elipses with same centers
        candidates = []
        for n in range(len(elps)):
            for m in range(n + 1, len(elps)):
                e1 = elps[n]
                e2 = elps[m]
                dist = np.sqrt(((e1[0][0] - e2[0][0]) ** 2 + (e1[0][1] - e2[0][1]) ** 2))
                #             print dist
                if dist < 5: # TODO: test parameter
                    candidates.append((e1,e2))

        print("Found", len(candidates), "candidates for rings")

        # Extract the depth from the depth image
        candidateNumber = 1
        for c in candidates:

            # the centers of the ellipses
            # c = (elps1, elps2)radiousX
            e1 = c[0]
            e2 = c[1]

            size = (e2[1][0] + e2[1][1])/2
            center = (e2[0][1], e2[0][0])

            x1 = int(center[0] - size / 2)
            x2 = int(center[0] + size / 2)
            x_min = x1 if x1 > 0 else 0
            x_max = x2 if x2 < rgb_image.shape[0] else rgb_image.shape[0]

            y1 = int(center[1] - size / 2)
            y2 = int(center[1] + size / 2)

            y_min = y1 if y1 > 0 else 0
            y_max = y2 if y2 < rgb_image.shape[1] else rgb_image.shape[1] 

            # Detect only real 3D rings --> there should be hole in the ring
            ringPixels = self.getRingPixels(rgb_image, e1, e2)
            
            avg_ringDepth = 0
            for ringPixel in ringPixels:
                avg_ringDepth += depth_image[ringPixel[0], ringPixel[1]]
            avg_ringDepth /= len(ringPixels)
            
            depth_diff = abs(depth_image[int(center[0]), int(center[1])] - avg_ringDepth)

            # Testing
            # print("XMin:", x_min)
            # print("XMax:", x_max)
            # print("YMin:", y_min)
            # print("YMax:", y_max)

            # for ringPixel in ringPixels:
            #     print("Ring pixel: ", ringPixel)
            #     print("RGB ring color: ", rgb_image[ringPixel[0], ringPixel[1]])

            # # # NOTE: fliped x and y pos
            # rgb_image = cv2.circle(rgb_image, (int(e2[0][0]), int(e2[0][1])), 2, (0, 0, 255), -1)
            
            # for ringPixel in ringPixels:
            #     rgb_image = cv2.circle(rgb_image, (ringPixel[1], ringPixel[0]), 1, (255, 0, 0), -1)  

            print("Depth diff:", depth_diff)
            
            if math.isnan(depth_diff):
            # if depth_diff > 0.1:
                print("Candidate", candidateNumber, "is true ring.")

                # Detect color
                color = self.detectColor(rgb_image, ringPixels)
                print("It's color is", self.getStringLabel(color))

                self.publish_detection(e1, avg_ringDepth, self.getStringLabel(color), self.depth_message.header.stamp)
            
            else:
                print("Candidate", candidateNumber, "is fake ring.")
            
            candidateNumber += 1

            # Drawing the ellipses on the image
            # cv2.ellipse(rgb_image, e1, (0, 255, 0), 1)
            # cv2.ellipse(rgb_image, e2, (0, 255, 0), 2)

        # if len(candidates)>-1:
        #     cv2.imshow("Image window", rgb_image)
        #     # cv2.imshow("Image window", rgb_image[x_min:x_max,y_min:y_max])
        #     cv2.waitKey(1)


    def getRingPixels(self, image, elipseSmall, elipseLarge):
        mask_image = np.copy(image)
        mask_image[:, :, :] = 0

        # Radious of elipse is average between radiouses of both contours
        elipse = copy.deepcopy(elipseLarge)
        elipseList = list(elipse)
        elipseList[1] = tuple([(elipseLarge[1][0] + elipseSmall[1][0]) / 2, (elipseLarge[1][1] + elipseSmall[1][1]) / 2])
        elipse = tuple(elipseList)

        cv2.ellipse(mask_image, elipse, (255, 255, 255), 1) # draw mask for elipse
        # cv2.imshow("Image window", mask_image)
        # cv2.waitKey(1) 

        mask_image = mask_image[:, :, 0] + mask_image[:, :, 1] + mask_image[:, :, 2] 
        xRing, yRing = np.where(mask_image > 0) # extract elipse pixels

        ringPoints = []
        for i in range(0, xRing.shape[0]):
            ringPoints.append((xRing[i], yRing[i]))

        return ringPoints


    def detectColor(self, image, ringPixels):
        # Generate histogram
        numBins = 256
        Channel1 = []
        Channel2 = []
        Channel3 = []

        for ringPixel in ringPixels:
            channels = image[ringPixel[0], ringPixel[1]]
            Channel1.append(channels[2]) # NOTE: BGR format 
            Channel2.append(channels[1])
            Channel3.append(channels[0])

        Hist1 = np.histogram(Channel1, numBins, range=(0, 255), density=True)[0] 
        Hist2 = np.histogram(Channel2, numBins, range=(0, 255), density=True)[0]
        Hist3 = np.histogram(Channel3, numBins, range=(0, 255), density=True)[0]

        histogram = np.concatenate((Hist1, Hist2, Hist3))

        # Detect
        return self.mlpClf.predict([histogram])

    def depthToGrayScale(self, depthImage):
        cv_image_array = np.array(depthImage, dtype = np.dtype('f8'))
        cv_image_norm = cv2.normalize(cv_image_array, cv_image_array, 0, 1, cv2.NORM_MINMAX)
        cv_image_norm *= 255

        # Resize to the desired size
        # cv_image_resized = cv2.resize(cv_image_norm, self.desired_shape, interpolation = cv2.INTER_CUBIC)

        return cv_image_norm


    def getNumericalLabel(self, color):
        if color == "Black":
            return 0
        elif color == "Blue":
            return 1
        elif color == "Green":
            return 2
        elif color == "Red":
            return 3
        elif color == "White":
            return 4
        elif color == "Yellow":
            return 5

    def getStringLabel(self, numLabel):
        if numLabel == 0:
            return "Black"
        elif numLabel == 1:
            return "Blue"
        elif numLabel == 2:
            return "Green"
        elif numLabel == 3:
            return "Red"
        elif numLabel == 4:
            return "White"
        elif numLabel == 5:
            return "Yellow"  


def main():
    ring_finder = The_Ring()

    rate = rospy.Rate(1)
    while not rospy.is_shutdown():
        ring_finder.find_ring()
        rate.sleep()

    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()
