import warnings

with warnings.catch_warnings():
    warnings.filterwarnings('ignore', category=FutureWarning)
    import h5py

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'

import tensorflow as tf

import queue
from threading import Thread, Event

_OUTPUT_TENSOR_NAMES = ['detection_boxes', 'detection_scores', 'num_detections', 'detection_classes']
_INPUT_TENSOR_NAME = 'image_tensor'


class DetectionThread(Thread):

    stop_thread = Event()

    def __init__(self, model: str, cpu: bool, image_q: queue.Queue, detection_q: queue.Queue):
        super(DetectionThread, self).__init__()
        self.model = model
        self.cpu = cpu
        self.image_q = image_q
        self.detection_q = detection_q

    def run(self):

        try:
            graph = tf.Graph()

            print('M Loading graph...')
            with open(self.model, 'rb') as f:
                graph_def = tf.GraphDef.FromString(f.read())
            print('M Loaded graph.')

            with graph.as_default():
                tf.import_graph_def(graph_def, name='')

                config = tf.ConfigProto()
                if self.cpu:
                    config.device_count['GPU'] = 0
                else:
                    config.gpu_options.per_process_gpu_memory_fraction = 0.2
                    config.gpu_options.allow_growth = True

                print('M Starting TF session...')
                with tf.Session(graph=graph, config=config) as sess:

                    output_tensors = {k: graph.get_tensor_by_name(k + ':0') for k in _OUTPUT_TENSOR_NAMES}
                    image_tensor = graph.get_tensor_by_name('image_tensor:0')

                    total_detections = 0
                    while not self.stop_thread.is_set():
                        total_detections += 1

                        if total_detections % 100 == 0:
                            print(f'D total_detections={total_detections}')

                        try:
                            batch = self.image_q.get(block=True, timeout=2.0)
                        except queue.Empty:
                            continue

                        result = sess.run(output_tensors, feed_dict={image_tensor: [batch]})
                        self.detection_q.put(result)

                        for i in range(0, len(result['num_detections'])):
                            n = int(result['num_detections'][i])
                            c = result['detection_classes'][i].astype(int)
                            b = result['detection_boxes'][i]
                            s = result['detection_scores'][i]
                            for j in range(0, n):
                                if s[j] > 0.7:
                                    print('%3d %3d %3d %0.2f %s' % (i, j, c[j], s[j], b[j]))

                        if self.stop_thread.is_set():
                            raise StopIteration()

        except StopIteration:
            print(f'D Stopping.')

        print(f'D Closed detector.')

    def interrupt(self):
        print(f'D Interrupting detector.')
        self.stop_thread.set()