import cv2
from cv_bridge import CvBridge, CvBridgeError
import rospy
from sensor_msgs.msg import Image

"""
Starter code for an image processing node. Subscribes to the downward camera
and publishes a processed cv image (i.e. colorized, overlayed with a detected line)
This node is not meant to be a full node that you will use in your final line tracking
solution, but it will teach you key concepts that you will need for line following
"""

class ImageToCV:
    def __init__(self):
        # A subscriber to the topic '/aero_downward_camera/image'. self.image_sub_cb is called when a message is recieved
        self.camera_sub = rospy.Subscriber("/aero_downward_camera/image", Image, self.image_cb)

        # A publisher which will publish a parametrization of the detected line to the topic '/image_to_cv/processed'
        self.image_pub = rospy.Publisher("/image_to_cv/processed", Image, queue_size=1)
        self.bridge = CvBridge()

    def image_cb(self, msg):
        """
        Callback function which is called when a new message of type Image is recieved by self.camera_sub.

            Args:
                - msg = ROS Image message in 8UC1 format

            Returns: NA
        """
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, "8UC1")
            self.process(cv_image)
        except CvBridgeError as e:
            print(e)

    def calculate_regression(points): # input is the result of np.argwhere
        # convert points to float

        #(see astype, https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.astype.html)
        points = points.astype(float)

        #num_points = points.shape[0]

        #print("num points: " + str(num_points))
        xs = points[:, 0]
        ys = points[:, 1]
        x_mean = np.average(xs)
        y_mean = np.average(ys)

        xy_mean = np.average(xs * ys)

        x_squared_mean = np.average(xs ** 2)

        if(x_mean ** 2 == x_squared_mean):
            m = sys.float_info.max
        else:
            m = ((x_mean * y_mean) - xy_mean) / ((x_mean ** 2) - x_squared_mean)

        b = y_mean - (m * x_mean)

        return (m,b)


    def find_inliers(m, b, shape):
        x1, y1, x2, y2 = None, None, None, None
        found_point_one = False

        for x_point in range(0, shape[0]):
            if 0 <= (m*x_point)+b <= shape[1] and not found_point_one:
                x1 = x_point
                y1 = (m*x1)+b
                found_point_one = True

            if 0 <= (m*x_point)+b <= shape[1]:
                x2 = x_point
                y2 = (m*x2)+b

        return x1, x2, y1, y2

    def process(self, img):
        """
        Takes in an img param and finds the line in it, if it exists

            Args:
                - img : OpenCV Image to process

            Publishes:
                - annotated image to the /image_to_cv/processed topic
        """

        # Get a colored copy of the image. This will be used solely to provide an annotated version
        # of the image for debuging purposes (for you to see in rqt_image_view). See slides for help here

        # Threshold the image to isolate the line and use it for your linear regression algorithm
        # (tip: check size to ensure you are detecting the LED line or a simple blip on the screen)

        # "Draw" a line if you have found one on the image

        # Publish your newly annotated, color image
        # NOTE: If you do not detect a line, or if you have decided to ignore detected lines of very small
        #       size, be sure to still publish the image. Also be sure to have some way to handle the camera
        #       seeing multiple "lines" at one time

        colored_img = img

        _, thresh_img = cv2.threshold(img, 245, 255, cv2.THRESH_BINARY)

        m, b = calculate_regression(np.argwhere(thresh_img))

        x1, x2, y1, y2 = find_inliers(m, b, thresh_img.shape)

        cv2.line(colored_img, (x1, y1), (x2, y2), color='lime', thickness = 3)

        self.image_pub.publish(colored_img)


if __name__ == "__main__":
    rospy.init_node("image_to_cv")
    a = ImageToCV()

    rospy.spin()