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

# In[ ]:


# requires python 3.6+ because I use format string literals

# get this library here: https://github.com/ZhuangLab/storm-control
from storm_control.sc_hardware.hamamatsu import hamamatsu_camera as hc

# install with pip install pyqtgraph
import pyqtgraph as pq
from pyqtgraph.Qt import QtGui

# stdlib
import time
from threading import Thread
from queue import LifoQueue

# this lets pyqtgraph without blocking execution in the notebook
get_ipython().run_line_magic('gui', 'qt')


# In[ ]:


# There are 2 camera classes, hc.HamamatsuCamera and hc.HamamatsuCameraMR, I'm not sure what the differences are yet 
# hc.HamamatsuCameraMR seems to have fewer ops? 
# I have 2 cameras, so I put them in a tuple
cams = hc.HamamatsuCameraMR(camera_id=0), hc.HamamatsuCameraMR(camera_id=1)

# Full frame at 100 FPS doesn't work, but that's a problem with pyqtgraph. 2x2 binning works.
binning = 2

# Set camera properties. I don't know what all of these mean yet. 
for cam in cams:
    cam_x, cam_y = 2048, 2048
    cam.setPropertyValue("defect_correct_mode", "OFF")
    cam.setPropertyValue("exposure_time", 0.01)
    cam.setPropertyValue("subarray_hsize", cam_x)
    cam.setPropertyValue("subarray_vsize", cam_y)
    cam.setPropertyValue("binning", f"{binning}x{binning}")
    cam.setPropertyValue("readout_speed", 2)


# In[ ]:


# Create windows for displaying images using pyqtgraph 
disps = [pq.image() for cam in cams]


# In[ ]:


# We store images in a last-in-first-out queue
queues = [LifoQueue(maxsize=10 ** 5) for cam in cams]

# Currently I start / stop acquisition based on the state of this boolean global. 
running = True

# A function that takes frames from a camera and puts them in a queue
def grabber(cam, queue):
    global running
    cam.startAcquisition()
    while running:
        frames, dims = cam.getFrames()
        for f in frames:
            queue.put(f.getData().reshape(*dims))        

    cam.stopAcquisition()    

# A function that takes frames from the queue and displays them    
def drawer(disp, queue, sleep_time=.032):
    global running
    
    while running:       
        im = queue.get() // binning
        # Optional: mess with im to reduce its size and reduce the load on pyqtgraph
        disp.imageItem.updateImage(im)
        QtGui.QApplication.processEvents()
        # need to give pyqtgraph some time to process the image before we give it another one
        time.sleep(sleep_time)
        queue.task_done()


# give our functions to threads with the right arguments 
frame_threads = [Thread(target=grabber, args=args) for args in zip(cams, queues)]
draw_threads = [Thread(target=drawer, args=args) for args in zip(disps, queues)]

# start the threads
for t in frame_threads: t.start()
for d in draw_threads: d.start()


# In[ ]:


# run this cell to stop acquisition / display
running = False


# In[ ]: