#!/usr/bin/python

import sys
import numpy as np
import cv2
import rospy
from sensor_msgs.msg import CompressedImage, Image
from cv_bridge import CvBridge, CvBridgeError
from std_msgs.msg import Float32
from curve_fit import *
from birds_eye_view import *
from lane_filter import *

class DetectLane(object):
    def __init__(self):

        self.img_sub = rospy.Subscriber('/raspicam_node/image/compressed', CompressedImage, self.callback)

        self.w = np.int(rospy.get_param('/raspicam_node/width') * 0.5)
        self.h = np.int(rospy.get_param('/raspicam_node/height') * 0.5)
        self.bridge = CvBridge()

        self.src = np.float32([
            (self.w * 0.25, self.h * 0.4), (self.w * 0.75, self.h * 0.4),
            (self.w * 0.9, self.h), (self.w * 0.1, self.h)
        ])

        self.dst = np.float32([
            (0, 0), (self.w, 0),
            (self.w, self.h), (0, self.h)
        ])

        self.birds_eye_view = BirdsEyeView(self.src, self.dst)
        
        self.lane_filter = LaneFilter()

        self.curve_fit = CurveFit(number_of_windows=10, margin=300, minimum_pixels=100)
        self.final_fit = None

        self.lane_error_pub = rospy.Publisher('/detect_lane/lane_error', Float32, queue_size=1)
        self.lane_error_msg = Float32()
        self.lane_error_msg.data = np.nan
        self.final_vis_img = rospy.Publisher('/detect_lane/final_vis_img', Image, queue_size=1)
        
        rospy.on_shutdown(self.shutdown_hook)

    def callback(self, data):
        try:
            cv_img = self.bridge.compressed_imgmsg_to_cv2(data, 'bgr8')
        except CvBridgeError as e:
            print(e)

        cv_img = cv2.resize(cv_img, (0, 0), fx=0.5, fy=0.5)
        warp_img = self.birds_eye_view.sky_view(cv_img)
        bin_warp_img = self.lane_filter.processed_img(warp_img)

        try:
            lane_fit = self.curve_fit.method_two(bin_warp_img, self.final_fit)
        except:
            lane_fit = self.curve_fit.sliding_windows(bin_warp_img)
        self.final_fit = lane_fit

        out_plot_img = self.curve_fit.plot()

        if self.final_fit is not None:
            lane_img = self.birds_eye_view.project_lanes(cv_img, warp_img, lane_fit)
            ploty = np.linspace(0, bin_warp_img.shape[0]-1, bin_warp_img.shape[0])
            lane_fitx = lane_fit[0] * (ploty**2) + lane_fit[1] * ploty + lane_fit[2]
            pts = np.array(np.transpose(np.vstack([lane_fitx, ploty])))
            lane_error = np.int(bin_warp_img.shape[1]/2 - pts[0][0])
            
            self.lane_error_msg.data = lane_error
        else:
            lane_img = cv_img
            self.lane_error_msg.data = np.nan
        self.lane_error_pub.publish(self.lane_error_msg)

        
        final_image_part1 = cv2.hconcat([cv_img, lane_img])
        final_image_part2 = cv2.hconcat([bin_warp_img, out_plot_img])
        final_image = cv2.vconcat([final_image_part2, final_image_part1])
        """
        try:
            self.final_vis_img.publish(self.bridge.cv2_to_imgmsg(final_image, 'bgr8'))
        except CvBridgeError as e:
            print(e)
        """
    def shutdown_hook(self):

        self.lane_error_msg.data = np.nan
        self.lane_error_pub.publish(self.lane_error_msg)

def main(args):
    rospy.init_node('detect_lane', anonymous=True)
    DetectLane()
    try:
        rospy.spin()
    except KeyboardInterrupt:
        print("Shutting down")
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main(sys.argv)



