"""
This code predicts CNN model creates contour and draw polynom
"""
import numpy as np
import cv2
from keras.models import load_model
import time


# Class to average lanes with
class Lanes():
    def __init__(self):
        self.recent_fit = []
        self.avg_fit = []


def final_viz(image, left_fit, right_fit):
    # function to draw image based on polynom parameters
    images1 = image.shape[1]
    images0 = image.shape[0]
    # Generate x and y values for plotting
    ploty = np.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]

    polynom_image = np.zeros((720, 1280, 3), dtype='uint8')  # NOTE: Hard-coded image dimensions
    # 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.polylines(polynom_image, np.int_([pts]), False, (0, 0, 255), 5)
    # cv2.fillPoly(polynom_image, np.int_([pts]), (0, 255, 0))

    return polynom_image


def findGreatesContour(contours):
    # this function pick the index of the largest contour
    largest_area = 0
    largest_contour_index = -1
    i = 0
    total_contours = len(contours)
    while i < total_contours:
        area = cv2.contourArea(contours[i])
        if area > largest_area:
            largest_area = area
            largest_contour_index = i
        i += 1

    return largest_contour_index


def polyfit_lane(image):
    # Function fits polynom to CNN model prediction
    # Height of the hood to fit the polynom
    height = 385
    # Loads and thresholds the CNN model prediction
    g_channel = image[:, :, 1]
    binary_output = np.zeros_like(g_channel)
    binary_output[(g_channel > 5) & (g_channel <= 255)] = 255
    # Finds the contours in the image and chose the largest one
    cnts, _ = cv2.findContours(binary_output, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    largest_contour_index = findGreatesContour(cnts)
    try:
        coords = cnts[largest_contour_index]
        coords = coords.tolist()
        # Divide the contour coordinates into left and right
        coords = [x[0] for x in coords if not x[0][1] < height-130 and x[0][1] < height-30]
        coords_right = [x for x in coords if not x[0] < 600]
        coords_left = [x for x in coords if x[0] < 600]

        # Empty lists of coordinates
        leftx = []
        lefty = []
        rightx = []
        righty = []

        for i in coords_left:
            leftx.append(i[0])
            lefty.append(i[1])
        for i in coords_right:
            rightx.append(i[0])
            righty.append(i[1])
        # Fit a second order polynomial to each
        left_fit = np.polyfit(lefty, leftx, 2)
        right_fit = np.polyfit(righty, rightx, 2)
    except:
        print("An exception occurred")
        left_fit = np.array([0, 1.0e-2, 6.8e+02])
        right_fit = np.array([0, -1.0e-2, 6.8e+02])

    result = final_viz(image, left_fit, right_fit)

    mask = np.zeros(result.shape[:2], dtype="uint8")
    cv2.rectangle(mask, (0, height-140), (1280, height+5), 255, -1)

    masked = cv2.bitwise_and(result, result, mask=mask)

    return masked


def road_lines(image):
    """
    Loads the road image,re-sizes for the model,
    predicts the lane adds the polynomial fit
    and merges with original image
    """

    t0 = time.time()
    # Get image ready for feeding into model
    small_img = cv2.resize(image, (320, 160))
    small_img = np.array(small_img)
    small_img = small_img[None, :, :, :]

    # Make prediction with neural network and un-normalize value by multiplying by 255
    prediction = model(small_img)[0] * 255

    # Add lane prediction to list for averaging
    lanes.recent_fit.append(prediction)
    # Only using last five for average
    if len(lanes.recent_fit) > 4:
        lanes.recent_fit = lanes.recent_fit[1:]
    # Calculate average detection
    lanes.avg_fit = np.mean(np.array([i for i in lanes.recent_fit]), axis=0)

    # Generate fake R & B color dimensions, stack with G
    blanks = np.zeros_like(lanes.avg_fit).astype(np.uint8)
    lane_drawn = np.dstack((blanks, lanes.avg_fit, blanks))

    # Re-size to match the original image
    lane_image = cv2.resize(lane_drawn, (1280, 720))
    image = cv2.resize(image, (1280, 720))
    image = np.asarray(image, np.uint8)
    lane_image = np.asarray(lane_image, np.uint8)
    # Create polynom image
    polynom_image = polyfit_lane(lane_image)

    # Merge the lane drawing onto the original image
    result = cv2.addWeighted(image, 0.6, polynom_image, 0.4, 0.0)
    result = cv2.addWeighted(result, 0.6, lane_image, 0.4, 0.0)


    return result

# Set the loop time in ms
loop_time = 500
# Load the model
model = load_model('models/model_n20.h5')
# Create lanes object
lanes = Lanes()
# Initialize camera or video file
cap = cv2.VideoCapture("video_test.mp4")
t_old = 0

while True:
    t_act = time.time() * 1000
    if t_act - t_old > loop_time:

        self_ret, image = cap.read()
        if not self_ret:
            break
        out_img = road_lines(image)
        cv2.imshow("Result", out_img)
        if t_act - t_old - 1 < loop_time:
            print("The image have been updated", t_act - t_old)
        else:
            print("ERROR the time was longer", t_act - t_old)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
        t_old = t_act

cap.release()
cv2.waitKey(0)
cv2.destroyAllWindows()
