import cv2
import numpy as np
import gradio as gr
import tempfile

# Define Utility Functions
def region_of_interest(img, vertices):
    """Select the region of interest (ROI) from a defined list of vertices."""
    mask = np.zeros_like(img)
    if len(img.shape) > 2:
        channel_count = img.shape[2]  # 3 or 4 depending on your image.
        ignore_mask_color = (255,) * channel_count
    else:
        ignore_mask_color = 255
    cv2.fillPoly(mask, [vertices], ignore_mask_color)
    masked_image = cv2.bitwise_and(img, mask)
    return masked_image

def draw_lines(img, lines, color=[255, 0, 0], thickness=2):
    """Utility for drawing lines."""
    if lines is not None:
        for line in lines:
            for x1,y1,x2,y2 in line:
                cv2.line(img, (x1, y1), (x2, y2), color, thickness)

def hough_lines(img, rho, theta, threshold, min_line_len, max_line_gap):
    """Utility for defining Line Segments."""
    lines = cv2.HoughLinesP(
        img, rho, theta, threshold, np.array([]),
        minLineLength=min_line_len, maxLineGap=max_line_gap)
    line_img = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)
    draw_lines(line_img, lines)
    return line_img, lines

def separate_left_right_lines(lines):
    """Separate left and right lines depending on the slope."""
    left_lines = []
    right_lines = []
    if lines is not None:
        for line in lines:
            for x1, y1, x2, y2 in line:
                if x1 == x2:
                    continue  # Skip vertical lines to avoid division by zero
                slope = (y2 - y1) / (x2 - x1)
                if slope < 0:  # Negative slope = left lane
                    left_lines.append([x1, y1, x2, y2])
                else:  # Positive slope = right lane
                    right_lines.append([x1, y1, x2, y2])
    return left_lines, right_lines

def cal_avg(values):
    """Calculate average value."""
    if values is not None and len(values) > 0:
        return sum(values) / len(values)
    else:
        return 0

def extrapolate_lines(lines, upper_border, lower_border):
    """Extrapolate lines keeping in mind the lower and upper border intersections."""
    slopes = []
    consts = []
    if lines is not None and len(lines) != 0:
        for x1, y1, x2, y2 in lines:
            if x1 == x2:
                continue  # Avoid division by zero
            slope = (y1 - y2) / (x1 - x2)
            slopes.append(slope)
            c = y1 - slope * x1
            consts.append(c)
        avg_slope = cal_avg(slopes)
        avg_consts = cal_avg(consts)
        if avg_slope == 0:
            return None
        x_lane_lower_point = int((lower_border - avg_consts) / avg_slope)
        x_lane_upper_point = int((upper_border - avg_consts) / avg_slope)
        return [x_lane_lower_point, lower_border, x_lane_upper_point, upper_border]
    else:
        return None

def extrapolated_lane_image(img, lines, roi_upper_border, roi_lower_border):
    """Main function called to get the final lane lines."""
    lanes_img = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)
    lines_left, lines_right = separate_left_right_lines(lines)
    lane_left = extrapolate_lines(lines_left, roi_upper_border, roi_lower_border)
    lane_right = extrapolate_lines(lines_right, roi_upper_border, roi_lower_border)
    if lane_left is not None and lane_right is not None:
        draw_con(lanes_img, [[lane_left], [lane_right]])
    return lanes_img

def draw_con(img, lines):
    """Fill in lane area."""
    points = []
    if lines is not None:
        for x1, y1, x2, y2 in lines[0]:
            points.append([x1, y1])
            points.append([x2, y2])
        for x1, y1, x2, y2 in lines[1]:
            points.append([x2, y2])
            points.append([x1, y1])
    if points:
        points = np.array([points], dtype='int32')
        cv2.fillPoly(img, [points], (0, 255, 0))

def process_image(image, low_threshold, high_threshold, kernel_size, rho, theta, hough_threshold, min_line_len, max_line_gap, roi_vertices):  
    # Convert to grayscale.
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    
    # Intensity selection.
    gray_select = cv2.inRange(gray, 150, 255)
    
    # Region masking: Select vertices according to the input image.
    roi_vertices_np = np.array(roi_vertices, dtype=np.int32)
    gray_select_roi = region_of_interest(gray_select, roi_vertices_np)
    
    # Canny Edge Detection.
    img_canny = cv2.Canny(gray_select_roi, low_threshold, high_threshold)
    
    # Remove noise using Gaussian blur.
    if kernel_size % 2 == 0:
        kernel_size += 1  # kernel_size must be odd
    canny_blur = cv2.GaussianBlur(img_canny, (kernel_size, kernel_size), 0)
    
    # Hough transform parameters set according to the input image.
    hough, lines = hough_lines(canny_blur, rho, theta, hough_threshold, min_line_len, max_line_gap)
    
    # Extrapolate lanes.
    roi_upper_border = min([vertex[1] for vertex in roi_vertices])  # smallest y value
    roi_lower_border = max([vertex[1] for vertex in roi_vertices])  # largest y value
    lane_img = extrapolated_lane_image(image, lines, roi_upper_border, roi_lower_border)
    
    # Combined using weighted image.
    image_result = cv2.addWeighted(image, 1, lane_img, 0.4, 0.0)
    return image_result

def process_video(input_video_path, low_threshold, high_threshold, kernel_size, rho, theta, hough_threshold, min_line_len, max_line_gap, roi_vertices):
    # Initialize video capture
    video_cap = cv2.VideoCapture(input_video_path)
    if not video_cap.isOpened():
        raise Exception("Error opening video stream or file")
    
    # Retrieve video frame properties.
    frame_w   = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_h   = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    frame_fps = video_cap.get(cv2.CAP_PROP_FPS)
    
    # Output video file
    temp_video = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
    output_video_path = temp_video.name
    
    # Video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    vid_out = cv2.VideoWriter(output_video_path, fourcc, frame_fps, (frame_w,frame_h))
    
    # Process each frame
    while True:
        ret, frame = video_cap.read()
        if not ret:
            break
        
        result = process_image(frame, low_threshold, high_threshold, kernel_size, rho, theta, hough_threshold, min_line_len, max_line_gap, roi_vertices)
        vid_out.write(result)
    
    # Release resources
    video_cap.release()
    vid_out.release()
    
    return output_video_path

def gradio_process_video(input_video_path, low_threshold, high_threshold, kernel_size, rho, theta_degree, hough_threshold, min_line_len, max_line_gap,
                         x1, y1, x2, y2, x3, y3, x4, y4):
    # Define ROI vertices from user input
    roi_vertices = [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]
    
    # Convert theta_degree to radians
    theta = np.deg2rad(theta_degree)
    
    # Process the video
    output_video_path = process_video(input_video_path, low_threshold, high_threshold, kernel_size, rho, theta, hough_threshold, min_line_len, max_line_gap, roi_vertices)
    
    return output_video_path

# Create the Gradio interface
iface = gr.Interface(
    fn=gradio_process_video,
    inputs=[
        gr.Video(label="Input Video"),
        gr.Slider(0, 255, step=1, value=50, label="Canny Low Threshold"),
        gr.Slider(0, 255, step=1, value=150, label="Canny High Threshold"),
        gr.Slider(1, 31, step=2, value=5, label="Gaussian Kernel Size"),
        gr.Slider(1, 10, step=1, value=1, label="Hough Transform Rho"),
        gr.Slider(0, 180, step=1, value=1, label="Hough Transform Theta (degrees)"),
        gr.Slider(1, 500, step=1, value=100, label="Hough Transform Threshold"),
        gr.Slider(1, 500, step=1, value=50, label="Hough Transform Min Line Length"),
        gr.Slider(1, 500, step=1, value=300, label="Hough Transform Max Line Gap"),
        gr.Number(value=100, label="ROI x1"),
        gr.Number(value=540, label="ROI y1"),
        gr.Number(value=900, label="ROI x2"),
        gr.Number(value=540, label="ROI y2"),
        gr.Number(value=525, label="ROI x3"),
        gr.Number(value=330, label="ROI y3"),
        gr.Number(value=440, label="ROI x4"),
        gr.Number(value=330, label="ROI y4"),
    ],
    outputs=gr.Video(label="Processed Video"),
    title="Lane Detection Video Processing",
    description="Upload a video and adjust parameters to process the video for lane detection.",
)

# Launch the interface
iface.launch()