import numpy as np
import cupy as cp
import cv2
import time
import matplotlib.pyplot as plt


def line_fit(binary_warped):
    """
	Find and fit lane lines
	"""

    # Take a histogram of the bottom half of the image
    histogram = cp.sum(binary_warped[binary_warped.shape[0] // 2:, :], axis=0)

    # Creates a histogram image
    # npHist = cp.asnumpy(histogram/255)
    # plt.plot(npHist)
    # plt.title("Histogram")
    # plt.xlabel("x")
    # plt.ylabel("počet pixelů")
    # plt.show()

    # 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 = cp.int(histogram.shape[0] / 2)
    leftx_base = cp.argmax(histogram[100:midpoint]) + 100
    rightx_base = cp.argmax(histogram[midpoint:-100]) + midpoint
    t1 = time.time()
    # Choose the number of sliding windows
    nwindows = 4
    # Set height of windows
    window_height = cp.int(binary_warped.shape[0] / nwindows)
    # Identify the x and y positions of all nonzero pixels in the image
    nonzero = binary_warped.nonzero()
    nonzeroy = cp.array(nonzero[0])
    nonzerox = cp.array(nonzero[1])
    # Set the width of the windows +/- margin
    margin = 50
    # Set minimum number of pixels found to recenter window
    minpix = 20
    # 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_base - margin
        win_xleft_high = leftx_base + margin
        win_xright_low = rightx_base - margin
        win_xright_high = rightx_base + margin

        # Draw the windows on the visualization image
        # 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_base = cp.int(cp.mean(nonzerox[good_left_inds]))
        if len(good_right_inds) > minpix:
            rightx_base = cp.int(cp.mean(nonzerox[good_right_inds]))

    # Concatenate the arrays of indices
    left_lane_inds = cp.concatenate(left_lane_inds)
    right_lane_inds = cp.concatenate(right_lane_inds)

    # 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]

    try:
        # Fit a second order polynomial to each
        left_fit = cp.polyfit(lefty, leftx, 2)
        right_fit = cp.polyfit(righty, rightx, 2)
    except:
        print("An exception occurred")
        left_fit = cp.array([0, 1.0e-02, 2.5e+02])
        right_fit = cp.array([0, -1.0e-02, 2.5e+02])

    # Return a dict of relevant variables
    ret = {}
    ret['left_fit'] = left_fit
    ret['right_fit'] = right_fit

    return ret


def final_viz(image, left_fit, right_fit, m_inv):
    """
	Final lane line prediction visualized and overlayed on top of original image
	"""
    images1 = image.shape[1]
    images0 = image.shape[0]
    # Generate x and y values for plotting
    ploty = cp.linspace(0, images1 - 1, images0)
    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]

    # Create an image to draw the lines on

    color_warp = np.zeros((720, 1280, 3), dtype='uint8')

    # Recast the x and y points into usable format for cv2.fillPoly()
    pts_left = cp.array([cp.transpose(cp.vstack([left_fitx, ploty]))])
    pts_right = cp.array([cp.flipud(cp.transpose(cp.vstack([right_fitx, ploty])))])
    pts = cp.hstack((pts_left, pts_right))
    pts = cp.asnumpy(pts)

    # Draw the lane onto the warped blank image
    cv2.fillPoly(color_warp, np.int_([pts]), (0, 255, 0))
    cv2.polylines(color_warp, np.int_([pts]), False, (0, 0, 255), 5)
    m_inv = cp.asnumpy(m_inv)

    # Warp the blank back to original image space using inverse perspective matrix (Minv)
    newwarp = cv2.warpPerspective(color_warp, m_inv, (images1, images0))
    # Combine the result with the original image
    result = cv2.addWeighted(image, 1, newwarp, 0.3, 0)

    return result, newwarp
