from styx_msgs.msg import TrafficLight

import tensorflow as tf
import numpy as np
import cv2
from functools import partial
import os
import sys

THRESHOLD = 0.50


class TLClassifierCV(object):
    def __init__(self):
        self.red_bounds = ([0,100,100],[10,255,255])
        self.yellow_bounds = ([18,150,66],[45,255,255])
        self.green_bounds = ([50,128,20],[90,255,255])

    def filter_range(self, img, lower, upper):
        filtered = np.copy(img)
        mask = cv2.inRange(filtered, lower, upper )
        return cv2.bitwise_and(filtered,filtered,mask=mask)

    def remove_noise(self, img):
        denoised_image = np.copy(img)
        #return cv2.fastNlMeansDenoisingColored(denoised_image,None,12,12,7,21)
        return cv2.medianBlur(denoised_image, 3)

    def find_circles(self,img):
        return cv2.HoughCircles(img,cv2.HOUGH_GRADIENT,1, 20,
              param1=30,
              param2=10,
              minRadius=10,
              maxRadius=40)

    def local_normalize_histogram(self, img):

        # CLAHE (Contrast Limited Adaptive Histogram Equalization)
        clahe = cv2.createCLAHE(clipLimit=3., tileGridSize=(8,8))

        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)  # convert from BGR to LAB color space
        l, a, b = cv2.split(lab)  # split on 3 different channels

        l2 = clahe.apply(l)  # apply CLAHE to the L-channel

        lab = cv2.merge((l2,a,b))  # merge channels
        return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)  # convert from LAB to BGR

    def has_color(self, img_denoise, color='Red'):

        if (color == 'Green'):
            # Filter for green
            img_thresh = self.filter_range(img_denoise,np.array(self.green_bounds[0]),np.array(self.green_bounds[1]))
        elif(color == 'Yellow'):
            # Filter for yellow
            img_thresh = self.filter_range(img_denoise,np.array(self.yellow_bounds[0]),np.array(self.yellow_bounds[1]))
        else:
            # Filter for red
            img_thresh = self.filter_range(img_denoise,np.array(self.red_bounds[0]),np.array(self.red_bounds[1]))


        #self.show_image(img_thresh, title ="filtered {}".format(color))

        # Apply Hough circles
        gray = cv2.cvtColor(np.copy(img_thresh),cv2.COLOR_BGR2GRAY)
        #self.show_image(gray, title ='gray filter')

        circles = self.find_circles(gray)
        #self.show_circles(img_denoise,circles)

        return circles is not None

    def show_image(self,img, title='image'):
        cv2.imshow(title,img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def show_circles(self,img,circles):
        if circles is not None:
        	# convert the (x, y) coordinates and radius of the circles to integers
        	circles = np.round(circles[0, :]).astype("int")

        	# loop over the (x, y) coordinates and radius of the circles
        	for (x, y, r) in circles:
        		# draw the circle in the output image, then draw a rectangle
        		# corresponding to the center of the circle
        		cv2.circle(img, (x, y), r, (0, 255, 0), 4)
        		cv2.rectangle(img, (x - 5, y - 5), (x + 5, y + 5), (0, 128, 255), -1)

        	# show the output image
        	cv2.imshow("output", img)
        	cv2.waitKey(0)

    def get_classification(self, image):

        #self.show_image(img, title='original')

        #normalize hist
        img_norm = self.local_normalize_histogram(image);
        #self.show_image(img_norm, title='normalized')

        # Convert to hsv_image
        img_hsv = cv2.cvtColor(np.copy(img_norm), cv2.COLOR_BGR2HSV)
        #self.show_image(img_hsv, title ='hsv')

        # Remove potential noise
        img_denoise = self.remove_noise(img_hsv)
        #self.show_image(img_denoise, title ='denoise')

        has_red = self.has_color(img_denoise,color='Red')

        has_green = self.has_color(img_denoise,color='Green')

        has_yellow = self.has_color(img_denoise,color='Yellow')

        if(has_red):
            return TrafficLight.RED
        elif(has_yellow):
            return TrafficLight.YELLOW
        elif(has_green):
            return TrafficLight.GREEN
        else:
            return TrafficLight.UNKNOWN


        #print("Light: R:{} G:{} Y:{}".format(has_red,has_green,has_yellow))



class TLClassifier(object):
    def __init__(self, model_path):
        self.tf_session = None
        self.predict = None
        self.model_path = model_path
        self.labels = [TrafficLight.UNKNOWN, TrafficLight.RED, TrafficLight.YELLOW, TrafficLight.GREEN, TrafficLight.UNKNOWN]
        self.readsize = 1024
        PATH_TO_CKPT = self.model_path + '/checkpoints/frozen_inference_graph.pb'
        # set up tensorflow and traffic light classifier
        if self.tf_session is None:
            # get the traffic light classifier
            self.config = tf.ConfigProto(log_device_placement=True)
            self.config.gpu_options.per_process_gpu_memory_fraction = 0.6  
            self.config.operation_timeout_in_ms = 30000 # terminate anything that don't return in 50 seconds
            self.tf_graph = tf.Graph()
            # load the model into memory
            with self.tf_graph.as_default():
                od_graph_def = tf.GraphDef()
                with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
                    serialized_graph = fid.read()
                    od_graph_def.ParseFromString(serialized_graph)
                    tf.import_graph_def(od_graph_def, name='')


                    self.tf_session = tf.Session(graph=self.tf_graph, config=self.config)
                    self.image_tensor = self.tf_graph.get_tensor_by_name('image_tensor:0') #define input
                    #define outputs:			
                    self.detection_scores = self.tf_graph.get_tensor_by_name('detection_scores:0')  
                    self.detection_classes = self.tf_graph.get_tensor_by_name('detection_classes:0')
                    self.num_detections = self.tf_graph.get_tensor_by_name('num_detections:0')
                    self.predict = True
        

    def get_classification(self, image_np):
        """Determines the color of the traffic light in the image

        Args:
            image_np (cv::Mat): image containing the traffic light

        Returns:
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)

        """

        predict = TrafficLight.UNKNOWN
        if self.predict is not None:
            # Expand dimensions to this shape: [1, None, None, 3]
            image_np_expanded = np.expand_dims(image_np, axis=0)
            # Perform the detection
            (scores, classes, num) = self.tf_session.run(
                [self.detection_scores, self.detection_classes, self.num_detections],
                feed_dict={self.image_tensor: image_np_expanded})

            # Visualization of the results of a detection.
            scores = np.squeeze(scores)
            classes = np.squeeze(classes).astype(np.int32)

            # calculate prediction
            index_prediction = 4  # TrafficLight.UNKNOWN
            predict = self.labels[index_prediction]
            class_predicted = classes[0]
            confidence = scores[0]
            
            if class_predicted > 0 and class_predicted < 4 and confidence is not None and confidence > THRESHOLD:
                predict = self.labels[class_predicted]
        return predict

