import zwoasi
import matplotlib.pyplot as plt
import numpy as np
import ctypes
import time
import os

from zwoasi import zwolib
if zwolib is None:
    zwoasi.init(os.path.expanduser('~/ASICAP/lib/libASICamera2.so'))
    from zwoasi import zwolib

print(zwoasi.get_num_cameras())
print(zwoasi.list_cameras())

class VideoFrame(object):
    def __init__(self):
        self.cam = zwoasi.Camera(0)
        self.cam.disable_dark_subtract()
        self.cam.set_control_value(zwoasi.ASI_GAIN, 100)
        self.cam.set_image_type(zwoasi.ASI_IMG_RAW16)
        self.cam.set_control_value(zwoasi.ASI_HIGH_SPEED_MODE,0)
        print(self.cam.get_camera_property())
        self.set_frame()

    def set_frame(self):
        self.whbi = zwoasi._get_roi_format(self.cam.id)
        self.sz = self.whbi[0] * self.whbi[1] * 2
        self.buffer_1 = bytearray(self.sz)
        self.cbuf_type = ctypes.c_char * len(self.buffer_1)
        self.cbuf1 = self.cbuf_type.from_buffer(self.buffer_1)
        self.buffer_2 = bytearray(self.sz)
        self.cbuf2 = self.cbuf_type.from_buffer(self.buffer_2)
        self.im1 = np.frombuffer(self.buffer_1, dtype='uint16')
        self.im2 = np.frombuffer(self.buffer_2, dtype='uint16')

    def video_frame(self, exptime=15e-3, nexp=100):
        ''' Take a video frame and gather the registered statistic on it
        
        Exptime given in seconds
        '''
        exptime = int(exptime*1e6)
        timeout = exptime + 50000
        self.cam.set_control_value(zwoasi.ASI_EXPOSURE, exptime)
        self.cam.start_video_capture()

        M = np.zeros(self.im1.shape, dtype='uint32')
        V = np.zeros(self.im1.shape, dtype='uint64')
        m = np.zeros(nexp, dtype='uint32')
        v = np.zeros(nexp, dtype='uint64')
        T = np.zeros(nexp, dtype='uint16')
        tic = time.time()
        i = 0
        try:
            for i in range(nexp):
                r = zwolib.ASIGetVideoData(self.cam.id, self.cbuf1, self.sz, int(timeout))
                if r:
                    raise zwo_errors[r]
                M += self.im1
                t2 = self.im1.astype('uint32')**2
                V += t2
                m[i] = self.im1.sum()
                v[i] = t2.sum()
                T[i] = self.cam.get_control_values()['Temperature']
        finally:
            self.cam.stop_video_capture()
            toc = time.time() - tic
            print(f'{i} images in {toc}s')
        return M.reshape(self.whbi[1::-1]), V.reshape(self.whbi[1::-1]), m, v, T


if __name__=='__main__':
    cam = VideoFrame()
    nexp=60000
    exptime=15
    a = cam.video_frame(exptime=exptime * 1e-3, nexp=nexp)
    import astropy.io.fits as pyfits
    Npix = a[0].shape[0] * a[0].shape[1]
    M = a[0] / len(a[2])
    Sig = np.sqrt(a[1]/len(a[2]) - M*M)
    m = a[2] / Npix
    sig = np.sqrt(a[3] / Npix - m * m)
    T = a[4]
    arr = np.rec.fromarrays([m,sig,T], names=['m', 'sig', 'T'])
    hlist = pyfits.HDUList([pyfits.PrimaryHDU(M),
                            pyfits.ImageHDU(Sig),
                            pyfits.TableHDU(arr)])
    hlist.writeto(f'exptime_{exptime}ms_nexp_{nexp}.fits')