#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
This script run neural network model on a camera live stream
"""

import argparse
import cv2
import numpy as np
import os
import time
import sys


COMMANDS = {0: "move_forward", 1: "go_down", 2: "rot_10_deg",
            3: "go_up", 4: "take_off", 5: "land", 6: "idle"}


def send_command(anafi, command_id):
    """
    Function to send commands to an Anafi drone in function of the command id
    """
    if command_id not in COMMANDS:
        raise f"Command id not in COMMANDS choices: {command_id}"

    print("The following command will be sent: ", COMMANDS[command_id])

    if COMMANDS[command_id] == "move_forward":
        anafi.move_relative(dx=1, dy=0, dz=0, dradians=0)
    if COMMANDS[command_id] == "go_down":
        anafi.move_relative(dx=0, dy=0, dz=-0.5, dradians=0)
    if COMMANDS[command_id] == "rot_10_deg":
        anafi.move_relative(dx=0, dy=0, dz=0, dradians=0.785)
    if COMMANDS[command_id] == "go_up":
        anafi.move_relative(dx=0, dy=0, dz=0.5, dradians=0)
    if COMMANDS[command_id] == "take_off":
        anafi.safe_takeoff(5)
    if COMMANDS[command_id] == "land":
        anafi.safe_land(5)
    return


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-p",
        "--weight_path",
        required=True,
        type=str,
        help="Path to load weights for the model."
    )
    parser.add_argument(
        "-a",
        "--pyparrot_path",
        required=True,
        type=str,
        help="Path to pyparrot module downloaded from amymcgovern on github."
    )
    parser.add_argument(
        "-w",
        "--img_width",
        required=False,
        default=28,
        type=int,
        help="Image width."
    )
    parser.add_argument(
        "-n",
        "--num_classes",
        required=False,
        default=7,
        type=int,
        help="Number of classes."
    )
    parser.add_argument(
        "-c",
        "--crop",
        required=False,
        default=None,
        type=str,
        help="Crop image, format: MinWidth,MaxWidth,MinHeight,MaxHeight.\
              Set -1 for the unchanged ones"
    )
    parser.add_argument(
        "-r",
        "--resize",
        required=False,
        default=None,
        type=str,
        help="Resize shape, format: height,width"
    )
    parser.add_argument(
        "-b",
        "--binarize",
        required=False,
        default=None,
        type=str,
        help="To binarize images, format for thresholding: min,max"
    )
    parser.add_argument(
        "-g",
        "--gray",
        required=False,
        action="store_true",
        help="To save 1-channel images"
    )
    parser.add_argument(
        "-e",
        "--erode",
        required=False,
        default=None,
        type=str,
        help="Erode option, format: kernel_size,iteration"
    )
    parser.add_argument(
        "-d",
        "--dilate",
        required=False,
        default=None,
        type=str,
        help="Dilate option, format: kernel_size,iteration"
    )
    parser.add_argument(
        "-m",
        "--camid",
        required=False,
        default=0,
        type=int,
        help="Camera ID, default is 0"
    )
    parser.add_argument(
        "-t",
        "--tensorflow",
        required=False,
        action="store_true",
        help="To specify if Tensorflow model is used."
    )
    parser.add_argument(
        "-z",
        "--number_of_confimation",
        required=False,
        default=3,
        type=int,
        help="Minimum number of identical commands before sending to drone."
    )

    args = parser.parse_args()

    """
    Drone connection
    """
    sys.path.append(args.pyparrot_path)
    from pyparrot.Anafi import Anafi
    print("Connecting to drone...")
    anafi = Anafi(drone_type="Anafi", ip_address="192.168.42.1")
    success = anafi.connect(10)
    print(success)
    print("Sleeping few seconds...")
    anafi.smart_sleep(3)

    """
    Load model
    """
    print("Loading model...")
    input_size = args.img_width**2
    num_class = args.num_classes
    hidden_size = 128

    if args.tensorflow:
        import tensorflow as tf
        model = tf.keras.models.load_model(args.weight_path)
    else:
        script_path = os.path.realpath(__file__)
        sys.path.append(os.path.dirname(script_path) + "/../")
        from homemade_framework import framework as NN
        model = NN.Sequential([NN.Linear(input_size, hidden_size),
                               NN.LeakyReLU(), NN.BatchNorm(),
                               NN.Linear(hidden_size, hidden_size),
                               NN.LeakyReLU(), NN.BatchNorm(),
                               NN.Linear(hidden_size, num_class),
                               NN.Softmax()], NN.LossMSE())
        model.load(args.weight_path)

    """
    Webcam process
    """
    print("Start webcam...")
    cam = cv2.VideoCapture(args.camid)
    ret, frame = cam.read()

    min_height, max_height = 0, frame.shape[0]
    min_width, max_width = 0, frame.shape[1]
    print("Cam resolution: {}x{}".format(max_width, max_height))
    if args.crop is not None:
        res = [int(x) for x in args.crop.split(',')]
        if res[0] != -1:
            min_width = res[0]
        if res[1] != -1:
            max_width = res[1]
        if res[2] != -1:
            min_height = res[2]
        if res[3] != -1:
            max_height = res[3]
        print("Image cropped to minWidth:maxWidth, minHeight:maxHeight: {}:{}\
              , {},{}".format(min_width, max_width, min_height, max_height))
    pause = False
    imgs = []

    while True:
        ret, frame = cam.read()
        if not ret:
            print("failed to grab frame")
            break
        if args.crop is not None:
            frame = frame[min_height:max_height, min_width:max_width]
        cv2.imshow("Original image", frame)

        k = cv2.waitKey(1)
        if k % 256 == 27:
            # ESC pressed
            print("Escape hit, closing...")
            break
        elif k % 256 == ord('p'):
            # p pressed
            if pause:
                pause = False
            else:
                pause = True

        if not pause:
            if args.gray:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            if args.binarize:
                frame = cv2.medianBlur(frame, 5)
                min_thresh, max_thresh = [int(x) for x in
                                          args.binarize.split(',')]
                ret, frame = cv2.threshold(frame, min_thresh, max_thresh,
                                           cv2.THRESH_BINARY)
            if args.erode is not None:
                k_size, iteration = [int(x) for x in args.erode.split(',')]
                kernel = np.ones((k_size, k_size), np.uint8)
                frame = cv2.erode(frame, kernel, iterations=int(iteration))
            if args.dilate is not None:
                k_size, iteration = [int(x) for x in args.dilate.split(',')]
                kernel = np.ones((k_size, k_size), np.uint8)
                frame = cv2.dilate(frame, kernel, iterations=int(iteration))

            if args.resize:
                height, width = [int(size) for size in args.resize.split(',')]
                frame = cv2.resize(frame, (height, width),
                                   interpolation=cv2.INTER_AREA)

            image = np.asarray(frame)/255.
            cv2.imshow("Input image for the model", frame)
            image = image.reshape([np.prod(image.shape)])
            if len(imgs) < args.number_of_confimation:
                imgs.append(image)
            else:
                if args.tensorflow:
                    results = np.argmax(model(np.asarray(imgs)), axis=1)
                else:
                    results = NN.get_inferences(model, np.asarray(imgs))
                print("Model's output on buffer: ", results)
                if np.unique(results).size == 1 and\
                        COMMANDS[results[0]] != "idle":
                    send_command(anafi, results[0])
                    imgs = []
                imgs = imgs[1:]
                imgs.append(image)
            time.sleep(0.3)

    cam.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()

[((45, 13, 45, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((153, 4, 153, 39), 'sys.path.append', 'sys.path.append', ({(153, 20, 153, 38): 'args.pyparrot_path'}, {}), '(args.pyparrot_path)', False, 'import sys\n'), ((156, 12, 156, 64), 'pyparrot.Anafi.Anafi', 'Anafi', (), '', False, 'from pyparrot.Anafi import Anafi\n'), ((189, 10, 189, 38), 'cv2.VideoCapture', 'cv2.VideoCapture', ({(189, 27, 189, 37): 'args.camid'}, {}), '(args.camid)', False, 'import cv2\n'), ((274, 4, 274, 27), 'cv2.destroyAllWindows', 'cv2.destroyAllWindows', ({}, {}), '()', False, 'import cv2\n'), ((172, 16, 172, 60), 'tensorflow.keras.models.load_model', 'tf.keras.models.load_model', ({(172, 43, 172, 59): 'args.weight_path'}, {}), '(args.weight_path)', True, 'import tensorflow as tf\n'), ((174, 22, 174, 48), 'os.path.realpath', 'os.path.realpath', ({(174, 39, 174, 47): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((217, 8, 217, 43), 'cv2.imshow', 'cv2.imshow', ({(217, 19, 217, 35): '"""Original image"""', (217, 37, 217, 42): 'frame'}, {}), "('Original image', frame)", False, 'import cv2\n'), ((219, 12, 219, 26), 'cv2.waitKey', 'cv2.waitKey', ({(219, 24, 219, 25): '1'}, {}), '(1)', False, 'import cv2\n'), ((182, 46, 182, 58), 'homemade_framework.framework.LossMSE', 'NN.LossMSE', ({}, {}), '()', True, 'from homemade_framework import framework as NN\n'), ((255, 12, 255, 58), 'cv2.imshow', 'cv2.imshow', ({(255, 23, 255, 50): '"""Input image for the model"""', (255, 52, 255, 57): 'frame'}, {}), "('Input image for the model', frame)", False, 'import cv2\n'), ((271, 12, 271, 27), 'time.sleep', 'time.sleep', ({(271, 23, 271, 26): '(0.3)'}, {}), '(0.3)', False, 'import time\n'), ((175, 24, 175, 52), 'os.path.dirname', 'os.path.dirname', ({(175, 40, 175, 51): 'script_path'}, {}), '(script_path)', False, 'import os\n'), ((177, 31, 177, 65), 'homemade_framework.framework.Linear', 'NN.Linear', ({(177, 41, 177, 51): 'input_size', (177, 53, 177, 64): 'hidden_size'}, {}), '(input_size, hidden_size)', True, 'from homemade_framework import framework as NN\n'), ((178, 31, 178, 45), 'homemade_framework.framework.LeakyReLU', 'NN.LeakyReLU', ({}, {}), '()', True, 'from homemade_framework import framework as NN\n'), ((178, 47, 178, 61), 'homemade_framework.framework.BatchNorm', 'NN.BatchNorm', ({}, {}), '()', True, 'from homemade_framework import framework as NN\n'), ((179, 31, 179, 66), 'homemade_framework.framework.Linear', 'NN.Linear', ({(179, 41, 179, 52): 'hidden_size', (179, 54, 179, 65): 'hidden_size'}, {}), '(hidden_size, hidden_size)', True, 'from homemade_framework import framework as NN\n'), ((180, 31, 180, 45), 'homemade_framework.framework.LeakyReLU', 'NN.LeakyReLU', ({}, {}), '()', True, 'from homemade_framework import framework as NN\n'), ((180, 47, 180, 61), 'homemade_framework.framework.BatchNorm', 'NN.BatchNorm', ({}, {}), '()', True, 'from homemade_framework import framework as NN\n'), ((181, 31, 181, 64), 'homemade_framework.framework.Linear', 'NN.Linear', ({(181, 41, 181, 52): 'hidden_size', (181, 54, 181, 63): 'num_class'}, {}), '(hidden_size, num_class)', True, 'from homemade_framework import framework as NN\n'), ((182, 31, 182, 43), 'homemade_framework.framework.Softmax', 'NN.Softmax', ({}, {}), '()', True, 'from homemade_framework import framework as NN\n'), ((233, 24, 233, 63), 'cv2.cvtColor', 'cv2.cvtColor', ({(233, 37, 233, 42): 'frame', (233, 44, 233, 62): 'cv2.COLOR_BGR2GRAY'}, {}), '(frame, cv2.COLOR_BGR2GRAY)', False, 'import cv2\n'), ((235, 24, 235, 48), 'cv2.medianBlur', 'cv2.medianBlur', ({(235, 39, 235, 44): 'frame', (235, 46, 235, 47): '5'}, {}), '(frame, 5)', False, 'import cv2\n'), ((238, 29, 239, 61), 'cv2.threshold', 'cv2.threshold', ({(238, 43, 238, 48): 'frame', (238, 50, 238, 60): 'min_thresh', (238, 62, 238, 72): 'max_thresh', (239, 43, 239, 60): 'cv2.THRESH_BINARY'}, {}), '(frame, min_thresh, max_thresh, cv2.THRESH_BINARY)', False, 'import cv2\n'), ((242, 25, 242, 60), 'numpy.ones', 'np.ones', ({(242, 33, 242, 49): '(k_size, k_size)', (242, 51, 242, 59): 'np.uint8'}, {}), '((k_size, k_size), np.uint8)', True, 'import numpy as np\n'), ((246, 25, 246, 60), 'numpy.ones', 'np.ones', ({(246, 33, 246, 49): '(k_size, k_size)', (246, 51, 246, 59): 'np.uint8'}, {}), '((k_size, k_size), np.uint8)', True, 'import numpy as np\n'), ((251, 24, 252, 64), 'cv2.resize', 'cv2.resize', (), '', False, 'import cv2\n'), ((254, 20, 254, 37), 'numpy.asarray', 'np.asarray', ({(254, 31, 254, 36): 'frame'}, {}), '(frame)', True, 'import numpy as np\n'), ((256, 35, 256, 55), 'numpy.prod', 'np.prod', ({(256, 43, 256, 54): 'image.shape'}, {}), '(image.shape)', True, 'import numpy as np\n'), ((263, 55, 263, 71), 'numpy.asarray', 'np.asarray', ({(263, 66, 263, 70): 'imgs'}, {}), '(imgs)', True, 'import numpy as np\n'), ((261, 46, 261, 62), 'numpy.asarray', 'np.asarray', ({(261, 57, 261, 61): 'imgs'}, {}), '(imgs)', True, 'import numpy as np\n'), ((265, 19, 265, 37), 'numpy.unique', 'np.unique', ({(265, 29, 265, 36): 'results'}, {}), '(results)', True, 'import numpy as np\n')]