import sys
import warnings
import pickle
import numpy as np

from PyQt5 import QtCore, QtWidgets


def create_data():
    n_images = 20

    # detection data types for the array that stores detections made in images.
    # the first data type is an object so that later on cross references can be
    # stored there.
    det_dtypes = [('dprop0', 'O')]
    det_dtypes += [(f'dprop{i+1}', 'i4') for i in range(36)]

    # create random number of detections for each image. in reality,
    # the data and data types are different, but for this MWE, integer ones
    # seem to be enough
    det_list = [
        np.ones(np.random.randint(2900, 3100), dtype=det_dtypes)
        for _ in range(n_images)
    ]

    # go through the detections list associated with an image
    for i, detections in enumerate(det_list):
        # go through all the detections of the respective image and store a
        # list of a random number of dictionaries with cross references to
        # detection entries from other images
        for detection in detections:
            choices = list(np.arange(n_images))
            choices.remove(i)
            idx = np.random.choice(choices)
            d = {
                'cr': detections[idx][
                    np.random.randint(0, len(detections[idx])-1)
                ]
            }
            detection['dprop0'] = [d for _ in range(np.random.randint(1, 20))]

    # track data types that that store different properties of the track,
    # among other things the detections that belong to that track
    track_dtypes = [('tprop0', 'O')]
    track_dtypes += [(f'tprop{i+1}', 'i4') for i in range(15)]

    # create the tracks with their properties. again, in reality, the data
    # and data types are different, but for this MWE, integer ones seem to
    # be enough
    tracks = np.ones(100, dtype=track_dtypes)

    # for each track, chose a random assembly of detections from the
    # detections list and assign them to the track
    for track in tracks:
        track['tprop0'] = [
            det_list[i][np.random.randint(0, len(det_list[i])-1)]
            for i in range(n_images)
        ]

    # the final data structure storing all the information, like the images,
    # detections, tracks and so on. for the MWE to work, only saving the
    # tracks seems to be enough.
    data = {
        # 'images': [
        #     np.random.uniform(0, 1, (2048, 2048)) for _ in range(n_images)
        # ],
        # 'times': [np.random.uniform(0, 900, n_images)],
        # 'imstack': np.random.uniform(0, 1, (2048, 2048)),
        # 'detections': det_list,
        'tracks': tracks,
    }

    return data


class Widget(QtWidgets.QWidget):
    def __init__(self, parent=None):
        super(Widget, self).__init__(parent)
        dump_btn = QtWidgets.QPushButton("dump", clicked=self.dump)
        load_btn = QtWidgets.QPushButton("load", clicked=self.load)
        create_data_btn = QtWidgets.QPushButton("create", clicked=self.create)

        self.file = '../test_files/pickle_test'
        self.data = None

        lay = QtWidgets.QVBoxLayout(self)
        lay.addWidget(dump_btn)
        lay.addWidget(load_btn)
        lay.addWidget(create_data_btn)
        lay.addStretch()

    @QtCore.pyqtSlot()
    def create(self):
        self.data = create_data()

    @QtCore.pyqtSlot()
    def dump(self):
        with open(self.file, 'wb') as output_file:
            pickle.dump(self.data, output_file, pickle.HIGHEST_PROTOCOL)

        # in the same session, the data can still be accessed without problems
        print(self.data['tracks'][0]['tprop0'][0]['dprop0'])

    @QtCore.pyqtSlot()
    def load(self):
        with open(self.file, 'rb') as input_file:
            load_data = pickle.load(input_file)

        # this should always work
        print(load_data['tracks'][0]['tprop0'][0].dtype.names)

        # this only works if the data is loaded in the same session. loading
        # data stored in a different session results in a crash
        print(load_data['tracks'][0]['tprop0'][0]['dprop0'])


if __name__ == '__main__':
    sys._excepthook = sys.excepthook

    def exception_hook(exctype, value, traceback):
        # noinspection PyProtectedMember,PyUnresolvedReferences
        sys._excepthook(exctype, value, traceback)
        sys.exit(1)

    sys.excepthook = exception_hook
    warnings.simplefilter('error', UserWarning)
    app = QtWidgets.QApplication(sys.argv)
    w = Widget()
    w.resize(640, 480)
    w.show()
    sys.exit(app.exec_())