# Uses DECTRIS Stream2 example

import zmq
import numpy as np
import sys

sys.path.insert(0, "/opt/dectris/albula/4.1/bin")
sys.path.insert(0, "/opt/dectris/albula/4.1/python")

import dectris.albula as albula
from dectris.albula import DNoObject, DDrawingEllipse, DDrawingString, DDrawingLine

import signal
import cbor2
import time
from dectris.compression import decompress
from tifffile import imwrite

TERMINATE = False


def decode_multi_dim_array(tag, column_major):
    dimensions, contents = tag.value
    if isinstance(contents, list):
        array = np.empty((len(contents),), dtype=object)
        array[:] = contents
    elif isinstance(contents, (np.ndarray, np.generic)):
        array = contents
    else:
        raise cbor2.CBORDecodeValueError("expected array or typed array")
    return array.reshape(dimensions, order="F" if column_major else "C")


def decode_typed_array(tag, dtype):
    if not isinstance(tag.value, bytes):
        raise cbor2.CBORDecodeValueError("expected byte string in typed array")
    return np.frombuffer(tag.value, dtype=dtype)


def decode_dectris_compression(tag):
    algorithm, elem_size, encoded = tag.value
    return decompress(encoded, algorithm, elem_size=elem_size)


tag_decoders = {
    40: lambda tag: decode_multi_dim_array(tag, column_major=False),
    64: lambda tag: decode_typed_array(tag, dtype="u1"),
    65: lambda tag: decode_typed_array(tag, dtype=">u2"),
    66: lambda tag: decode_typed_array(tag, dtype=">u4"),
    67: lambda tag: decode_typed_array(tag, dtype=">u8"),
    68: lambda tag: decode_typed_array(tag, dtype="u1"),
    69: lambda tag: decode_typed_array(tag, dtype="<u2"),
    70: lambda tag: decode_typed_array(tag, dtype="<u4"),
    71: lambda tag: decode_typed_array(tag, dtype="<u8"),
    72: lambda tag: decode_typed_array(tag, dtype="i1"),
    73: lambda tag: decode_typed_array(tag, dtype=">i2"),
    74: lambda tag: decode_typed_array(tag, dtype=">i4"),
    75: lambda tag: decode_typed_array(tag, dtype=">i8"),
    77: lambda tag: decode_typed_array(tag, dtype="<i2"),
    78: lambda tag: decode_typed_array(tag, dtype="<i4"),
    79: lambda tag: decode_typed_array(tag, dtype="<i8"),
    80: lambda tag: decode_typed_array(tag, dtype=">f2"),
    81: lambda tag: decode_typed_array(tag, dtype=">f4"),
    82: lambda tag: decode_typed_array(tag, dtype=">f8"),
    83: lambda tag: decode_typed_array(tag, dtype=">f16"),
    84: lambda tag: decode_typed_array(tag, dtype="<f2"),
    85: lambda tag: decode_typed_array(tag, dtype="<f4"),
    86: lambda tag: decode_typed_array(tag, dtype="<f8"),
    87: lambda tag: decode_typed_array(tag, dtype="<f16"),
    1040: lambda tag: decode_multi_dim_array(tag, column_major=True),
    56500: lambda tag: decode_dectris_compression(tag),
}


def tag_hook(decoder, tag):
    tag_decoder = tag_decoders.get(tag.tag)
    return tag_decoder(tag) if tag_decoder else tag


def signal_handler(signal, frame):
    global TERMINATE
    TERMINATE = True


signal.signal(signal.SIGINT, signal_handler)

context = zmq.Context()
socket = context.socket(zmq.SUB)
socket.RCVTIMEO = 20000  # in milliseconds
socket.connect("tcp://127.0.0.1:5400")
socket.setsockopt(zmq.SUBSCRIBE, b"")

albulaMain = albula.openMainFrame()
albulaSubFrame = albulaMain.openSubFrame()
startTime = time.time()

OPTIONAL_DATA = albula.DImageOptionalData()

OPTIONAL_DATA.set_x_pixel_size(0.000075)
OPTIONAL_DATA.set_y_pixel_size(0.000075)

while not TERMINATE:
    try:
        print("a")
        msg = socket.recv()
        if round(time.time() - startTime) >= 60 * 60:
            albulaSubFrame.close()
            albulaSubFrame = albulaMain.openSubFrame()
            startTime = time.time()
        print("b")
        msg = cbor2.loads(msg, tag_hook=tag_hook)
        print(msg)
        #OPTIONAL_DATA.set_wavelength(12398/msg["incident_energy"])
        #OPTIONAL_DATA.set_beam_center_x(msg['beam_center_x'])
        #OPTIONAL_DATA.set_beam_center_y(msg['beam_center_y'])
        #OPTIONAL_DATA.set_detector_distance(msg["detector_distance"])

        dimage = albula.DImage(msg['data']['default'])
        dimage.setOptionalData(OPTIONAL_DATA)
        albulaSubFrame.loadImage(dimage)
        #albulaSubFrame.setTitle("JUNGFRAU PREVIEW Dataset: %s Image: %d" % (msg["unique_series_id"], msg['number']))
    except:
        pass