# 6_4.py
import mycamera
import cv2
import time
from gpiozero import DigitalOutputDevice
from gpiozero import PWMOutputDevice

# motor setting
PWMA = PWMOutputDevice(18)
AIN1 = DigitalOutputDevice(22)
AIN2 = DigitalOutputDevice(27)

PWMB = PWMOutputDevice(23)
BIN1 = DigitalOutputDevice(25)
BIN2 = DigitalOutputDevice(24)

# motor direction
def motor_go(speed):
    AIN1.value = 0
    AIN2.value = 1
    PWMA.value = speed
    BIN1.value = 0
    BIN2.value = 1
    PWMB.value = speed

def motor_back(speed):
    AIN1.value = 1
    AIN2.value = 0
    PWMA.value = speed
    BIN1.value = 1
    BIN2.value = 0
    PWMB.value = speed

def motor_left(speed):
    AIN1.value = 1
    AIN2.value = 0
    PWMA.value = speed
    BIN1.value = 0
    BIN2.value = 1
    PWMB.value = speed

def motor_right(speed):
    AIN1.value = 0
    AIN2.value = 1
    PWMA.value = speed
    BIN1.value = 1
    BIN2.value = 0
    PWMB.value = speed

def motor_stop():
    AIN1.value = 0
    AIN2.value = 1
    PWMA.value = 0.0
    BIN1.value = 0
    BIN2.value = 1
    PWMB.value = 0.0

speedSet = 0.5

def main():
    camera = mycamera.MyPiCamera(640, 480)
    filepath = '/home/pi/DeepRoad/video/train'
    i = 0
    carState = 'stop'

    while( camera.isOpened() ):
        keyValue = cv2.waitKey(10)

        # move by keyboard
        if keyValue == ord('q'):
            motor_stop()
            break
        # up => go
        elif keyValue == 82:
            print('up')
            carState = 'go'
            motor_go(speedSet)
        # down => back
        elif keyValue == 84:
            print('down')
            carState = 'back'
            motor_back(speedSet)
        # left => left
        elif keyValue == 81:
            print('left')
            carState = 'left'
            motor_left(speedSet)
        # right => right
        elif keyValue == 83:
            print('right')
            carState = 'right'
            motor_right(speedSet)

        _, image = camera.read()
        image = cv2.flip(image, -1)
        cv2.imshow('Original', image)

        # (1) ROI
        height, _, _ = image.shape
        save_image = image[int(height/2):, :, :]
        # (2) RGB(BGR) -> YUV
        save_image = cv2.cvtColor(save_image, cv2.COLOR_BGR2YUV)
        # (3) blur
        save_image = cv2.GaussianBlur(save_image, (3, 3), 0)
        # (4) resize for nvidia model
        save_image = cv2.resize(save_image, (200, 66))
        cv2.imshow('Save', save_image)

        # save images
        if carState == 'left':
            cv2.imwrite('%s_%05d_%03d.png' % (filepath, i, 45), save_image)
            i += 1
        elif carState == 'right':
            cv2.imwrite('%s_%05d_%03d.png' % (filepath, i, 135), save_image)
            i += 1
        elif carState == 'go':
            cv2.imwrite('%s_%05d_%03d.png' % (filepath, i, 90), save_image)
            i += 1

    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()