# This is a CamiTK python action
#
# Video acquisition using OpenCV.
#
# This action creates a new  image component and continuously replaces
# its image data using the video feed. 
#
# Three types of video can be acquired:
# - grayscale video
# - RGB video
# - HSV mode video
# 
# You can pause/restart the video acquisition at any time.
# 
# This action uses a timer to refresh the image.

import camitk
from PySide2.QtCore import QTimer
import cv2

def capture(self):
    ret, frame = self.capture_device.read()

    if not ret:
        camitk.warning("Error: Could not read frame.")
        return None
    else:
        # Convert the image from BGR to RGB
        rgb_frame = cv2.cvtColor(frame, self.mode)
        return rgb_frame

def capture_loop(self):
    rgb_frame = capture(self)
    if rgb_frame is not None and camitk.Application.isAlive(self.image_component):
        self.image_component.replaceImageData(rgb_frame)
    else:
        self.timer.stop()
        # Release the camera
        self.capture_device.release()

def init(self:camitk.Action):
    self.timer:QTimer = None
    self.image_component_color:camitk.ImageComponent = None
    self.image_component_grayscale:camitk.ImageComponent = None
    self.image_component:camitk.ImageComponent = None
    self.camera_index = 0
    self.setApplyButtonText("Start")

def process(self:camitk.Action):
    # get the type of image to acquire
    enumValues = self.getProperty("Mode").getAttribute("enumNames")
    mode = enumValues[self.getParameterValue("Mode")]
    if mode == "Color":
        self.mode = cv2.COLOR_BGR2RGB
    elif mode == "Grayscale":
        self.mode = cv2.COLOR_BGR2GRAY
    else:
        self.mode = cv2.COLOR_BGR2HSV

    # create timer if needed
    if not self.timer:
        self.timer = QTimer()
        connected = self.timer.timeout.connect(lambda: capture_loop(self))
        
    if not self.timer.isActive():
        self.capture_device = cv2.VideoCapture(self.camera_index)
        if not self.capture_device.isOpened():
            camitk.warning(f"Error: Could not open capturing device #{self.camera_index}.")
            return

        rgb_frame = capture(self)
        if rgb_frame is not None:            
            if not camitk.Application.isAlive(self.image_component_grayscale):
                self.image_component_grayscale = None
            if not camitk.Application.isAlive(self.image_component_color):
                self.image_component_color = None
            
            if not self.image_component or not camitk.Application.isAlive(self.image_component):
                if mode == "Grayscale":
                    self.image_component_grayscale = camitk.newImageComponentFromNumpy(rgb_frame, "Webcam Video")
                    self.image_component = self.image_component_grayscale
                else:
                    self.image_component_color = camitk.newImageComponentFromNumpy(rgb_frame, "Webcam Video")
                    self.image_component = self.image_component_color   
            else:
                if mode == "Grayscale" and self.image_component != self.image_component_grayscale:
                    if self.image_component_grayscale is not None:
                        self.image_component = self.image_component_grayscale
                    else:
                        self.image_component_grayscale = camitk.newImageComponentFromNumpy(rgb_frame, "Webcam Video")
                        self.image_component = self.image_component_grayscale
                elif mode != "Grayscale" and self.image_component != self.image_component_color:
                    if self.image_component_color is not None:
                        self.image_component = self.image_component_color
                    else:
                        self.image_component_color = camitk.newImageComponentFromNumpy(rgb_frame, "Webcam Video")
                        self.image_component = self.image_component_color
                        
            self.image_component.setName(f"Webcam Video {mode}")

            # refresh application to show the newly created image component        
            self.refreshApplication()
        
        self.setApplyButtonText("Stop")
        self.timer.start(10)
    else:
        self.setApplyButtonText("Restart")
        self.timer.stop()
        # Release the camera
        self.capture_device.release()
