# camera.py
#
# Sample Webots controller file for receiving images from a fixed camera.

# No copyright, 2020-2024, Garth Zeglin.  This file is
# explicitly placed in the public domain.

print("camera.py waking up.")

# Import the Webots simulator API.
from controller import Robot, Display

# Import the standard Python math library.
import math, random, time

# Import the third-party numpy library for matrix calculations.
# Note: this can be installed using 'pip3 install numpy' or 'pip3 install scipy'.
import numpy as np

# Import the third-party OpenCV library for image operations.
# Note: this can be installed using 'pip3 install opencv-python'
import cv2 as cv

# Import the vision module from the same directory.  Note: this requires OpenCV.
import vision

# Define the time step in milliseconds between controller updates.
EVENT_LOOP_DT = 50

################################################################
class Camera(Robot):
    def __init__(self):

        super(Camera, self).__init__()
        self.robot_name = self.getName()
        print("%s: controller connected." % (self.robot_name))

        # Connect to the radio emitter and receiver.
        self.emitter  = self.getDevice('emitter')
        # self.receiver = self.getDevice('receiver')
        # self.radio_interval = 200
        # self.receiver.enable(self.radio_interval)
        self.message_base_address = "/" + self.robot_name

        # Enable the camera and set the sampling frame rate.
        self.camera = self.getDevice('camera')
        self.camera_frame_dt = EVENT_LOOP_DT # units are milliseconds
        self.camera_frame_timer = 0
        self.camera.enable(self.camera_frame_dt)
        self.frame_count = 0
        self.vision_target = None

        # Connect to the display object to show debugging data within the Webots GUI.
        # The display width and height should match the camera.
        self.display = self.getDevice('display')

        return

    #================================================================
    # Polling function to process camera input.  The API doesn't seem to include
    # any method for detecting frame capture, so this just samples at the
    # expected frame rate.  It is possible it might retrieve duplicate frames or
    # experience time aliasing.

    def poll_camera(self):
        self.camera_frame_timer -= EVENT_LOOP_DT
        if self.camera_frame_timer < 0:
            self.camera_frame_timer += self.camera_frame_dt
            # read the camera
            image = self.camera.getImage()
            if image is not None:
                # print("read %d byte image from camera" % (len(image)))
                self.frame_count += 1

                # convert the image data from a raw bytestring into a 3D matrix
                # the pixel format is BGRA, which is the default for OpenCV
                width  = self.camera.getWidth()
                height = self.camera.getHeight()
                frame  = np.frombuffer(image, dtype=np.uint8).reshape(height, width, 4)

                # temporary: write image to file for offline testing
                # cv.imwrite("capture-images/%04d.png" % (self.frame_count), frame)

                # convert to grayscale and normalize levels
                frame = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)

                # extract features from the image
                render = self._find_round_image_targets(frame)

                # display a debugging image in the Webots display
                imageRef = self.display.imageNew(render.tobytes(), format=Display.RGB, width=width, height=height)
                self.display.imagePaste(imageRef, 0, 0, False)

    #================================================================
    # Analyze a grayscale image for circular targets.
    # Updates the internal state.  Returns and annotated debugging image.
    def _find_round_image_targets(self, grayscale):
        circles = vision.find_circles(grayscale)

        # generate controls from the feature data
        if circles is None:
            if self.vision_target is not None:
                # print("target not visible, tracking dropped")
                self.vision_target = None

                # send a radio message
                # emitter requires a bytestring, not a Python Unicode string
                status = self.message_base_address + "/target/dropped"
                self.emitter.send(status.encode())

        else:
            # find the offset of the first circle center from the image center in pixels,
            # ignoring the radius value
            height, width = grayscale.shape
            halfsize = np.array((width/2, height/2))
            offset = circles[0][0:2] - halfsize

            # normalize to unit scale, i.e. range of each dimension is (-1,1)
            self.vision_target = offset / halfsize

            # invert the Y coordinate so that positive values are 'up'
            self.vision_target[1] *= -1.0
            # print("target: ", self.vision_target)

            # send a radio message
            status = self.message_base_address + "/target " + ("%f %f" % (self.vision_target[0], self.vision_target[1]))
            self.emitter.send(status.encode())

        # produce a debugging image
        return vision.annotate_circles(grayscale, circles)

   #================================================================
    def run(self):
        # Run loop to execute a periodic script until the simulation quits.
        # If the controller returns -1, the simulator is quitting.
        while self.step(EVENT_LOOP_DT) != -1:
            # Read simulator clock time.
            self.sim_time = self.getTime()

            # Process available camera data.
            self.poll_camera()

################################################################
# Start the script.
robot = Camera()
robot.run()