import cv2
import numpy as np

import persp_transform


def draw_on_orig(binary_warped, undistorted, leftx, lefty, rightx, righty):
    # Fit a second order polynomial to each
    left_fit = np.polyfit(lefty, leftx, 2)
    right_fit = np.polyfit(righty, rightx, 2)
    # LeftX Base = Bottom
    # RightX Base = Bottom
    ploty = np.linspace(0, 719, num=720)  # to cover same y-range as image

    # Fit a second order polynomial to pixel positions in each fake lane line
    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]
    # Calculate the position of the vehicle
    rightx_int = right_fit[0] * 720 ** 2 + right_fit[1] * 720 + right_fit[2]
    leftx_int = left_fit[0] * 720 ** 2 + left_fit[1] * 720 + left_fit[2]
    center = abs(640 - ((rightx_int+leftx_int)/2))
    # Create an image to draw the lines on
    warp_zero = np.zeros_like(binary_warped).astype(np.uint8)
    color_warp = np.dstack((warp_zero, warp_zero, warp_zero))

    # Recast the x and y points into usable format for cv2.fillPoly()
    pts_left = np.array([np.transpose(np.vstack([left_fitx, ploty]))])
    pts_right = np.array([np.flipud(np.transpose(np.vstack([right_fitx, ploty])))])
    pts = np.hstack((pts_left, pts_right))
    # Draw the lane onto the warped blank image
    cv2.fillPoly(color_warp, np.int_([pts]), (0, 255, 0))
    # Warp the blank back to original image space using inverse perspective matrix (Minv)
    newwarp = persp_transform.get_warped_perspective(color_warp, reverse_persp=True)
    # Combine the result with the original image
    final = cv2.addWeighted(undistorted, 1, newwarp, 0.3, 0)
    # cv2.imwrite('final_ud.png', final)
    return final, center
