# Define a class to represent the camera
import numpy as np
import cv2
import glob
import matplotlib
import matplotlib.pyplot as plt

class Camera():
    def __init__(self):
        self.nx = 9
        self.ny = 6
        # Arrays to store object points and image points from all the images.
        self.objpoints = [] # 3d points in real world space
        self.imgpoints = [] # 2d points in image plane.
        
        
    def calibrate(self, path = '../data/camera_cal/calibration*.jpg', plot_figures = False):
        # Make a list of calibration images
        images = glob.glob(path)
        
        # prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
        objp = np.zeros((self.ny*self.nx,3), np.float32)
        objp[:,:2] = np.mgrid[0:self.nx,0:self.ny].T.reshape(-1,2)

        # Step through the list and search for chessboard corners
        i = 0
        if plot_figures == True:
            fig = plt.figure()
        row = 5
        column = 4
        gray = None
        for fname in images:
            i += 1
            if plot_figures == True:
                fig.add_subplot(row, column, i)

            print('Loading calibration images at ' + str(int(100 * i/(len(images)))) + ' %...')
            img = cv2.imread(fname)
            if plot_figures == True:
                plt.imshow(img)
                        
            gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
        
            # Find the chessboard corners
            ret, corners = cv2.findChessboardCorners(gray, (self.nx,self.ny),None)
        
            # If found, add object points, image points
            if ret == True:
                self.objpoints.append(objp)
                self.imgpoints.append(corners)
        
                #Draw and display the corners
                #img = cv2.drawChessboardCorners(img, (self.nx,self.ny), corners, ret)
                #cv2.imshow('img',img)
                #cv2.waitKey(200) 

        # Do camera calibration given object points and image points, by getting the matrix coefficients
        self.ret, self.mtx, self.dist, self.rvecs, self.tvecs = cv2.calibrateCamera(self.objpoints, self.imgpoints, gray.shape[::-1], None, None)
        
        if plot_figures == True:
            plt.show()

        #cv2.waitKey(1)
        #cv2.destroyAllWindows()

    # Given an image return it's undistorted version
    def undistort(self, img):
        return cv2.undistort(img, self.mtx, self.dist, None, self.mtx)

    def sobel_threshold(self, gray_img, orient='x', sobel_kernel=3, thresh=(0, 255)):
        # Calculate directional gradient
        if orient == 'x':
            sobel = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=sobel_kernel)
        else:
            sobel = cv2.Sobel(gray_img, cv2.CV_64F, 0, 1, ksize=sobel_kernel)
        abs_sobel = np.absolute(sobel) # Absolute x derivative to accentuate lines away from horizontal
        scaled_sobel = np.uint8(255*abs_sobel/np.max(abs_sobel))
        
        # Apply threshold
        binary = np.zeros_like(scaled_sobel)
        binary[(scaled_sobel >= thresh[0]) & (scaled_sobel <= thresh[1])] = 1
        return binary

    def mag_threshold(self, gray_image, sobel_kernel=3, thresh=(0, 255)):
        # Calculate gradient magnitude
        # Apply threshold
        sobelx = self.sobel_threshold(gray_image, 'x', sobel_kernel)
        sobely = self.sobel_threshold(gray_image, 'y', sobel_kernel)
        # Calculate the gradient magnitude
        gradmag = np.sqrt(sobelx**2 + sobely**2)
        # Rescale to 8 bit
        scale_factor = np.max(gradmag)/255 
        gradmag = (gradmag/scale_factor).astype(np.uint8) 
        # Create a binary image of ones where threshold is met, zeros otherwise
        binary_output = np.zeros_like(gradmag)
        binary_output[(gradmag >= thresh[0]) & (gradmag <= thresh[1])] = 1

        # Return the binary image
        return binary_output

    def dir_threshold(self, gray_image, sobel_kernel=3, thresh=(0, np.pi/2)):
        # Calculate gradient magnitude
        # Apply threshold
        sobelx = self.sobel_threshold(gray_image, 'x', sobel_kernel)
        sobely = self.sobel_threshold(gray_image, 'y', sobel_kernel)
        # Take the absolute value of the gradient direction, 
        # apply a threshold, and create a binary image result
        absgraddir = np.arctan2(np.absolute(sobely), np.absolute(sobelx))
        binary_output =  np.zeros_like(absgraddir)
        binary_output[(absgraddir >= thresh[0]) & (absgraddir <= thresh[1])] = 1

        # Return the binary image
        return binary_output

    # Define a function that takes an image and returns the combined color and gradient
    # thresholds and the staked image
    def binary_from_combined_thresholds(self, img, sobel_kernel = 3, thresh_color_s_channel = (0, 255), 
            thresh_sobel_x = (20, 100), thresh_dir_gradient = (0, np.pi/2), thresh_magnitude = (0, 255)):
    # Convert to HLS color space and separate the S channel
        # Note: img is the undistorted image
        cp_img = np.copy(img)
        hls = cv2.cvtColor(img, cv2.COLOR_RGB2HLS)
        s_channel = hls[:,:,2]
        l_channel = hls[:,:,1]

        # Grayscale image
        # NOTE: we already saw that standard grayscaling lost color information for the lane lines
        # Explore gradients in other colors spaces / color channels to see what might work better
        gray = cv2.cvtColor(cp_img, cv2.COLOR_RGB2GRAY)

        # Sobel x & threshold x gradient
        #sobel_binary = self.sobel_threshold(gray, 'x', sobel_kernel, thresh_sobel_x)
        sobel_binary = self.sobel_threshold(l_channel, 'x', sobel_kernel, thresh_sobel_x)

        # Threshold color channel
        s_binary = np.zeros_like(s_channel)
        s_binary[(s_channel >= thresh_color_s_channel[0]) & (s_channel <= thresh_color_s_channel[1])] = 1

        thresh_l = (130, 255)
        l_binary = np.zeros_like(l_channel)
        l_binary[(l_channel > thresh_l[0]) & (l_channel <= thresh_l[1])] = 1
        
        # Stack each channel to view their individual contributions in green and blue respectively
        # This returns a stack of the two binary images, whose components you can see as different colors
        color_binary = np.dstack(( np.zeros_like(sobel_binary), sobel_binary, s_binary)) * 255

        # gradient Direction 
        dir_binary = self.dir_threshold(gray, sobel_kernel, thresh_dir_gradient)
        mag_binary = self.mag_threshold(gray, sobel_kernel, thresh_magnitude)

        # Combine the two binary thresholds
        combined_binary = np.zeros_like(sobel_binary)
        #combined_binary[((s_binary == 1) | (sobel_binary == 1)) & l_binary == 1 &  ((dir_binary == 1) & (mag_binary == 1))] = 1
        combined_binary[((s_binary == 1) | (sobel_binary == 1)) & l_binary == 1 &  (dir_binary == 1)] = 1

        return combined_binary, color_binary
        
    # Define a function that takes an image, number of x and y points, 
    # camera matrix and distortion coefficients and return the birdeye-view
    def corners_unwarp(self, img_undistorted, corners_source, corners_destination):
        # Use the OpenCV undistort() function to remove distortion
        img_size = (img_undistorted.shape[1], img_undistorted.shape[0])
        self.perspectiveM = cv2.getPerspectiveTransform(corners_source, corners_destination)
        # Warp the image using OpenCV warpPerspective()
        warped = cv2.warpPerspective(img_undistorted, self.perspectiveM, img_size)
        # Return the resulting image and matrix
        return warped, self.perspectiveM

    
    def identify_lane_lines_once(self, binary_warped, nwindows = 9, margin = 100, minpix = 50):
        histogram = np.sum(binary_warped[binary_warped.shape[0]//2:,:], axis=0)
        # Create an output image to draw on and visualize the result
        out_img = np.dstack((binary_warped, binary_warped, binary_warped))*255
        # Find the peak of the left and right halves of the histogram
        # These will be the starting point for the left and right lines
        midpoint = np.int(histogram.shape[0]//2)
        leftx_base = np.argmax(histogram[:midpoint])
        rightx_base = np.argmax(histogram[midpoint:]) + midpoint

        # HYPERPARAMETERS
        # Choose the number of sliding windows
        #nwindows = 9
        # Set the width of the windows +/- margin
        #margin = 100
        # Set minimum number of pixels found to recenter window
        #minpix = 50

        # Set height of windows - based on nwindows above and image shape
        window_height = np.int(binary_warped.shape[0]//nwindows)
        # Identify the x and y positions of all nonzero (i.e. activated) pixels in the image
        nonzero = binary_warped.nonzero()
        nonzeroy = np.array(nonzero[0])
        nonzerox = np.array(nonzero[1])
        # Current positions to be updated later for each window in nwindows
        leftx_current = leftx_base
        rightx_current = rightx_base

        # Create empty lists to receive left and right lane pixel indices
        left_lane_inds = []
        right_lane_inds = []

        # Step through the windows one by one
        for window in range(nwindows):
            # Identify window boundaries in x and y (and right and left)
            win_y_low = binary_warped.shape[0] - (window+1)*window_height
            win_y_high = binary_warped.shape[0] - window*window_height
            win_xleft_low = leftx_current - margin
            win_xleft_high = leftx_current + margin
            win_xright_low = rightx_current - margin
            win_xright_high = rightx_current + margin
            
            # Draw the windows on the visualization image
            cv2.rectangle(out_img,(win_xleft_low,win_y_low),
            (win_xleft_high,win_y_high),(0,255,0), 2) 
            cv2.rectangle(out_img,(win_xright_low,win_y_low),
            (win_xright_high,win_y_high),(0,255,0), 2) 
            
            # Identify the nonzero pixels in x and y within the window #
            good_left_inds = ((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) & 
            (nonzerox >= win_xleft_low) &  (nonzerox < win_xleft_high)).nonzero()[0]
            good_right_inds = ((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) & 
            (nonzerox >= win_xright_low) &  (nonzerox < win_xright_high)).nonzero()[0]
            
            # Append these indices to the lists
            left_lane_inds.append(good_left_inds)
            right_lane_inds.append(good_right_inds)
            
            # If you found > minpix pixels, recenter next window on their mean position
            if len(good_left_inds) > minpix:
                leftx_current = np.int(np.mean(nonzerox[good_left_inds]))
            if len(good_right_inds) > minpix:        
                rightx_current = np.int(np.mean(nonzerox[good_right_inds]))

        # Concatenate the arrays of indices (previously was a list of lists of pixels)
        #try:
        left_lane_inds = np.concatenate(left_lane_inds)
        right_lane_inds = np.concatenate(right_lane_inds)
        #except ValueError:
            # Avoids an error if the above is not implemented fully
        #    pass

        # Extract left and right line pixel positions
        leftx = nonzerox[left_lane_inds]
        lefty = nonzeroy[left_lane_inds] 
        rightx = nonzerox[right_lane_inds]
        righty = nonzeroy[right_lane_inds]

        return leftx, lefty, rightx, righty, out_img

    def fit_polynomial(self, binary_warped, nwindows = 9, margin = 100, minpix = 50):
        # Find our lane pixels first
        leftx, lefty, rightx, righty, out_img = self.identify_lane_lines_once(binary_warped)

        # Fit a second order polynomial to each using `np.polyfit`


        # Generate x and y values for plotting
        ploty = np.linspace(0, binary_warped.shape[0]-1, binary_warped.shape[0] )
        left_detected = True
        right_detected = True
        try:
            left_fit = np.polyfit(lefty, leftx, 2)
            left_fitx = left_fit[0]*ploty**2 + left_fit[1]*ploty + left_fit[2]
        except TypeError:
            # Avoids an error if `left` and `right_fit` are still none or incorrect
            #print('The function failed to fit a line!')
            left_detected = False
            left_fit = [10, 20, 0]
            left_fitx = 1*ploty**2 + 1*ploty
        
        try:
            right_fit = np.polyfit(righty, rightx, 2)
            right_fitx = right_fit[0]*ploty**2 + right_fit[1]*ploty + right_fit[2]
        except TypeError:
            # Avoids an error if `left` and `right_fit` are still none or incorrect
            #print('The function failed to fit a line!')
            right_detected = False
            right_fit = [10, 20, 0]
            right_fitx = 1*ploty**2 + 1*ploty

        ## Visualization ##
        # Colors in the left and right lane regions
        out_img[lefty, leftx] = [255, 0, 0]
        out_img[righty, rightx] = [0, 0, 255]

        # Plots the left and right polynomials on the lane lines
        # plt.plot(left_fitx, ploty, color='yellow')
        # plt.plot(right_fitx, ploty, color='yellow')

        return left_detected, right_detected, leftx, lefty, rightx, righty, ploty, left_fit, right_fit, left_fitx, right_fitx, out_img

    def fit_poly(self, img_shape, leftx, lefty, rightx, righty):
        ### TO-DO: Fit a second order polynomial to each with np.polyfit() ###
        left_fit = np.polyfit(lefty, leftx, 2)
        right_fit = np.polyfit(righty, rightx, 2)
        # Generate x and y values for plotting
        ploty = np.linspace(0, img_shape[0]-1, img_shape[0])
        ### TO-DO: Calc both polynomials using ploty, left_fit and right_fit ###
        left_fitx = left_fit[0]*ploty**2 + left_fit[1]*ploty + left_fit[2]
        right_fitx = right_fit[0]*ploty**2 + right_fit[1]*ploty + right_fit[2]
        
        return left_fitx, right_fitx, ploty

    def search_around_poly(self, binary_warped, left_fit, right_fit, margin = 100):
        # HYPERPARAMETER
        # Choose the width of the margin around the previous polynomial to search
        # The quiz grader expects 100 here, but feel free to tune on your own!
        # margin = 100

        # Grab activated pixels
        nonzero = binary_warped.nonzero()
        nonzeroy = np.array(nonzero[0])
        nonzerox = np.array(nonzero[1])
        
        ### TO-DO: Set the area of search based on activated x-values ###
        ### within the +/- margin of our polynomial function ###
        ### Hint: consider the window areas for the similarly named variables ###
        ### in the previous quiz, but change the windows to our new search area ###
        left_lane_inds = ((nonzerox > (left_fit[0]*(nonzeroy**2) + left_fit[1]*nonzeroy + 
                        left_fit[2] - margin)) & (nonzerox < (left_fit[0]*(nonzeroy**2) + 
                        left_fit[1]*nonzeroy + left_fit[2] + margin)))
        right_lane_inds = ((nonzerox > (right_fit[0]*(nonzeroy**2) + right_fit[1]*nonzeroy + 
                        right_fit[2] - margin)) & (nonzerox < (right_fit[0]*(nonzeroy**2) + 
                        right_fit[1]*nonzeroy + right_fit[2] + margin)))
        
        # Again, extract left and right line pixel positions
        leftx = nonzerox[left_lane_inds]
        lefty = nonzeroy[left_lane_inds] 
        rightx = nonzerox[right_lane_inds]
        righty = nonzeroy[right_lane_inds]

        right_detected = True
        left_detected = True

        # Fit new polynomials
        try:
            left_fitx, right_fitx, ploty = self.fit_poly(binary_warped.shape, leftx, lefty, rightx, righty)
        except TypeError:
            # Avoids an error if `left` and `right_fit` are still none or incorrect
            #print('The function failed to fit a line!')
            right_detected = False
            left_detected = False
        ## Visualization ##
        # Create an image to draw on and an image to show the selection window
        out_img = np.dstack((binary_warped, binary_warped, binary_warped))*255
        window_img = np.zeros_like(out_img)
        # Color in left and right line pixels
        out_img[nonzeroy[left_lane_inds], nonzerox[left_lane_inds]] = [255, 0, 0]
        out_img[nonzeroy[right_lane_inds], nonzerox[right_lane_inds]] = [0, 0, 255]

        # Generate a polygon to illustrate the search window area
        # And recast the x and y points into usable format for cv2.fillPoly()
        left_line_window1 = np.array([np.transpose(np.vstack([left_fitx-margin, ploty]))])
        left_line_window2 = np.array([np.flipud(np.transpose(np.vstack([left_fitx+margin, 
                                ploty])))])
        left_line_pts = np.hstack((left_line_window1, left_line_window2))
        right_line_window1 = np.array([np.transpose(np.vstack([right_fitx-margin, ploty]))])
        right_line_window2 = np.array([np.flipud(np.transpose(np.vstack([right_fitx+margin, 
                                ploty])))])
        right_line_pts = np.hstack((right_line_window1, right_line_window2))

        # Draw the lane onto the warped blank image
        cv2.fillPoly(window_img, np.int_([left_line_pts]), (0,255, 0))
        cv2.fillPoly(window_img, np.int_([right_line_pts]), (0,255, 0))
        result = cv2.addWeighted(out_img, 1, window_img, 0.3, 0)
        
        # Plot the polynomial lines onto the image
        #plt.plot(left_fitx, ploty, color='yellow')
        #plt.plot(right_fitx, ploty, color='yellow')
        ## End visualization steps ##
    
        #return result
        return left_detected, right_detected, leftx, lefty, rightx, righty, ploty, left_fit, right_fit, left_fitx, right_fitx, result