from threading import Thread

import cv2
import time
import apriltag
import logging

class VideoStream:
    def __init__(self, cam_id=0):
        # self.stream = cv2.VideoCapture(cam_id)
        # self.stream.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
        # self.stream.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) test.mp4
        self.stream = self.open_cam_usb(cam_id, 1280, 720)
        _, img = self.stream.read()
        self.buffer = img
        self.frames = 0
        self.running = False

    def start(self):
        t = Thread(target=self.update, args=())
        t.daemon = True
        self.running = True
        t.start()
        self.start_time = time.time()
        return self

    def update(self):
        while self.running:
            _, frame = self.stream.read()
            self.buffer = frame

            self.frames += 1
            if self.frames == 30:
                time_taken = time.time() - self.start_time
                self.start_time = time.time()
                self.frames = 0
                fps = 30 / time_taken
                print(f"##### INGEST FPS: {fps} #####")
            # print('#')
        self.stream.release()

    def read(self):
        return self.buffer

    def stop(self):
        self.running = False

    def open_cam_usb(self, dev, width, height):
    # We want to set width and height here, otherwise we could just do:
    #     return cv2.VideoCapture(dev)
    # We might get better framerates with this implementation...
        gst_str = ('v4l2src device=/dev/video{} ! '
                'video/x-raw, width=(int){}, height=(int){} ! '
                'videoconvert ! appsink').format(dev, width, height)
        return cv2.VideoCapture(gst_str, cv2.CAP_GSTREAMER)


def cuda_gray_conveter(frame):
    gpu_frame = cv2.cuda_GpuMat()
    gpu_frame.upload(frame)
    gray_cuda = cv2.cuda.cvtColor(gpu_frame, cv2.COLOR_BGR2GRAY)
    gray = gray_cuda.download()
    return gray


options = apriltag.DetectorOptions(families="tag36h11", nthreads=2)
detector = apriltag.Detector(options)

logging.basicConfig(level='INFO')

stream_0 = VideoStream(0).start()
# stream_2 = VideoStream(1).start()

frames = 0
start = time.time()
running = True
while running:
    try:
        frame = stream_0.read()
        if frame.any():
            gray_frame = cuda_gray_conveter(frame)
            result = detector.detect(gray_frame)
            if result:
                logging.info(result[0].center)
            frames += 1
            if frames == 30:
                time_taken = time.time() - start
                start = time.time()
                frames = 0
                fps = 30 / time_taken
                print(f"##### TOTAL PROCESSING FPS: {fps} #####")
    except Exception as e:
        print('Exiting cleanly...')
        stream_0.stop()
        running = False

print('Exited')