import cv2
import argparse
import numpy as np
from adafruit_motorkit import MotorKit
from imutils.video import FPS
from sys import stdout
from gpiozero import CPUTemperature

# parse args from cli
parser = argparse.ArgumentParser()
parser.add_argument('-src', '--source', required=True, type=str)
parser.add_argument('-v', '--view', action='store_true')
parser.add_argument('-d', '--debug', nargs='+')
args = vars(parser.parse_args())
src = int(args['source']) if args['source'].isdigit() else args['source']
view = args['view']
debug = 'debug' in args.keys()
db_fps = True if args['debug'] is not None and 'fps' in args['debug'] else False
db_lists = True if args['debug'] is not None and 'lists' in args['debug'] else False
db_values = True if args['debug'] is not None and 'values' in args['debug'] else False
db_temp = True if args['debug'] is not None and 'temp' in args['debug'] else False
db_pid = True if args['debug'] is not None and 'pid' in args['debug'] else False
frame = []

# set constants and variables
VIDEO_WIDTH = 640
VIDEO_HEIGHT = 480
VIDEO_FPS = 30
VIDEO_CENTER = (int(VIDEO_WIDTH/2), int(VIDEO_HEIGHT/2))
PROP_ID = {'WIDTH': 3, 'HEIGHT': 4, 'FPS': 5, 'BRIGHTNESS': 10, 'CONTRAST': 11, 'SATURATION': 12, 'EXPOSURE': 15, 'WHITEBAL': 45}
# B: 120, C: 30, S: 70, E:140, WB: 4000
BRIGHTNESS_DEFMAX = (120, 255)
CONTRAST_DEFMAX = (30, 255)
SATURATION_DEFAX = (70, 255)
EXPOSURE_DEFMAX = (140, 10000)
WHITEBAL_DEFMAX = (4000, 10000)
CONTROLS_WINDOW = 'Camera Controls'
EXIT_KEY = ord('q')
TARGET_POWER = 1.0 
LB_COLOR = np.array([120.0, 45, 100])
UB_COLOR = np.array([180.0, 255, 255])
FILTER_THRESHOLD = 0.60
CONTOUR_OUTLINE_COLOR = (0, 255, 255)
CONTOUR_CENTER_COLOR = (255, 0, 0)
CONTOUR_LINE_COLOR = (0, 255, 0)
KP = 0.0007
ERROR_THRESHOLD = 15

# init camera
try:
    print('Initializing Camera...')
    cap = cv2.VideoCapture(src)
    if view and not is_video:
        def brightness_callback(val):
            cap.set(PROP_ID['BRIGHTNESS'], val)
        def contrast_callback(val):
            cap.set(PROP_ID['CONTRAST'], val)
        def saturation_callback(val):
            cap.set(PROP_ID['SATURATION'], val)
        def exposure_callback(val):
            cap.set(PROP_ID['EXPOSURE'], val)
        def whitebal_callback(val):
            cap.set(PROP_ID['WHITEBAL'], val)

        cv2.namedWindow(CONTROLS_WINDOW)
        CTRL_WIN = np.zeros((1, 600, 3), np.uint8)
        cv2.createTrackbar('BRIGHTNESS', CONTROLS_WINDOW, BRIGHTNESS_DEFMAX[0], BRIGHTNESS_DEFMAX[1], brightness_callback)
        cv2.createTrackbar('CONTRAST', CONTROLS_WINDOW, CONTRAST_DEFMAX[0], CONTRAST_DEFMAX[1], contrast_callback)
        cv2.createTrackbar('SATURATION', CONTROLS_WINDOW, SATURATION_DEFAX[0], SATURATION_DEFAX[1], saturation_callback)
        cv2.createTrackbar('EXPOSURE', CONTROLS_WINDOW, EXPOSURE_DEFMAX[0], EXPOSURE_DEFMAX[1], exposure_callback)

    if not is_video:
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, VIDEO_HEIGHT)
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, VIDEO_WIDTH)
        cap.set(cv2.CAP_PROP_FPS, VIDEO_FPS)
        cap.set(cv2.CAP_PROP_BRIGHTNESS, BRIGHTNESS_DEFMAX[0])
        cap.set(cv2.CAP_PROP_CONTRAST, CONTRAST_DEFMAX[0])
        cap.set(cv2.CAP_PROP_SATURATION, SATURATION_DEFAX[0])
        cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, 1)
        cap.set(cv2.CAP_PROP_EXPOSURE, EXPOSURE_DEFMAX[0])
        cap.set(cv2.CAP_PROP_AUTO_WB, False)
        cap.set(cv2.CAP_PROP_WB_TEMPERATURE, WHITEBAL_DEFMAX[0])
    for prop in PROP_ID.items():
        print(f'{prop[0]}:{int(cap.get(prop[1]))}')
    def exit_vision():
        cap.release()
        cv2.destroyAllWindows()
    print('Camera initialized\n')
except Exception as e:
    print('ERROR: An error occured trying to initialize the camera')
    print(f'ERROR: {e}')

# init motors and PID
try:
    print('Initializing Motors...')
    kit = MotorKit()
    motorL = kit.motor1
    motorR = kit.motor2

    def drive_stop():
            motorL.throttle = 0
            motorR.throttle = 0

    print('Motors initialized')
except:
    print('ERROR: Unable to initalize motors')

# init lists
rect_list = []
sorted_contours = []
cord_list = []
append = cord_list.append
cpu = CPUTemperature()
print('Starting Vision...')
print('-' * 75)
try:
    while True:
        fps = FPS().start() if db_fps else None
        largest_contour_area = -1
        target_coord = (-1, -1)
        index = -1
        lat_error = 0
        control = 0
        if debug:
            output = ''

        _, frame = cap.read()
        if frame is None:
            continue
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        mask = cv2.inRange(hsv, LB_COLOR, UB_COLOR)
        unfiltered_contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        filtered_contours = [c for c in unfiltered_contours if 50 < cv2.contourArea(c)]
        filtered_contours_area = [cv2.contourArea(c) for c in unfiltered_contours if 50 < cv2.contourArea(c)]
        if len(unfiltered_contours) > 0:
            largest_contour_area = max([cv2.contourArea(c) for c in unfiltered_contours])
        filtered_contours = [c for c in filtered_contours if cv2.contourArea(c) > FILTER_THRESHOLD * largest_contour_area]
        if len(filtered_contours) > 0:
            bounding_boxes = [cv2.boundingRect(c) for c in filtered_contours]
            sorted_contours, _ = zip(*sorted(zip(filtered_contours, bounding_boxes), key=lambda b: b[1][0], reverse=False))
            sorted_contours = list(sorted_contours)
            # ((cx, cy), (w, h), theta)
            rect_list = [cv2.minAreaRect(c) for c in sorted_contours]
            for pos, rect in enumerate(rect_list):
                if view:
                    box = np.int0(cv2.boxPoints(rect))
                    cv2.drawContours(frame, [box], 0, CONTOUR_OUTLINE_COLOR, 2)
                cx = int(rect[0][0])
                cy = int(rect[0][1])
                append((cx, cy))

        if len(cord_list) == 1:
            target_coord = (cord_list[index])
            index = 0
            if view:
                cv2.line(frame, target_coord, VIDEO_CENTER, CONTOUR_LINE_COLOR, 2)
                cv2.line(frame, target_coord, target_coord, CONTOUR_CENTER_COLOR, 5)

        elif len(cord_list) > 1:
            best_center_x = min(cord_list, key=lambda xy: abs(xy[0] - VIDEO_WIDTH / 2))[0]
            index = [coord[0] for coord in cord_list].index(best_center_x)
            best_center_y = cord_list[index][1]
            target_coord = (best_center_x, best_center_y)
            if view:
                cv2.line(frame, target_coord, VIDEO_CENTER, CONTOUR_LINE_COLOR, 2)
                for coord in cord_list:
                    cv2.line(frame, coord, coord, CONTOUR_CENTER_COLOR, 5)

        if view:
            cv2.imshow('Mask', mask)
            cv2.imshow('Contour Window', frame)
            cv2.imshow(CONTROLS_WINDOW, CTRL_WIN)
            if cv2.waitKey(1) & 0xFF == EXIT_KEY:
                exit_vision()
                drive_stop()
                break

        if fps is not None:
            fps.update()
            fps.stop()
            curr_fps = fps.fps()

        filtered_contours.clear()
        sorted_contours.clear()
        rect_list.clear()
        cord_list.clear()

        if target_coord != (-1,-1) and not is_video:
            lat_error = VIDEO_CENTER[0] - target_coord[0]
            control = abs(lat_error * KP)
            if abs(lat_error) >= ERROR_THRESHOLD:
                if lat_error > 0:
                    # need to turn right
                    motorL.throttle = abs(TARGET_POWER - control)
                    motorR.throttle = TARGET_POWER
                else:
                    # need to turn left
                    motorL.throttle = TARGET_POWER
                    motorR.throttle = abs(TARGET_POWER - control)
            else:
                motorL.throttle = TARGET_POWER / 2
                motorR.throttle = TARGET_POWER / 2

        if debug:
            if db_fps:
                output += f'FPS: {curr_fps} '
            if db_lists:
                output += f'RectL: {len(rect_list)} UnfltL: {len(unfiltered_contours)} CntrL: {len(sorted_contours)} CoordL: {len(cord_list)} '
            if db_values:
                output += f'Coord: {target_coord} Index: {index} Largest: {largest_contour_area} '
            if db_temp:
                output += f'Temp: {cpu.temperature}'
            if db_pid:
                output += f'Error: {lat_error} Control: {control}'
            output += '\r'
            stdout.write(output)
except:
    print('Exiting vision...')
    drive_stop()

drive_stop()
