#! /usr/bin/python3.7

from datetime import datetime
from functools import partial
import queue
import time

from vimba import *
import cv2


def setup_camera(cam):
    cam.set_pixel_format(PixelFormat.BayerRG8)
    cam.ExposureTimeAbs.set(10000)
    cam.BalanceWhiteAuto.set('Off')
    cam.Gain.set(0)
    cam.AcquisitionMode.set('Continuous')
    cam.GainAuto.set('Off')
    # NB: Following adjusted for my Manta G-033C
    cam.Height.set(492)
    cam.Width.set(656)

# Called periodically as frames are received by Vimba's capture thread
# NB: This is invoked in a different thread than the rest of the code!
def frame_handler(frame_queue, cam, frame):
    img = frame.as_numpy_ndarray()
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BAYER_RG2RGB)
    try:
        # Try to put the frame in the queue...
        frame_queue.put_nowait(img_rgb)
    except queue.Full:
        # If that fials (queue is full), just drop the frame
        # NB: You may want to handle this better...
        print('Dropped Frame')
    cam.queue_frame(frame)

def do_something(img, count):
    filename = 'data/IMG_' + str(count) + '.jpg'
    cv2.putText(img, str(datetime.now()), (20, 40)
        , cv2.FONT_HERSHEY_PLAIN, 2, (255, 255, 255)
        , 2, cv2.LINE_AA)
    cv2.imwrite(filename, img)

def run_processing(cam):
    try:
        # Create a queue to use for communication between Vimba's capture thread
        # and the main thread, limit capacity to 10 entries
        frame_queue = queue.Queue(maxsize=10)
        # Start asynchronous capture, using frame_handler
        # Bind the first parameter of frame handler to our frame_queue
        cam.start_streaming(handler=partial(frame_handler,frame_queue)
            , buffer_count=10)

        start = time.time()
        frame_count = 0
        while True:
            if frame_queue.qsize() > 0:
                # If there's something in the queue, try to fetch it and process
                try:
                    frame = frame_queue.get_nowait()
                    frame_count += 1
                    cv2.imshow('Live feed', frame)
                    do_something(frame, frame_count)
                except queue.Empty:
                    pass

            key = cv2.waitKey(1)
            if (key == ord('q')) or (frame_count >= 100):
                cv2.destroyAllWindows()
                break

        fps = int((frame_count + 1)/(time.time() - start))
        print('FPS:', fps)
    finally:
        # Stop the asynchronous capture
        cam.stop_streaming()

#@profile
def main():
    with Vimba.get_instance() as vimba:
        with vimba.get_all_cameras()[0] as cam:
            setup_camera(cam)
            run_processing(cam)

if __name__ == "__main__":
    main()