# import the necessary packages
from django.conf import settings
import torch

import redis
import time
import json
import numpy as np
import math
from threading import Thread
import random
from dl.yolo.yolo import YOLO
from dl.face_dlib.FaceAlignment import FaceAlignment
from dl.insightFace.insightFace import InsightFace
from dl.hnsw.pyw_hnswlib import HNSWIndex
from dl import utils
from .faceBank import FaceBank
import time

def classify_process():
    # connect to Redis server
    time.sleep(5)

    db = redis.StrictRedis(host=settings.REDIS_HOST,
                           port=settings.REDIS_PORT, db=settings.REDIS_DB)

    print("===* Loading model... *===")
    yolo_model = YOLO()

    fa = FaceAlignment()

    insightFace = InsightFace(use_mobilefacenet=settings.USE_MOBILEFACENET)

    hnsw = HNSWIndex('cosine', 512)

    db.flushdb()
    try:
        hnsw.load_index()
    except:
        hnsw.init_index(max_elements=15000, ef_construction=200, M=512)

    hnsw.set_ef(512)
#    if hnsw.cur_ind == 0:
    db.set("hnswUpdateSignal", 1)

    try:
        hnswDict = np.load("hnswdict.npy").item()
    except:
        hnswDict = dict()

    delete_mask = []

    #    fb=FaceBank()
    #    if hnsw.cur_ind != hnswDict.get("len"):
    #        t=Thread(target=fb.updateFaceBank)
    #        t.start()

    print("===* Model loaded     *===")

    fb = FaceBank()

    stream = "FRStream"
    group = "FRGroup"
    consumer = "FRConsumer"
    # try:
    db.xgroup_create(stream, group, id="0-0", mkstream=True)
    # except:
    #    pass
    NewHNSWIndex=None
    # continually pool for new images to classify
    while True:
        # attempt to grab a batch of images from the database
        # db.xread(streams={stream:'>'},
        #          count=settings.BATCH_SIZE//2,
        #          block=300)

        groupData = db.xreadgroup(groupname=group,
                                  consumername=consumer,
                                  streams={stream: '>'},
                                  count=settings.BATCH_SIZE // 2,
                                  block=300)

        hnswUpdateSignal = int(db.get("hnswUpdateSignal")) \
            if db.get("hnswUpdateSignal") is not None else 0
        if hnswUpdateSignal:
            NewHNSWIndex = HNSWIndex('cosine', 512)
            NewHNSWIndex.init_index(max_elements=20000, ef_construction=100, M=512)
            db.set("hnswUpdateSignal", 0)
            t = Thread(target=fb.updateFaceBank, args=(NewHNSWIndex, db))
            t.daemon = True
            t.start()

        hnswDictUpdateSignal = int(db.get("hnswDictUpdateSignal")) \
            if db.get("hnswDictUpdateSignal") is not None else 0
        if hnswDictUpdateSignal:
            hnsw = NewHNSWIndex
            hnsw.save_index()
            delete_mask = []
            db.set("hnswDictUpdateSignal", 0)
            hnswDict = json.loads(db.get("hnswDict")) if db.get("hnswDict") is not None else hnswDict

        if len(groupData) == 0 or len(groupData[0][1]) == 0:
            continue
        # groupData[0][0] is the stream name
        groupData = groupData[0][1]

        Data = []
        for groupdata in groupData:
            id = groupdata[0].decode('utf8')
            data = json.loads(list(groupdata[1].values())[0])

            if data.get('type') == "Signal":            #Signal 项什么都不做
                db.xack(stream, group, id)
                db.xdel(stream, id)

                pass
            elif data.get("type") == "Delete":
                db.xack(stream, group, id)
                db.xdel(stream, id)
                md5 = data.get("md5")
                name = data.get("name")
                if md5 in hnswDict and name in hnswDict[md5]["name"]:
                    index = hnswDict[md5]["name"].index(name)
                    delete_mask.append(hnswDict[md5]["id"][index])
                    dictLen = hnswDict["len"]
                    hnswLen = hnsw.cur_ind - len(delete_mask)
                    if dictLen > hnswLen:
                        db.set("hnswUpdateSignal", 1)

            elif data.get("type") == "Recognition":
                Data.append((id, data))
        groupData = Data

        # Now groupData[0] is the id ang groupData[1] is the data

        batch_for_detect = []
        image_sizes_for_detect = []
        image_ids_for_detect = []
        for index, groupdata in enumerate(groupData):
            data = groupdata[1]
            if not data.get("detected"):
                # Get the image and image shape
                image = utils.base64_decode_image(data["image"])
                image_size = image.shape

                # Get the image size and id in the groupData
                image_sizes_for_detect.append(image_size)
                image_ids_for_detect.append(index)

                # data prepare and batch
                batch = yolo_model.prepare_image(image)
                batch_for_detect.append(batch)

        if len(batch_for_detect) > 0:
            start = time.time()
            data = torch.stack(batch_for_detect, 0)
            results = yolo_model.predict(data, image_sizes_for_detect)

            for (imageID, resultSet) in zip(image_ids_for_detect, results):
                # initialize the list of output predictions
                groupData[imageID][1]["face"] = []
                # loop over the results and add them to the list of
                # output predictions

                for *xyxy, conf, score, label in resultSet:
                    r = {"xmin": int(xyxy[0]),
                         "ymin": int(xyxy[1]),
                         "xmax": int(xyxy[2]),
                         "ymax": int(xyxy[3]),
                         "score": float(conf) * float(score)}
                    groupData[imageID][1]["face"].append(r)
            end = time.time()
            print("* YOLO : {} / {}".format(data.shape, end - start))

        # It's the map from ImageID to FaceImage , the ImageID is from groupData.
        image2faceMap = [None] * len(groupData)

        batch_for_recognize = []
        # It's the map from FaceID to ImageID , the FaceID is from batch_for_recognize
        # and the ImageID is from groupData.
        face2ImageIDMap = []
        face2ImageLocationMap = []

        for index, groupdata in enumerate(groupData):
            data = groupdata[1]

            image = utils.base64_decode_image(data["image"])

            face_rectangles = []
            #            if isinstance(data["face"],str):
            #                data['face']=json.loads(data["face"])

            for face in data["face"]:
                face_rectangles.append((face["xmin"], face["ymin"], face["xmax"], face["ymax"]))

            if len(face_rectangles) == 0:
                continue
            aligned_face = fa.faceAlignment(image, face_rectangles=face_rectangles)

            image2faceMap[index] = aligned_face

            for loc, face in enumerate(aligned_face):
                face_transed = insightFace.prepare_transform(face)

                # Append to the batch for recognize
                batch_for_recognize.append(face_transed)
                face2ImageIDMap.append(index)
                face2ImageLocationMap.append(loc)

        BATCH_SIZE = settings.BATCH_SIZE
        for i in range(int(math.ceil(len(batch_for_recognize) / BATCH_SIZE))):
            start = time.time()
            start_index = i * BATCH_SIZE
            end_index = (i + 1) * BATCH_SIZE
            data = torch.stack(batch_for_recognize[start_index:end_index], 0)

            embedings = insightFace.predict(data)

            for index, embeding in enumerate(embedings):

                embeding_numpy = embeding.cpu().detach().numpy()

                id = face2ImageIDMap[start_index + index]
                loc = face2ImageLocationMap[start_index + index]

                timeid = groupData[id][0]
                jsonData = groupData[id][1]

                if jsonData["saveVec"]:
                    hnsw.add_items([embeding_numpy],
                                   ids=[jsonData.get("name")])
                    hnsw.save_index()
                    if jsonData.get("md5"):
                        hnswDict = json.loads(db.get("hnswDict")) if db.get("hnswDict") is not None else hnswDict
                        md5 = jsonData["md5"]
                        name = jsonData["name"]
                        if md5 not in hnswDict:
                            hnswDict[md5] = {"name": [],
                                             "id": []}
                        hnswDict[md5]["name"].append(name)
                        hnswDict[md5]["id"].append(hnsw.cur_ind)
                        hnswDict["len"] += 1
                        db.set("hnswDict", json.dumps(hnswDict))
                    jsonData["vec"] = embeding_numpy.tolist()

                ret = {
                    "match": None,
                    "score": 0
                }

                if "recognize" not in jsonData or jsonData["recognize"] == True:
                    try:
                        if len(delete_mask) != 0:
                            labelsid, labels, distances = hnsw.knn_query_id(embeding_numpy, k=3)
                            labelsid, labels, distances = list(labelsid[0]), labels[0], distances[0]
                            new_labels = []
                            new_distances = []
                            for id in labelsid:
                                if id not in delete_mask:
                                    index = labelsid.index(id)
                                    new_labels.append(labels[index])
                                    new_distances.append(distances[index])
                            labels = new_labels
                            distances = np.array(new_distances)
                        else:
                            labels, distances = hnsw.knn_query(embeding_numpy, k=1)
                            labels, distances = labels[0], distances[0]
                        ret = {"match": labels[0],
                               "score": float(1 - distances[0])}
                        print(labels, 1 - distances)
                    except:
                        pass

                jsonData["face"][loc]["match"] = ret["match"]
                jsonData["face"][loc]["similarity"] = ret["score"]
                jsonData["retVal"] = 1

            end = time.time()
            print("* InsightFace : {} / {}".format(data.shape, end - start))

        for groupdata in groupData:
            timeid = groupdata[0]
            jsonData = groupdata[1]

            # 人脸检测失败
            if len(jsonData["face"]) == 0:
                jsonData["retVal"] = 101
            # 验证成功
            elif jsonData.get("match") is not None:
                jsonData["retVal"] = 1
            # 添加成功
            elif jsonData.get("saveVec"):
                jsonData["retVal"] = 2

            jsonData.pop("image")
            jsonData.pop("detected")
            jsonData.pop("recognize")
            jsonData.pop("savePic")
            jsonData.pop("saveVec")

            # loop over the results and add them to the list of
            # output predictions
            # store the output predictions in the database, using
            # the image ID as the key so we can fetch the results
            db.set(timeid, json.dumps(jsonData))
            db.xack(stream, group, timeid)


"""
.. _tut-sensor-locations:

Working with sensor locations
=============================

This tutorial describes how to read and plot sensor locations, and how
MNE-Python handles physical locations of sensors.

As usual we'll start by importing the modules we need and loading some
:ref:`example data <sample-dataset>`:
"""

# %%

import os
import numpy as np
import matplotlib.pyplot as plt

# the following import is required for matplotlib < 3.2:
from mpl_toolkits.mplot3d import Axes3D  # noqa
import mne

sample_data_folder = mne.datasets.sample.data_path()
sample_data_raw_file = os.path.join(sample_data_folder, 'MEG', 'sample',
                                    'sample_audvis_raw.fif')
raw = mne.io.read_raw_fif(sample_data_raw_file, preload=True, verbose=False)

# %%
# About montages and layouts
# ^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# `Montages <mne.channels.DigMontage>` contain sensor positions in 3D (x, y, z
# in meters), which can be assigned to existing EEG/MEG data. By specifying the
# locations of sensors relative to the brain,
# `Montages <mne.channels.DigMontage>` play an important role in computing the
# forward solution and inverse estimates.
#
# In contrast, `Layouts <mne.channels.Layout>` are *idealized* 2D
# representations of sensor positions. They are primarily used for arranging
# individual sensor subplots in a topoplot or for showing the *approximate*
# relative arrangement of sensors as seen from above.
#
# Working with built-in montages
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# The 3D coordinates of MEG sensors are included in the raw recordings from MEG
# systems. They are automatically stored in the ``info`` attribute of the
# `~mne.io.Raw` object upon loading. EEG electrode locations are much more
# variable because of differences in head shape. Idealized montages for many
# EEG systems are included in MNE-Python; these files are stored in your
# ``mne-python`` directory in the :file:`mne/channels/data/montages` folder:

montage_dir = os.path.join(os.path.dirname(mne.__file__),
                           'channels', 'data', 'montages')
print('\nBUILT-IN MONTAGE FILES')
print('======================')
print(sorted(os.listdir(montage_dir)))

# %%
# .. sidebar:: Computing sensor locations
#
#     If you are interested in how standard (idealized) EEG sensor positions
#     are computed on a spherical head model, make sure to check out the
#     `eeg_positions`_ repository.
#
# These built-in EEG montages can be loaded with
# `mne.channels.make_standard_montage` (note that you need to provide the
# filename *without* its extension):

ten_twenty_montage = mne.channels.make_standard_montage('standard_1020')
print(ten_twenty_montage)

# %%
# Once loaded, a montage can be applied to data with the
# `~mne.io.Raw.set_montage` method, for example
# `raw.set_montage <mne.io.Raw.set_montage>`. It is also possible to skip the
# loading step by passing the filename string directly to the
# `~mne.io.Raw.set_montage` method. This will not work with our sample
# data, because its channel names do not match the channel names in the
# standard 10–20 montage. Therefore, we do not run the following commands here:

# these will be equivalent:
# raw_1020 = raw.copy().set_montage(ten_twenty_montage)
# raw_1020 = raw.copy().set_montage('standard_1020')

# %%
# `Montage <mne.channels.DigMontage>` objects have a
# `~mne.channels.DigMontage.plot` method for visualizing the sensor locations
# in 2D or 3D:

fig = ten_twenty_montage.plot(kind='3d')
fig.gca().view_init(azim=70, elev=15)  # set view angle
ten_twenty_montage.plot(kind='topomap', show_names=False)

# %%
# .. _control-chan-projection:
#
# Controlling channel projection (MNE vs EEGLAB)
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# Channel positions in 2D space are obtained by projecting their actual 3D
# positions onto a sphere, then projecting the sphere onto a plane. Because the
# ``'standard_1020'`` montage contains realistic (as opposed to idealized
# spherical) channel positions, we will use a different montage to demonstrate
# how channels are projected to 2D:

biosemi_montage = mne.channels.make_standard_montage('biosemi64')
biosemi_montage.plot(show_names=False)

# %%
# By default, a sphere with origin at ``(0, 0, 0)`` (x, y, z coordinates) and
# radius of ``0.095`` meters (9.5 cm) is used. You can use a different sphere
# radius by passing a single value as the  ``sphere`` argument in any function
# that plots channels in 2D (like `~mne.channels.DigMontage.plot` that we use
# here, but also for example `mne.viz.plot_topomap`):

biosemi_montage.plot(show_names=False, sphere=0.07)

# %%
# To change not only the radius, but also the sphere origin, pass a
# ``(x, y, z, radius)`` tuple as the ``sphere`` argument:

biosemi_montage.plot(show_names=False, sphere=(0.03, 0.02, 0.01, 0.075))

# %%
# In MNE-Python, the head center and therefore the sphere center are calculated
# using :term:`fiducial points <fiducial>`. This means that the head circle
# represents the head circumference at the nasion and ear level, and not where
# it is commonly measured in the 10–20 EEG system (above the nasion at T4/T8,
# T3/T7, Oz, and Fz). Notice below that by default T7 and Oz are placed
# *within* the head circle:

biosemi_montage.plot()

# %%
# If you prefer to draw the head circle using 10–20 conventions (which are also
# used by EEGLAB), you can move the sphere origin a few centimeters up along
# the z dimension:

biosemi_montage.plot(sphere=(0, 0, 0.035, 0.094))

# %%
# Alternatively, you can calculate the sphere origin from Oz, Fpz, T3/T7 or
# T4/T8 channels. This is easier once the montage has been applied to the data
# and channel positions are in the head space (see
# :ref:`this example <ex-topomap-eeglab-style>`).

# %%
# .. _reading-dig-montages:
#
# Reading sensor digitization files
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# In the sample data, the sensor positions are already available in the
# ``info`` attribute of the `~mne.io.Raw` object (see the documentation of the
# reading functions and `~mne.io.Raw.set_montage` for details on how that
# works). Therefore, we can plot sensor locations directly from the
# `~mne.io.Raw` object using `~mne.io.Raw.plot_sensors`, which provides similar
# functionality to `montage.plot() <mne.channels.DigMontage.plot>`. In
# addition, `~mne.io.Raw.plot_sensors` supports channel selection by type,
# color-coding channels in various ways (by default, channels listed in
# ``raw.info['bads']`` will be plotted in red), and drawing in an existing
# Matplotlib ``Axes`` object (so the channel positions can easily be added as a
# subplot in a multi-panel figure):

# sphinx_gallery_thumbnail_number = 8
fig = plt.figure()
ax2d = fig.add_subplot(121)
ax3d = fig.add_subplot(122, projection='3d')
raw.plot_sensors(ch_type='eeg', axes=ax2d)
raw.plot_sensors(ch_type='eeg', axes=ax3d, kind='3d')
ax3d.view_init(azim=70, elev=15)

# %%
# The previous 2D topomap reveals irregularities in the EEG sensor positions in
# the :ref:`sample dataset <sample-dataset>` — this is because the sensor
# positions in that dataset are digitizations of actual sensor positions on the
# head rather than idealized sensor positions based on a spherical head model.
# Depending on the digitization device (e.g., a Polhemus Fastrak digitizer),
# you need to use different montage reading functions (see :ref:`dig-formats`).
# The resulting `montage <mne.channels.DigMontage>` can then be added to
# `~mne.io.Raw` objects by passing it as an argument to the
# `~mne.io.Raw.set_montage` method (just as we did before with the name of the
# predefined ``'standard_1020'`` montage). Once loaded, locations can be
# plotted with the `~mne.channels.DigMontage.plot` and saved with the
# `~mne.channels.DigMontage.save` methods of the
# `montage <mne.channels.DigMontage>` object.
#
# .. note::
#
#     When setting a montage with `~mne.io.Raw.set_montage`, the measurement
#     info is updated in two places (both ``chs`` and ``dig`` entries are
#     updated) – see :ref:`tut-info-class` for more details. Note that ``dig``
#     may contain HPI, fiducial, or head shape points in addition to electrode
#     locations.
#
#
# Visualizing sensors in 3D surface renderings
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# It is also possible to render an image of an MEG sensor helmet using 3D
# surface rendering instead of matplotlib. This works by calling
# `mne.viz.plot_alignment`:

fig = mne.viz.plot_alignment(raw.info, dig=False, eeg=False,
                             surfaces=[], meg=['helmet', 'sensors'],
                             coord_frame='meg')
mne.viz.set_3d_view(fig, azimuth=50, elevation=90, distance=0.5)

# %%
# Note that `~mne.viz.plot_alignment` requires an `~mne.Info` object, and can
# also render MRI surfaces of the scalp, skull, and brain (by passing a dict
# with keys like ``'head'``, ``'outer_skull'`` or ``'brain'`` to the
# ``surfaces`` parameter). This makes the function useful for
# :ref:`assessing coordinate frame transformations <tut-source-alignment>`.
# For examples of various uses of `~mne.viz.plot_alignment`, see
# :ref:`plot_montage`, :ref:`ex-eeg-on-scalp`, and :ref:`ex-plot-meg-sensors`.
#
#
# Working with layout files
# ^^^^^^^^^^^^^^^^^^^^^^^^^
#
# Similar to montages, many layout files are included with MNE-Python. They are
# stored in the :file:`mne/channels/data/layouts` folder:

layout_dir = os.path.join(os.path.dirname(mne.__file__),
                          'channels', 'data', 'layouts')
print('\nBUILT-IN LAYOUT FILES')
print('=====================')
print(sorted(os.listdir(layout_dir)))

# %%
# The file formats (and therefore file extensions) of the built-in layout and
# montage files vary considerably (because manufacturers like to use different
# conventions). However, the montage and layout loading functions in MNE-Python
# take the filename *without its extension* so you do not have to keep track of
# which file format is used by which manufacturer.
#
# To load a layout file, use the `mne.channels.read_layout` function and
# provide the filename *without* its file extension. You can then visualize the
# layout using its `~mne.channels.Layout.plot` method or equivalently passing
# the layout to `mne.viz.plot_layout`:

biosemi_layout = mne.channels.read_layout('biosemi')
biosemi_layout.plot()  # same result as mne.viz.plot_layout(biosemi_layout)

# %%
# Similar to the ``picks`` argument for selecting channels from `~mne.io.Raw`
# objects, the `~mne.channels.Layout.plot` method of `~mne.channels.Layout`
# objects also has a ``picks`` argument. However, because layouts only contain
# information about sensor name and location (not sensor type), the
# `~mne.channels.Layout.plot` method only supports picking channels by index
# (not by name or by type). In the following example, we find the desired
# indices using `numpy.where`; selection by name or type is possible with
# `mne.pick_channels` or `mne.pick_types`.

midline = np.where([name.endswith('z') for name in biosemi_layout.names])[0]
biosemi_layout.plot(picks=midline)

# %%
# If you have a `~mne.io.Raw` object that contains sensor positions, you can
# create a `~mne.channels.Layout` object with either
# `mne.channels.make_eeg_layout` or `mne.channels.find_layout`.

layout_from_raw = mne.channels.make_eeg_layout(raw.info)
# same result as mne.channels.find_layout(raw.info, ch_type='eeg')
layout_from_raw.plot()

# %%
# .. note::
#
#     There is no corresponding ``make_meg_layout()`` function because sensor
#     locations are fixed in an MEG system (unlike in EEG, where sensor caps
#     deform to fit snugly on a specific head). Therefore, MEG layouts are
#     consistent (constant) for a given system and you can simply load them
#     with `mne.channels.read_layout` or use `mne.channels.find_layout` with
#     the ``ch_type`` parameter (as previously demonstrated for EEG).
#
# All `~mne.channels.Layout` objects have a `~mne.channels.Layout.save` method
# that writes layouts to disk as either :file:`.lout` or :file:`.lay` formats
# (inferred from the file extension contained in the ``fname`` argument). The
# choice between :file:`.lout` and :file:`.lay` format only matters if you need
# to load the layout file in some other application (MNE-Python can read both
# formats).
#
#
# .. LINKS
#
# .. _`eeg_positions`: https://github.com/sappelhoff/eeg_positions


import base64
import io
import zmq
import cv2
import os
import numpy as np
import tensorflow as tf
import sys
import time

class CarDetector(object):
    def __init__(self, frozen_graph):

# Tensorflow localization/detection model
# Single-shot-dectection with mobile net architecture trained on COCO dataset
# setup tensorflow graph
        self.detection_graph = tf.Graph()

# configuration for possible GPU use
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True

# load frozen tensorflow detection model and initialize
# the tensorflow graph

        with tf.gfile.GFile(frozen_graph, "rb") as f:
            self.graph_def = tf.GraphDef()
            self.graph_def.ParseFromString(f.read())

# import the graph_def into a new Graph and returns it
        with self.detection_graph.as_default():
            tf.import_graph_def(self.graph_def, name="")

            self.sess = tf.Session(graph=self.detection_graph, config=config)
            self.image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
# Each box represents a part of the image where a particular object was detected.
            self.boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
# Each score represent how level of confidence for each of the objects.
# Score is shown on the result image, together with the class label.
            self.scores =self.detection_graph.get_tensor_by_name('detection_scores:0')
            self.classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
            self.num_detections =self.detection_graph.get_tensor_by_name('num_detections:0')
        self.image_width = 640
        self.image_height = 480

    def detect(self, image):
        """
        Args:
            image: camera image

        Returns:
            list of bounding boxes: coordinates [y_up, x_left, y_down, x_right]

        """
        self.image_width = image.shape[0]
        self.image_height = image.shape[1]

        with self.detection_graph.as_default():
            image_expanded = np.expand_dims(image, axis=0)
            (boxes, scores, classes, num_detections) = self.sess.run(
                [self.boxes, self.scores, self.classes, self.num_detections],
                feed_dict={self.image_tensor: image_expanded})

            self.out_boxes = np.squeeze(boxes)
            self.out_classes = np.squeeze(classes).astype("int")
            self.out_scores = np.squeeze(scores)
            self.out_number = len(self.out_classes)

    def box_isvalid(self, box, score, category, confidence = 0.3, size_limit = 0.5, ratio_limit = 0.1):
#       Examine whether a detected box is valid (good confidence, reasonable shape and size, etc).
        h = box[2] - box[0]
        w = box[3] - box[1]
        ratio = h/(w+0.01)
        return ratio>ratio_limit and ratio<1/ratio_limit and h<size_limit and w<size_limit and category<9 and score>confidence

    def output(self, confidence =0.3, quiet = True):
        self.car_boxes = []
        category_index={1: {'id': 1, 'name': u'person'},
                        2: {'id': 2, 'name': u'bicycle'},
                        3: {'id': 3, 'name': u'car'},
                        4: {'id': 4, 'name': u'motorcycle'},
                        5: {'id': 5, 'name': u'airplane'},
                        6: {'id': 6, 'name': u'bus'},
                        7: {'id': 7, 'name': u'train'},
                        8: {'id': 8, 'name': u'truck'},
                        9: {'id': 9, 'name': u'boat'},
                        10: {'id': 10, 'name': u'traffic light'},
                        11: {'id': 11, 'name': u'fire hydrant'},
                        13: {'id': 13, 'name': u'stop sign'},
                        14: {'id': 14, 'name': u'parking meter'}}

              # The ID for car in COCO data set is 3
            #idx_vec = [i for i, v in enumerate(cls) if ((v==3) and (scores[i]>0.3))]
        for i in range(self.out_number):
            category = self.out_classes[i]
            score = self.out_scores[i]
            box = self.out_boxes[i]
            if self.box_isvalid(box, score, category):
                box[0] = self.image_width*box[0]
                box[2] = self.image_width*box[2]
                box[1] = self.image_height*box[1]
                box[3] = self.image_height*box[3]
                box = box.astype("int")
                self.car_boxes.append([box[0],box[1],box[2],box[3],score,category])
                if not quiet:
                    print(box, ', confidence: ', self.out_scores[i], ', category: ',category_index[category]['name'])
        if len(self.car_boxes) ==0 and not quiet:
            print('no detection!')

        return self.car_boxes

    def output_as_str(self, confidence =0.3, quiet = True):
        self.output(confidence = confidence, quiet = quiet)
        return [[str(item) for item in box] for box in self.car_boxes]

def DL_init(frozen_graph = "/home/ubuntu/frozen_inference_graph.pb", dummy = cv2.imread("/home/ubuntu/dummy.jpg")):
# Initialize the SSD detector.
    det = CarDetector(frozen_graph)
    det.detect(dummy)
    return det

def speedtest(path):
    limit = 100
    detector = DL_init()
    start = time.time()
    for i,item in enumerate(os.listdir(path)):
        img = cv2.imread(path+"/"+item)
        detector.detect(img)
        boxes = detector.output_as_str()
        if i>limit:
            end = time.time()
            elapsed = end-start
            print("%d files in %f seconds." % (limit,elapsed))
            break

if __name__=="__main__":
    speedtest(sys.argv[1])


import pickle
import os
import time
import numpy as np
import sys
import shutil
from openmmlib import openmmlib
from openmmlib import polymerutils
from openmmlib.polymerutils import scanBlocks
from openmmlib.openmmlib import Simulation
from openmmlib.polymerutils import grow_rw
import pyximport; pyximport.install()
from smcTranslocator_MovingBarrier import smcTranslocatorDirectional
import tools
import random

# -------defining parameters----------
# -- basic loop extrusion parameters--

logname="log.txt"
GPU = 0
LIFETIME = 100   # Processivity of cohesin, Default: 100 for WT, 1000 for Wapl KO
SEPARATION = 200 # Separation LEFs in number of monomers, Default: 200 for WT, 100 for Wapl KO, assuming a  monomer size of 1kb
N = 10000                     # System size in number of monomers
smcStepsPerBlock = 1          # Number of LEF steps between blocks of polymer simulations, Default: 1
stiff = 1                     # Polymer siffness in unit of bead size, default: 1
dens = 0.2                    # density in beads / volume. The density can roughly be estimated by looking at the amount of DNA in a nuclear volume, Default: 0.2
box = (N / dens) ** 0.33      # Define size of the bounding box for Periodic Boundary Conditions, Default: 0.33
data = polymerutils.grow_rw(N, int(box) - 2)  # creates a compact conformation
block = 0  # starting block
stg = 0.8# same as Banigan, van den Berg, Brandao eLife 2020 #0.1  #stall probability at ctcf sites
unstallLEFRate = 0.005 # about once per two typical LEF lifetimes.
ctcf_interval = 300 # 300 kb, e.g., see Busslinger et al. Nature 2017

# -- LEF and transcription dynamics --

# Note that the each rate can have a maximum value of 1.
lef_speed=1.0
lefperm = 0. # controls amount of LEF-LEF bypassing we have
pauseArray = np.ones(N,dtype=np.double) # The speed of LEFS at each position. Default is an array of ones, which means that LEFS go on maximum speed everywhere. To simulate the presence of CTCF Aafke reduced the speed at CTCF sites to 0.005. However, this doesn't account for directionality of CTCF!
shrinkPauseArray=np.zeros(N,dtype=np.double)# speed at which loops shrink
shrink_speed=0.0
kinPol= 0.001     # PolII initiation rate, Default range: 0.00025-0.002
kterPol=0.002                 # Termination rate of PolII, Default range: 0.002-1.0. Normally transcription is initiation limited, so make sure that kinPol < kterPol
kterPolArray=np.zeros(N,dtype=np.double) # fix this later based on gene stucture
PolSpeed=0.1                  # The speed of PolII as a fraction of the speed of cohesin. Default value: 0.1
poldissoc=0.
#pauseArrayPol=np.zeros(N,dtype=np.double) + PolSpeed # The speed of PolII at each lattice site #initialize later
polloading=np.array([950,1950,2950,3950,4950,5950,6950,7950,8950]) # PolII initiation sites. The direction of the gene is set by the relative position of the polloading and poltermination sites.
#np.array([700,903,1100,1303,1700,1903,2100,2303,2700,2903,3100,3303,3700,3903,4100,4303,4700,4903,5100,5303,5700,5903,6100,6303])
poltermination=np.array([1850,2850,3850,4850,5850,6850,7850,8850,9850]) #PolII termination sites. At a termination site, PolII stalls with probability 1 and then unloads with rate kterPol. To simulate a broad PolII unloading area, I would set the poltermination site far beyond the gene length and then define wide region where PolII can stall, as shown on the lines below
#np.array([800,803,1200,1203,1800,1803,2200,2203,2800,2803,3200,3203,3800,3803,4200,4203,4800,4803,5200,5203,5800,5803,6200,6203])
stalProbPol=np.zeros(N,dtype=np.double) # The rate of PolII stalling. Once stalled, PolII unloads with a rate 'kterPol'. One can choose a single stall site, or a range of stall sites. If this array is set to zero everywhere, PolII will stall at the defined termination sites.
stalProb=0.001
stall_in_gene=0.
unstall_in_gene=1.
unstallArray=np.zeros(N,dtype=np.double) #array for unstalling in of Pol II in gene

STALL_FROM_FILE = False # option to take stall probabilities from file for spatially varying patterns
stallfile=""
gene_stall=[]
ctcf_left_list=[]
ctcf_right_list=[]
strongCTCFstall = 0

TSSloadbias = 1.0
TSSloadstart=0 #use this and TSSloadend to control offset/width of loading near TSS,  note: positive TSSloadstart denotes num sites before TSS, positive TSSloadend denotes sites after TSS
TSSloadend=0

base_genelen=200
#genelength=[base_genelen]*len(polloading)
variable_genelength = 0
fixed_variable_genelength=0
variable_permeability = 0 #This varaible toggles on/off variable permeability of cohesin through RNAp depending on position in gene, 0 or 1 for constant permeability or variable perm, respectively
variable_type = 0 #type of variable permeability: 0- linearly decreasing, 1- step function
variable_offset = 0 #param for var perm function: for linear this is a constant offset from 0, for step this is offset for the bottom step
variable_pos = base_genelen // 2 #param: irrelevant for linear, position of step for step function.
variable_return = 0 #return to max permeability after passing end of gene but before reaching termination site
variable_TSSfactor = 1.0  # factor by which permeability at TSS is lower (or higher)

permLeftArray=np.zeros(N,dtype=np.double)
permRightArray=np.zeros(N,dtype=np.double)

collisionLifeFactor=1.0 # factor by which head-on collision with RNAP changes lifetime of cohesin
collisionLife=LIFETIME*collisionLifeFactor
TTSunload=1.0 # factor by which cohesin life time is changed near TTS
unloadZone=10 # width of zone in which cohesin life is changed by TTS

PolPause=0.002   # Step rate PolII at TSS, Default: 0.002
#later:
#for i in polloading:
#    pauseArrayPol[i]=PolPause

L=0.          # Set permeability of PolII to cohesin. L=0 means PolII is impermeable.
R=0.          # Set permeability of PolII to cohesin coming from the right

run_id=1

# -- polymer simulation settings --

steps = int(200*(smcStepsPerBlock)) # nr of 3D simulation blocks btw advancing LEFs. For deterministic stepping choose 200-250 steps per block, otherwise, rescale with stepping probability. When genes are sparse, smcStepsPerBLock is approximately the number of smc steps per smc block.

saveRNAP=True # whether or not to print RNAP positions w/ each printed block

saveEveryBlocks = int(200/(smcStepsPerBlock))  # number of blocks until polymer configuration is saved
skipSavedBlocksBeginning = int(20/(smcStepsPerBlock))  # how many blocks (saved) to skip after you restart LEF positions
#totalSavedBlocks = 5000  # how many blocks to save (number of blocks done is totalSavedBlocks * saveEveryBlocks)
totalSavedBlocks = 4000  # how many blocks to save (number of blocks done is totalSavedBlocks * saveEveryBlocks)
#restartMilkerEveryBlocks = int(200/(smcStepsPerBlock))
restartMilkerEveryBlocks = int(400/(smcStepsPerBlock))
#Only one Hamiltonian can be loaded at a time to the simkt, but we want to update the bonds every time a LEF steps. Loading a new Hamiltonian costs a lot of time. Instead we precalculate bonds and load all positions at once as one big Hamiltonian and just change the prefactors.

# parameters for smc bonds

smcBondWiggleDist = 0.2
smcBondDist = 0.5

#if len(sys.argv)!=8:
#    print("Warning: Number of input arguments != 8")
#    sys.exit('Number of input arguments is not correct')

FLAG="" #extra label for directory name

#######use custom class to parse inputs with keywords#######################

params= tools.argsList()
for p in params.arg_dict:
    print(p, params.arg_dict[p])

if "gpu" in params.arg_dict:
    GPU = int(params.arg_dict["gpu"])
if "lifetime" in params.arg_dict:
    LIFETIME = float(params.arg_dict["lifetime"])
if "separation" in params.arg_dict:
    SEPARATION = float(params.arg_dict["separation"])
if "initiation" in params.arg_dict:
    kinPol = float(params.arg_dict["initiation"])
if "termination" in params.arg_dict:
    kterPol = float(params.arg_dict["termination"])
if "dissociation" in params.arg_dict:
    poldissoc= float(params.arg_dict["dissociation"])
if "stall" in params.arg_dict: #pol stall
    stalProb = float(params.arg_dict["stall"])
if "stallgene" in params.arg_dict:
    stall_in_gene = float(params.arg_dict["stallgene"])
if "unstall" in params.arg_dict: # pol unstall
    unstall_in_gene=float(params.arg_dict["unstall"])
if "lefspeed" in params.arg_dict:
    lef_speed=float(params.arg_dict["lefspeed"])
    pauseArray= lef_speed*pauseArray
if "lefperm" in params.arg_dict:
    lefperm=float(params.arg_dict["lefperm"])
if "shrink" in params.arg_dict:
    shrink_speed=float(params.arg_dict["shrink"])
    shrinkPauseArray= shrinkPauseArray + shrink_speed
    if shrink_speed+lef_speed>1.:
        print("WARNING: step+shrink > 1!! extrusion probabilities will not be computed correctly.\n")
if "lefstall" in params.arg_dict:
    stg = float(params.arg_dict["lefstall"])
if "lefunstall" in params.arg_dict:
    unstallLEFRate = float(params.arg_dict["lefunstall"])
if "polspeed" in params.arg_dict:
    PolSpeed = float(params.arg_dict["polspeed"])
if "polpause" in params.arg_dict:
    PolPause = float(params.arg_dict["polpause"])
if "permL" in params.arg_dict:
    L=float(params.arg_dict["permL"])
if "permR" in params.arg_dict:
    R=float(params.arg_dict["permR"])
if "collisionlife" in params.arg_dict:
    collisionLifeFactor=float(params.arg_dict["collisionlife"])
if "tssload" in params.arg_dict:
    TSSloadbias=float(params.arg_dict["tssload"])
if "tssloadstart" in params.arg_dict:
    TSSloadstart=int(params.arg_dict["tssloadstart"])
if "tssloadend" in params.arg_dict:
    TSSloadend=int(params.arg_dict["tssloadend"])
if "ttsunload" in params.arg_dict:
    TTSunload=float(params.arg_dict["ttsunload"])
if "ttszone" in params.arg_dict:
    unloadZone=int(params.arg_dict["ttszone"])
if "genelen" in params.arg_dict:
    base_genelen=int(params.arg_dict["genelen"])
if "vperm" in params.arg_dict:
    variable_permeability = int(params.arg_dict["vperm"])
if "vpermtype" in params.arg_dict:
    variable_type = int(params.arg_dict["vpermtype"])
if "vpermoffset" in params.arg_dict:
    variable_offset = float(params.arg_dict["vpermoffset"])
if "vpermpos" in params.arg_dict:
    variable_pos = int(params.arg_dict["vpermpos"])
if "vreturn" in params.arg_dict:
    variable_return = int(params.arg_dict["vreturn"])
if "vtss" in params.arg_dict:
    variable_TSSfactor = float(params.arg_dict["vtss"])
if "vgene" in params.arg_dict:
    variable_genelength = int(params.arg_dict["vgene"])
if "fixed_vgene" in params.arg_dict:
    fixed_variable_genelength=int(params.arg_dict["fixed_vgene"])
if "convergent" in params.arg_dict:
    if int(params.arg_dict["convergent"]):
        #print("warning only correct for genelength = 110 right now.\n")
        polloading=     np.array([840,  1160, 1640, 1960, 2440, 2760, 3240, 3560, 4040, 4360, 4840, 5160, 5640, 5960, 6440, 6760, 7240, 7560, 8040, 8360, 8840, 9160])
        poltermination= np.array([1000, 1001, 1800, 1801, 2600, 2601, 3400, 3401, 4200, 4201, 5000, 5001, 5800, 5801, 6600, 6601, 7400, 7401, 8200, 8201, 9000, 9001])
if "sparse" in params.arg_dict:
    if int(params.arg_dict["sparse"]):
        #print("warning only correct for genelength = 110 right now.\n")
        polloading= np.array([1950,3950,5950,7950])
        poltermination=np.array([2850,4850,6850,8850])
if "ctcfint" in params.arg_dict:
    ctcf_interval=int(params.arg_dict["ctcfint"])
if "ctcf" in params.arg_dict:
    if params.arg_dict["ctcf"] == "tss":
        #puts ctcf just before each tss, alternating left/right
        ctcf_left_list = polloading[0::2] - 1
        ctcf_right_list = polloading[1::2] - 1
    elif params.arg_dict["ctcf"] == "body":
        #puts ctcf halfway through the gene, assuming fixed gene length and genes pointing downstream
        ctcf_left_list = polloading[0::2] + base_genelen // 2
        ctcf_right_list = polloading[1::2] + base_genelen // 2
    elif params.arg_dict["ctcf"] == "distributed":
        #put in sites every ~300 kb
        Nctcf = int(N/ctcf_interval/2)
        ctcf_left_list = np.arange(1,Nctcf+2)*2*ctcf_interval-ctcf_interval
        ctcf_right_list = np.arange(1,Nctcf+1)*2*ctcf_interval
        ctcf_left_list=np.delete(ctcf_left_list, np.where(ctcf_left_list>=N)[0])
        ctcf_right_list=np.delete(ctcf_right_list, np.where(ctcf_right_list>=N)[0])
    elif params.arg_dict["ctcf"] == "distributed2":
        Nctcf = int(N/ctcf_interval/2)
        ctcf_left_list = np.arange(1,Nctcf+2)*2*ctcf_interval-int(3*ctcf_interval/2)
        ctcf_right_list = np.arange(1,Nctcf+1)*2*ctcf_interval-int(ctcf_interval/2)
        ctcf_left_list=np.delete(ctcf_left_list, np.where(ctcf_left_list>=N)[0])
        ctcf_right_list=np.delete(ctcf_right_list, np.where(ctcf_right_list>=N)[0])
    else:
        with open(params.arg_dict["ctcf"], "r") as ctcffile:
            ctcfdata=ctcffile.readlines()
            #should only be 2 lines
            entries=ctcfdata[0].split()
            ctcf_left_list = np.array([int(x) for x in entries])
            entries=ctcfdata[1].split()
            ctcf_right_list = np.array([int(x) for x in entries])
if "strongctcf" in params.arg_dict:
    strongCTCFstall= int(params.arg_dict["strongctcf"])
#granting command line control of these variables for flexibility:
if "save" in params.arg_dict:
    saveEveryBlocks = int(int(params.arg_dict["save"])/(smcStepsPerBlock))
if "skip" in params.arg_dict:
    skipSavedBlocksBeginning = int(int(params.arg_dict["skip"])/(smcStepsPerBlock))
if "total" in params.arg_dict:
    totalSavedBlocks = int(params.arg_dict["total"])
if "stallfile" in params.arg_dict:
    stallfile=params.arg_dict["stallfile"]
    STALL_FROM_FILE=True
if "restart" in params.arg_dict:
    restartMilkerEveryBlocks = int(int(params.arg_dict["restart"])/(smcStepsPerBlock))
if "log" in params.arg_dict:
    logname=params.arg_dict["log"]
if "flag" in params.arg_dict:
    FLAG=params.arg_dict["flag"]

###### a few variables to be initialized last b/c they depend on others ###
collisionLife= LIFETIME*collisionLifeFactor

pauseArrayPol=np.zeros(N,dtype=np.double) + PolSpeed ## The speed of PolII at each lattice site
for i in polloading:
    pauseArrayPol[i]=PolPause

genelength=[base_genelen]*len(polloading)
if variable_genelength:
    if fixed_variable_genelength:
        genelength=random.shuffle([80,81,82,84,86,90,95,100,110])
    else:
        for ii in range(len(genelength)):
            genelength[ii] = genelength[ii] + int(np.random.normal(0, np.sqrt(genelength[ii])))
    print("genelengths:", genelength)

#stalling, unstalling and termination in gene and after gene ends
for i,j, gl in zip(polloading,poltermination, genelength):
    if j>i:
        stalProbPol[i:i+gl]=stall_in_gene
        unstallArray[i:i+gl]=unstall_in_gene
        stalProbPol[i+gl:j]=stalProb # This line adds stall probability to all sites between end of gene (loading site + gene len) and the "termination site" (where pol stops with prob=1)
               # int(poltermination[0]-polloading[0])
        kterPolArray[i+gl:j]=kterPol
        kterPolArray[j]=kterPol
    else: # flipped genes
        stalProbPol[i-gl+1:i+1]=stall_in_gene
        unstallArray[i-gl:i+1]=unstall_in_gene
        stalProbPol[j+1:i-gl+1]=stalProb
        kterPolArray[j+1:i-gl+1]=kterPol
        kterPolArray[j]=kterPol

#permeabilities
if not variable_permeability:
        permLeftArray += L
        permRightArray += R
else:
    for ii in range(len(genelength)):
        direc = 2*int(poltermination[ii]>polloading[ii])-1
        gene_end=genelength[ii]*direc+polloading[ii]
        RL_factor=1.
        if not L==0.:
            RL_factor=R/L
        if variable_type == 0:
            for j in range(polloading[ii], gene_end):
                permLeftArray[j] = L*(gene_end-j)/(gene_end-polloading[ii])*direc + variable_offset
                permRightArray[j] = R*(gene_end-j)/(gene_end-polloading[ii])*direc + variable_offset*RL_factor
            permLeftArray[polloading[ii]] *= variable_TSSfactor
            permRightArray[polloading[ii]] *= variable_TSSfactor
            if gene_end<poltermination[ii]:
                s1=gene_end
                s2=poltermination[ii]
            else:
                s1=poltermination[ii]+1
                s2=gene_end+1
            for j in range(s1,s2):
                if not variable_return:
                    permLeftArray[j] = variable_offset
                    permRightArray[j] = variable_offset*RL_factor  # note that RL_factor adjusts right-left permeability of RNAP, not head-on/co-directional collision permeability. should change that.
                else:
                    permLeftArray[j] = L + variable_offset
                    permRightArray[j] = R + variable_offset*RL_factor
        elif variable_type == 1:
            perm_break=polloading[ii]+direc*variable_pos
            if polloading[ii] < perm_break:
                s1=polloading[ii]
                s2=perm_break
            else:
                s1=perm_break+1
                s2=polloading[ii]+1
            for j in range(s1,s2):
                permLeftArray[j] = L
                permRightArray[j] = R
            permLeftArray[polloading[ii]] *= variable_TSSfactor
            permRightArray[polloading[ii]] *= variable_TSSfactor
            if perm_break<gene_end:
                s1=perm_break
                s2=gene_end
            else:
                s1=gene_end+1
                s2=perm_break+1
            for j in range(s1,s2):
                permLeftArray[j] = variable_offset
                permRightArray[j] = variable_offset*RL_factor
            if gene_end<poltermination[ii]:
                s1=gene_end
                s2=poltermination[ii]
            else:
                s1=poltermination[ii]+1
                s2=gene_end+1
            for j in range(s1,s2):
                if not variable_return:
                    permLeftArray[j] = variable_offset
                    permRightArray[j] = variable_offset*RL_factor
                else:
                    permLeftArray[j] = L
                    permRightArray[j] = R

# -- data folder -
folder="data/"
#folder = "/net/levsha/scratch/ebanigan/txn/"

FullFileName=folder
while True:
    fname="extr_life{0}_sep{1}_density{2}_N{3}_kinpol{4}_kterPol{5}_TSSPause{6}_save{7}_total{8}_PSpeed{9}_stf{10}_stg{11}_th0.001_genelen{12}_stall{13}_L{14}_R{15}".format(LIFETIME, SEPARATION, dens,N,kinPol,kterPol,PolPause,saveEveryBlocks,totalSavedBlocks,PolSpeed,stiff,stg,base_genelen,stalProb,L,R)
    if variable_permeability:
        fname=fname+"_vperm"+str(variable_type)+"_offset"+str(variable_offset)
        if variable_type == 1:
            fname=fname+"_pos"+str(variable_pos)
        if variable_return == 1:
            fname=fname+"_vret"
        if not (variable_TSSfactor == 1.):
            fname=fname+"_vtss"+str(variable_TSSfactor)
    if not (TSSloadbias == 1.):
        fname=fname+"_tssload"+str(TSSloadbias)
        if not ((TSSloadstart==0) and (TSSloadend==0)):
            fname=fname+"_st{0}end{1}".format(TSSloadstart,TSSloadend)
    if variable_genelength:
        fname=fname+"_vgenelen"
    if not (collisionLifeFactor == 1.0):
        fname=fname+"_collife"+str(collisionLife)
    if not (TTSunload == 1.):
        fname=fname+"_ttslife"+str(TTSunload*LIFETIME)+"width"+str(unloadZone)
    if poldissoc > 0.:
        fname=fname+"_dissoc"+str(poldissoc)
    if lef_speed < 1.0:
        fname=fname+"_lefspeed"+str(lef_speed)
    if shrink_speed > 0.0:
        fname=fname+"_shrink"+str(shrink_speed)
    if stall_in_gene > 0.:
        fname=fname+"_stallgene{0}_unstall{1}".format(stall_in_gene, unstall_in_gene)
    if lefperm > 0.:
        fname=fname+"_lefperm{0}".format(lefperm)
    fname=fname+FLAG
    fname=fname+"_"+str(run_id)
    FullFileName=os.path.join(folder, fname)
    sleeptime=np.random.uniform(0,10)
    os.system("sleep {0}".format(sleeptime))
    if not os.path.exists(FullFileName):
        os.system("mkdir {0}".format(FullFileName))
        print("directory is {0}".format(FullFileName))
        break
    else:
        run_id = run_id+1

# -- Assertions to make sure parameters have been chosen correctly --

assert restartMilkerEveryBlocks % saveEveryBlocks == 0
assert (skipSavedBlocksBeginning * saveEveryBlocks) % restartMilkerEveryBlocks == 0
assert (totalSavedBlocks * saveEveryBlocks) % restartMilkerEveryBlocks == 0
assert smcStepsPerBlock<6 # max number of steps per smc block should not be too large to prevent 'jerky' polymer motion

savesPerMilker = restartMilkerEveryBlocks // saveEveryBlocks
milkerInitsSkip = saveEveryBlocks * skipSavedBlocksBeginning  // restartMilkerEveryBlocks
milkerInitsTotal  = (totalSavedBlocks + skipSavedBlocksBeginning) * saveEveryBlocks // restartMilkerEveryBlocks
print("Milker will be initialized {0} times, first {1} will be skipped".format(milkerInitsTotal, milkerInitsSkip))

# create filenames for Ekin, Epot and time

Ekin_fname = os.path.join(FullFileName,'Ekin.txt')
Epot_fname = os.path.join(FullFileName,'Epot.txt')
time_fname = os.path.join(FullFileName,'time.txt')
Par_fname  = os.path.join(FullFileName,'Pars.txt')

def save_Es_ts_Rg():
    with open(time_fname, "a+") as time_file:
        time_file.write('%f\n'%(a.state.getTime()/openmmlib.ps))
    with open(Ekin_fname, "a+") as Ekin_file:
        Ekin_file.write('%f\n'%((a.state.getKineticEnergy())/a.N/a.kT))
    with open(Epot_fname, "a+") as Epot_file:
        Epot_file.write('%f\n'%((a.state.getPotentialEnergy()) /a.N /a.kT))

class smcTranslocatorMilker(object):

    def __init__(self, smcTransObject):
        """
        :param smcTransObject: smc translocator object to work with
        """
        self.smcObject = smcTransObject
        self.allBonds = []

    def setParams(self, activeParamDict, inactiveParamDict):
        """
        A method to set parameters for bonds.
        It is a separate method because you may want to have a Simulation object already existing

        :param activeParamDict: a dict (argument:value) of addBond arguments for active bonds
        :param inactiveParamDict:  a dict (argument:value) of addBond arguments for inactive bonds

        """
        self.activeParamDict = activeParamDict
        self.inactiveParamDict = inactiveParamDict

    def setup(self, bondForce,  blocks = 100, smcStepsPerBlock = 1):
        """
        A method that milks smcTranslocator object
        and creates a set of unique bonds, etc.

        :param bondForce: a bondforce object (new after simulation restart!)
        :param blocks: number of blocks to precalculate
        :param smcStepsPerBlock: number of smcTranslocator steps per block
        :return:
        """

        if len(self.allBonds) != 0:
            raise ValueError("Not all bonds were used; {0} sets left".format(len(self.allBonds)))

        self.bondForce = bondForce

        #precalculating all bonds
        allBonds = []
        for dummy in range(blocks):
            self.smcObject.steps(smcStepsPerBlock)
            left, right = self.smcObject.getSMCs()
            bonds = [(int(i), int(j)) for i,j in zip(left, right)]
            allBonds.append(bonds)

        self.allBonds = allBonds
        self.uniqueBonds = list(set(sum(allBonds, []))) # 'sum' preserves order and makes one long list with bonds, 'set' creates a set with left bonds from different time points ordered from small to large and eliminates two equal bonds (also if they were created by different LEFs at different times). List turns set into a list with unique bonds at different time points.

        # adding forces and getting bond indices
        self.bondInds = []
        self.curBonds = allBonds.pop(0) # pop(0) removes and returns first list of bonds

        for bond in self.uniqueBonds:
            paramset = self.activeParamDict if (bond in self.curBonds) else self.inactiveParamDict
            ind = bondForce.addBond(bond[0], bond[1], **paramset)
            self.bondInds.append(ind)
        self.bondToInd = {i:j for i,j in zip(self.uniqueBonds, self.bondInds)}
        return self.curBonds,[]

    def step(self, context, verbose=False):
        """
        Update the bonds to the next step.
        It sets bonds for you automatically!
        :param context:  context
        :return: (current bonds, previous step bonds); just for reference
        """
        if len(self.allBonds) == 0:
            raise ValueError("No bonds left to run; you should restart simulation and run setup  again")

        pastBonds = self.curBonds
        self.curBonds = self.allBonds.pop(0)  # getting current bonds
        bondsRemove = [i for i in pastBonds if i not in self.curBonds]
        bondsAdd = [i for i in self.curBonds if i not in pastBonds]
        bondsStay = [i for i in pastBonds if i in self.curBonds]
        if verbose:
            print("{0} bonds stay, {1} new bonds, {2} bonds removed".format(len(bondsStay),
                                                                            len(bondsAdd), len(bondsRemove)))
        bondsToChange = bondsAdd + bondsRemove
        bondsIsAdd = [True] * len(bondsAdd) + [False] * len(bondsRemove)
        for bond, isAdd in zip(bondsToChange, bondsIsAdd):
            ind = self.bondToInd[bond]
            paramset = self.activeParamDict if isAdd else self.inactiveParamDict
            self.bondForce.setBondParameters(ind, bond[0], bond[1], **paramset)  # actually updating bonds
        self.bondForce.updateParametersInContext(context)  # now run this to update things in the context
        return self.curBonds, pastBonds

    def getRNAP(self):
        return self.smcObject.getPolOccupied()

def initModel():
    birthArray = np.zeros(N,dtype=np.double) + 0.1
    #birthArray[polloading] *= TSSloadbias
    for nn, tt  in zip(polloading,poltermination):
        #the following accounts for different gene directions:
        if tt > nn:
            loadbegin=min(max(nn-TSSloadstart,0), N) #max/min construction here ensures range will not exceed ends of polymer, noting that TSSloadstart can be + or -
            loadend=max(min(nn+TSSloadend+1,N),0)
        else:
            loadbegin= min(max(nn-TSSloadend,0),N)
            loadend= max(min(nn+TSSloadstart+1,N),0)
        for jj in range(loadbegin,loadend):
            birthArray[jj] *= TSSloadbias
    deathArray = np.zeros(N, dtype=np.double) + 1. / (LIFETIME)
    collisionDeathArray= np.zeros(N,dtype=np.double) + 1./collisionLife
    if not (TTSunload==1.):
        for nn,tt,gg in zip(polloading, poltermination, genelength):
            if tt>nn:
                zonebegin=nn+gg
                zoneend=min(nn+gg+unloadZone,N)
            else:
                zonebegin= max(nn-gg+1-unloadZone,0)
                zoneend= nn-gg+1
            for jj in range(zonebegin, zoneend):
                deathArray[jj] = 1./(TTSunload*LIFETIME)
                collisionDeathArray[jj] =  1./(TTSunload*LIFETIME)

    stallLeftArray = np.zeros(N, dtype=np.double) #stall prob for left LEF legs, i.e. CTCF sites that point rightward
    stallRightArray = np.zeros(N, dtype=np.double)
    if STALL_FROM_FILE: #read in stall rates from file; file should provide list of LEF stalling rates starting from the TSS, extending as long as the configuration of repeats permits
        with open("stallfiles/"+stallfile,"r") as infile:
            stall_list = infile.readlines()
        for entry in stall_list:
            gene_stall.append(float(entry.split()[0]))
        for nn,tt in zip(polloading,poltermination):
            ii=0
            if nn<tt:
                gene_direc=1
            else:
                gene_direc=-1
            for ii in range(len(gene_stall)):
                stallLeftArray[nn+ii*gene_direc]=gene_stall[ii]
                stallRightArray[nn+ii*gene_direc]=gene_stall[ii]
    else: #stall rates set by stall params & CTCF position options
        stallLeftArray[ctcf_left_list] = stg
        stallRightArray[ctcf_right_list] = stg

    unstallLEFArray = np.ones(N, dtype=np.double) * unstallLEFRate # generally have a single unstall rate.  doesn't matter if non-stall sites have unstall>0, since there won't be stalls there anyway
    stallDeathArray = np.zeros(N, dtype=np.double) + 1. / (LIFETIME) # unbinding rate during stall can be different, but we usually leave it the same
    smcNum = int(N / SEPARATION)
    curPos = 0

    SMCTran = smcTranslocatorDirectional(birthArray, deathArray, stallLeftArray, stallRightArray, unstallLEFArray, pauseArray, stallDeathArray, smcNum,
                                         kinPol,kterPolArray,pauseArrayPol,shrinkPauseArray, polloading,poltermination, stalProbPol, unstallArray,
                                         PolPermL=permLeftArray, PolPermR=permRightArray, collisionFalloffProb=collisionDeathArray,
                                         poldissoc= poldissoc,
                                         LefPerm=lefperm,
                                         strongCTCF=strongCTCFstall
                                        )
    return SMCTran

SMCTran = initModel()  # defining actual smc translocator object

# now polymer simulation code starts

# ------------feed smcTran to the milker---
SMCTran.steps(1000000)  # first steps to "equilibrate" SMC dynamics. If desired of course.
milker = smcTranslocatorMilker(SMCTran)   # now feed this thing to milker (do it once!)
#--------- end new code ------------

for milkerCount in range(milkerInitsTotal):
    doSave = milkerCount >= milkerInitsSkip

    # simulation parameters are defined below
    a = Simulation(timestep=80, thermostat=0.001)#Collision rate in inverse picoseconds, low collistion rate means ballistic like motion, default in openmmpolymer is 0.001. Motion polymer is not diffusive, this is ok for statistical average,
    #but not for dynamics of the polymer
    a.setup(platform="CUDA", PBC=True, PBCbox=[box, box, box], GPU=GPU, precision="mixed")  # set up GPU here, PBC=Periodic Boundary Conditions. Default integrator is langevin with 300 K, friction coefficient of 1/ps, step size 0.002ps
    a.saveFolder(FullFileName)
    a.load(data)
    a.addHarmonicPolymerBonds(wiggleDist=0.1) # WiggleDist controls distance at which energy of bond equals kT
    if stiff > 0:
        a.addGrosbergStiffness(stiff) # Chain stiffness is introduced by an angular potential U(theta)=stiff(1-cos(theta-Pi))
    a.addPolynomialRepulsiveForce(trunc=1.5, radiusMult=1.05) #Polynomial repulsive potential between particles. Has value trunc=3.0 at zero, stays flat until 0.6-0.7 and then drops to zero. For attraction between a selective set of particles, use LeonardJones or addSelectiveSSWForce (see blocks.py or ask Johannes)
    a.step = block

    # ------------ initializing milker; adding bonds ---------
    # copied from addBond
    kbond = a.kbondScalingFactor / (smcBondWiggleDist ** 2)
    bondDist = smcBondDist * a.length_scale

    activeParams = {"length":bondDist,"k":kbond}
    inactiveParams = {"length":bondDist, "k":0}
    milker.setParams(activeParams, inactiveParams)

    # this step actually puts all bonds in and sets first bonds to be what they should be
    milker.setup(bondForce=a.forceDict["HarmonicBondForce"],
                 blocks=restartMilkerEveryBlocks,   # default value; milk for 100 blocks
                 smcStepsPerBlock=smcStepsPerBlock)  #
    print("Restarting milker")

    a.doBlock(steps=steps, increment=False)  # do block for the first time with first set of bonds in
    #print('done 1')
    for i in range(restartMilkerEveryBlocks - 1):
        #print(i)
        curBonds, pastBonds = milker.step(a.context)  # this updates bonds. You can do something with bonds here
        if i % saveEveryBlocks == (saveEveryBlocks - 2):
            a.doBlock(steps=steps, increment = doSave)
            if doSave:
                a.save()
                pickle.dump(curBonds, open(os.path.join(a.folder, "SMC{0}.dat".format(a.step)),'wb'))
                save_Es_ts_Rg() # save energies and time
                if saveRNAP:
                    pickle.dump(milker.getRNAP(), open(os.path.join(a.folder, "RNAP{0}.dat".format(a.step)), 'wb'))
        else:
            a.integrator.step(steps)  # do steps without getting the positions from the GPU (faster)

    data = a.getData()  # save data and step, and delete the simulation
    block = a.step
    del a

    time.sleep(0.2)  # wait 200ms for sanity (to let garbage collector do its magic)

with open(Par_fname,"a+") as Parfile:
    Parfile.write(" tau="+str(LIFETIME)+"\n Separation="+str(SEPARATION)+"\n N="+str(N)+"\n smcStepsPerBlock="+str(smcStepsPerBlock)+"\n stiff="+str(stiff)+"\n dens="+str(dens)+"\n block="+str(block)+"\n  SaveEveryBlocks="+str(saveEveryBlocks)+"\n skipSavedBlocksBeginning="+str(skipSavedBlocksBeginning)+"\n totalSavedBlocks="+str(totalSavedBlocks)+"\n restartMilkerEveryBlocks="+str(restartMilkerEveryBlocks)+"\n smcBondWiggleDist="+str(smcBondWiggleDist)+"\n smcBondDist="+str(smcBondDist)+"\n SmcTimestep=1\n NumMonos"+str(N))

os.system("mv {0} {1}".format(logname, FullFileName))


import logging
import os
import pickle
import uuid
from collections import namedtuple
from glob import glob
from typing import Tuple

import jsonpickle
import numpy as np
import python_on_whales
import requests
import time
from PIL import Image
from emissor.representation.annotation import AnnotationType
from emissor.representation.entity import Gender, Person, Object
from emissor.representation.ldschema import emissor_dataclass
from emissor.representation.scenario import ImageSignal, Mention, Annotation

from .cv.plots import Annotator, Colors

logging.basicConfig(
    level=os.environ.get("LOGLEVEL", "INFO").upper(),
    format="%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)

def cosine_similarity(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    """Compute the cosine similarity of the two vectors.

    Args
    ----
    x, y: vectors

    Returns
    -------
    similarity: a similarity score between -1 and 1, where 1 is the most
        similar.

    """

    similarity = np.dot(x, y) / (np.sqrt(np.dot(x, x)) * np.sqrt(np.dot(y, y)))

    return similarity

def start_docker_container(
    image: str, port_id: int, sleep_time=5
) -> python_on_whales.Container:
    """Start docker container given the image name and port number.

    Args
    ----
    image: docker image name
    port_id: port id
    sleep_time: warmup time

    Returns
    -------
    container: a docker container object.

    """
    container = python_on_whales.docker.run(
        image=image, detach=True, publish=[(port_id, port_id)]
    )

    logging.info(f"starting a {image} container ...")
    logging.debug(f"warming up the container ...")
    time.sleep(sleep_time)

    return container

def kill_container(container: python_on_whales.Container) -> None:
    """Kill docker container.

    Args
    ----
    container:
        a docker container object.

    """
    container.kill()
    logging.info(f"container killed.")
    logging.info(f"DONE!")

def unpickle(path: str):
    """Unpickle the pickled file, and return it.

    Args
    ----
    path: path to the pickle

    Returns
    -------
    returned: un unpickled object to be returned.

    """
    with open(path, "rb") as stream:
        returned = pickle.load(stream)

    return returned

def load_embeddings(paths: str = "./friend_embeddings/*.pkl") -> dict:
    """Load pre-defined face embeddings.

    Args
    ----
    paths: paths to the face embedding vectors.

    Returns
    -------
    embeddings_predefined: predefined embeddings

    """
    embeddings_predefined = {}
    for path in glob(paths):
        name = path.split("/")[-1].split(".pkl")[0]
        unpickled = unpickle(path)

        embeddings_predefined[unpickled["uuid"]] = {
            "embedding": unpickled["embedding"],
            "name": name,
        }

    return embeddings_predefined

def face_recognition(
    friends_path: str, embeddings: list, COSINE_SIMILARITY_THRESHOLD=0.65
) -> list:
    """Perform face recognition based on the cosine similarity.

    Args
    ----
    embeddings: a list of embeddings
    COSINE_SIMILARITY_THRESHOLD: Currently fixed to The cosine similarity
        threshold is fixed to 0.65. Feel free to play around with this number.

    Returns
    -------
        faces_detected: list of faces (uuids and names) detected.

    """
    embeddings_predefined = load_embeddings(friends_path + "/*.pkl")

    cosine_similarities = []
    for embedding in embeddings:
        cosine_similarities_ = {
            uuid_
            + " "
            + embedding_name["name"]: cosine_similarity(
                embedding, embedding_name["embedding"]
            )
            for uuid_, embedding_name in embeddings_predefined.items()
        }
        cosine_similarities.append(cosine_similarities_)

    logging.debug(f"cosine similarities: {cosine_similarities}")

    faces_detected_ = [max(sim, key=sim.get) for sim in cosine_similarities]
    faces_detected = []
    for uuid_name, sim in zip(faces_detected_, cosine_similarities):
        uuid_, name = uuid_name.split()
        if sim[uuid_name] > COSINE_SIMILARITY_THRESHOLD:
            faces_detected.append({"uuid": uuid_, "name": name})
        else:
            logging.info("new face!")
            faces_detected.append({"uuid": str(uuid.uuid4()), "name": None})
            pass

    return faces_detected

def load_binary_image(image_path: str) -> bytes:
    """Load encoded image as a binary string and return it.

    Args
    ----
    image_path: path to the image to load.

    Returns
    -------
    binary_image: encoded binary image in bytes

    """
    logging.debug(f"{image_path} loading image ...")
    with open(image_path, "rb") as stream:
        binary_image = stream.read()
    logging.info(f"{image_path} image loaded!")

    return binary_image

def run_face_api(to_send: dict, url_face: str = "http://127.0.0.1:10002/") -> tuple:
    """Make a RESTful HTTP request to the face API server.

    Args
    ----
    to_send: dictionary to send to the server. In this function, this will be
        encoded with jsonpickle. I know this is not conventional, but encoding
        and decoding is so easy with jsonpickle somehow.
    url_face: the url of the face recognition server.

    Returns
    -------
    face_bboxes: (list) boudning boxes
    det_scores: (list) detection scores
    landmarks: (list) landmarks
    embeddings: (list) face embeddings
    """
    logging.debug(f"sending image to server...")
    to_send = jsonpickle.encode(to_send)
    response = requests.post(url_face, json=to_send)
    logging.info(f"got {response} from server!...")

    response = jsonpickle.decode(response.text)

    face_detection_recognition = response["face_detection_recognition"]
    logging.info(f"{len(face_detection_recognition)} faces deteced!")

    face_bboxes = [fdr["bbox"] for fdr in face_detection_recognition]
    det_scores = [fdr["det_score"] for fdr in face_detection_recognition]
    landmarks = [fdr["landmark"] for fdr in face_detection_recognition]

    embeddings = [fdr["normed_embedding"] for fdr in face_detection_recognition]

    return face_bboxes, det_scores, landmarks, embeddings

def run_age_gender_api(
    embeddings: list, url_age_gender: str = "http://127.0.0.1:10003/"
) -> tuple:
    """Make a RESTful HTTP request to the age-gender API server.

    Args
    ----
    embeddings: a list of embeddings. The number of elements in this list is
        the number of faces detected in the frame.
    url_age_gender: the url of the age-gender API server.

    Returns
    -------
    ages: (list) a list of ages
    genders: (list) a list of genders.

    """
    # -1 accounts for the batch size.
    data = np.array(embeddings).reshape(-1, 512).astype(np.float32)
    data = pickle.dumps(data)

    data = {"embeddings": data}
    data = jsonpickle.encode(data)
    logging.debug(f"sending embeddings to server ...")
    response = requests.post(url_age_gender, json=data)
    logging.info(f"got {response} from server!...")

    response = jsonpickle.decode(response.text)
    ages = response["ages"]
    genders = response["genders"]

    return ages, genders

def run_yolo_api(to_send: dict, url_yolo: str = "http://127.0.0.1:10004/") -> list:
    """Make a RESTful HTTP request to the face API server.

    Args
    ----
    to_send: dictionary to send to the server. In this function, this will be
        encoded with jsonpickle. I know this is not conventional, but encoding
        and decoding is so easy with jsonpickle somehow.
    url_yolo: the url of the YOLO server.

    Returns
    -------
    results: yolo results. Each element in this list is a dictionary. e.g.,
        {'yolo_bbox': [752, 46, 1148, 716], 'det_score': 0.875, 'label_num': 0,
        'label_string': 'person'}

    """
    logging.debug(f"Running yolo ...")
    logging.debug(f"sending image to server at {url_yolo}...")

    to_send = jsonpickle.encode(to_send)
    print("url", url_yolo)
    response = requests.post(url_yolo, json=to_send)
    logging.info(f"got {response} from server!...")
    response = jsonpickle.decode(response.text)
    results = response["yolo_results"]

    for result in results:
        result["yolo_bbox"] = result.pop("bbox")
    results = [{key: val for key, val in result.items()} for result in results]

    return results

def annotate_yolo(image: Image.Image, yolo_results: list) -> Image.Image:
    """Annotate YOLO Image.

    Args
    ----
    image: PIL image object.
    yolo_results: yolo prediction results.

    Returns
    -------
    image_annotated: Annotated PIL image object.

    """
    logging.debug("Annotating yolo image ...")
    annotator = Annotator(np.ascontiguousarray((image)))
    colors = Colors()  # create instance for 'from utils.plots import colors'

    for result in yolo_results:
        box = result["yolo_bbox"]
        label_num = result["label_num"]
        label_string = result["label_string"]

        color = colors(label_num)
        annotator.box_label(box, label_string, color=color)

    image_annotated = Image.fromarray(annotator.im)
    logging.info(f"YOLO image annotation is done!")

    return image_annotated

def annotate_face(
    image: Image.Image, genders, ages, face_bboxes, faces_detected, det_scores
):
    """Annotate face nicely.

    Args
    ----

    Returns
    -------

    """
    logging.debug("Annotating face, genders, and ages ...")

    assert (
        len(genders)
        == len(ages)
        == len(face_bboxes)
        == len(faces_detected)
        == len(det_scores)
    )
    annotator = Annotator(np.ascontiguousarray((image)))
    colors = Colors()  # create instance for 'from utils.plots import colors'

    for gender, age, face_bbox, uuid_name, faceprob in zip(
        genders, ages, face_bboxes, faces_detected, det_scores
    ):
        box = face_bbox.tolist()
        if gender["m"] > gender["f"]:
            binary_gender = "male"
        else:
            binary_gender = "female"

        try:
            short_name = str(uuid_name["name"].split("_")[0])
        except:
            short_name = str(uuid_name["name"])

        label_string = (
            f"{short_name}, " f"{round(age['mean'])} years old, " f"{binary_gender}."
        )

        color = colors(81)
        annotator.box_label(box, label_string, color=color)

    image_annotated = Image.fromarray(annotator.im)
    logging.info("Annotating face, genders, and ages is done!")

    return image_annotated

FaceInfo = namedtuple(
    "FaceInfo",
    ("gender", "age", "bbox", "face_id", "det_score", "embedding", "yolo_result"),
)

def detect_faces(
    friends_path: str,
    image_path: str,
    url_face: str = "http://127.0.0.1:10002/",
    url_age_gender: str = "http://127.0.0.1:10003/",
    url_yolo: str = "http://127.0.0.1:10004",
) -> Tuple[FaceInfo]:
    """Detect faces in an image.

    Args
    ----
    image_path: path to the image in disk
    url_face: the url of the face recognition server.
    url_age_gender: the url of the age-gender API server.
    url_yolo: the url of the YOLO5 API server

    Returns
    -------
    List[FaceInfo]
    """
    MAXIMUM_ENTROPY = {"gender": 0.6931471805599453, "age": 4.615120516841261}

    data = {"image": load_binary_image(image_path)}

    face_bboxes, det_scores, landmarks, embeddings = run_face_api(data, url_face)

    faces_detected = face_recognition(friends_path, embeddings)

    ages, genders = run_age_gender_api(embeddings, url_age_gender)

    yolo_results = run_yolo_api(data, url_yolo)

    logging.debug("annotating image ...")
    image = Image.open(image_path)
    image = annotate_yolo(image, yolo_results)
    image = annotate_face(image, genders, ages, face_bboxes, faces_detected, det_scores)
    logging.info("image annotation is done!")

    image_path = image_path + ".ANNOTATED.jpg"
    logging.debug(f"saving image at {image_path}...")
    image.save(image_path)
    logging.info(f"image saved at {image_path}")

    return tuple(
        FaceInfo(*info)
        for info in zip(
            genders,
            ages,
            face_bboxes,
            faces_detected,
            det_scores,
            embeddings,
            yolo_results,
        )
    )

def detect_objects(
    image_path: str,
    url_yolo: str = "http://127.0.0.1:10004",
):
    """Detect objects in an image.

    Args
    ----
    image_path: path to the image in disk
    url_yolo: the url of the YOLO5 API server

    Returns
    -------
    List[ObjectInfo]
    """

    data = {"image": load_binary_image(image_path)}

    yolo_results = run_yolo_api(data, url_yolo)

    logging.debug("annotating image ...")
    image = Image.open(image_path)
    image = annotate_yolo(image, yolo_results)
    logging.info("image annotation is done!")

    image_path = image_path + ".ANNOTATED.jpg"
    logging.debug(f"saving image at {image_path}...")
    image.save(image_path)
    logging.info(f"image saved at {image_path}")

    return yolo_results

def create_face_mention(
    image_signal: ImageSignal,
    source: str,
    current_time: int,
    bbox: Tuple[int, int, int, int],
    uri: str,
    name: str,
    age: str,
    gender: str,
    face_prob: float,
) -> Mention:
    bbox = [
        max(x, lower) for x, lower in zip(bbox[:2], image_signal.ruler.bounds[:2])
    ] + [min(x, upper) for x, upper in zip(bbox[-2:], image_signal.ruler.bounds[-2:])]
    face_segment = image_signal.ruler.get_area_bounding_box(*bbox)
    face_annotation = Annotation(
        AnnotationType.PERSON.name,
        FacePerson(uri, name, age, Gender[gender.upper()], face_prob),
        source,
        current_time,
    )

    return Mention(str(uuid.uuid4()), [face_segment], [face_annotation])

def create_object_mention(
    image_signal: ImageSignal,
    source: str,
    current_time: int,
    bbox: Tuple[int, int, int, int],
    name: str,
    obj_prob: float,
) -> Mention:
    bbox = [
        max(x, lower) for x, lower in zip(bbox[:2], image_signal.ruler.bounds[:2])
    ] + [min(x, upper) for x, upper in zip(bbox[-2:], image_signal.ruler.bounds[-2:])]
    object_segment = image_signal.ruler.get_area_bounding_box(*bbox)
    object_annotation = Annotation(
        AnnotationType.OBJECT.name,
        ImageObject(str(uuid.uuid4()), name, obj_prob),
        source,
        current_time,
    )

    return Mention(str(uuid.uuid4()), [object_segment], [object_annotation])

@emissor_dataclass(namespace="http://cltl.nl/leolani/n2mu")
class FacePerson(Person):
    face_prob: float

@emissor_dataclass(namespace="http://cltl.nl/leolani/n2mu")
class ImageObject(Object):
    obj_prob: float


import sys
import logging
import os.path as op
import pandas as pd
import numpy as np
from tqdm import tqdm
from copy import copy
from sklearn.preprocessing import OneHotEncoder, LabelEncoder
from sklearn.model_selection import train_test_split

sys.path.append('src')
from noise_ceiling import convert_doubles_to_single_labels
from utils import argmax_with_random_ties

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(funcName)-8.8s] [%(levelname)-7.7s]  %(message)s",
    datefmt='%Y-%m-%d %H:%M:%S',
    handlers=[
        logging.StreamHandler()
    ]
)

DU2EN = dict(
    Verrassing='surprise',
    Verdrietig='sadness',
    Bang='fear',
    Walging='disgust',
    Blij='happiness',
    Boos='anger'
)

class DataLoader:

    def __init__(self, sub='01', data_dir=None, rnd_seed=42, log_level=20):
        """ Initializes a DataLoader object.

        Parameters
        ----------
        sub : str
            Subject ID (zero-padded)
        data_dir : str
            Path to data directory
        rnd_seed : int
            Random seed for train-test set split
        """
        self.sub = sub
        self.y = None
        self.X = None
        self.target_name = None

        if data_dir is None:
            data_dir = op.abspath('data')
            if not op.isdir(data_dir):
                raise ValueError(f"Directory {data_dir} does not exist.")

        self.data_dir = data_dir
        self.rnd_seed = rnd_seed
        self.le = LabelEncoder().fit(['happiness', 'surprise', 'fear', 'sadness', 'disgust', 'anger'])
        self.ohe = OneHotEncoder(
            categories=self.le.classes_,
            sparse=False
        )
        self.log = logging.getLogger(__name__)
        self.log.setLevel(log_level)

    def load_y(self, target='emotion', data_split='train', filter_gva=True, strategy_doubles='soft', return_doubles=False):
        """ Loads the target variable (y).

        Parameters
        ----------
        target : str
            Name of target variable ("emotion", "valence", or "arousal")
        data_split : str
            Either "train" or "test"
        filter_gva : bool
            Whether to remove the "geen van allen" ratings
        strategy_doubles : str
            Strategy to handle doubles (either 'soft', 'hard', or 'none').
        """

        self.target_name = target

        f = op.join(self.data_dir, 'ratings_complete', f'sub-{self.sub}_task-expressive_ratings.tsv')
        df = pd.read_csv(f, sep='\t', index_col=0)
        df = df.query("rating_type == @target")  # filter rating type

        n_orig = df.shape[0]
        df = df.query("data_split == @data_split")
        self.log.info(f"Removed {n_orig - df.shape[0]} test trials (N = {df.shape[0]}).")

        if filter_gva:  # remove "geen van allen" (none of all)
            n_orig = df.shape[0]
            df = df.query("rating != 'Geen van allen'")
            n_remov = n_orig - df.shape[0]
            self.log.info(f"Removed {n_remov} 'Geen van allen' trials (N = {df.shape[0]}).")

        with pd.option_context('mode.chained_assignment', None):  # suppress stupid warning
            df.loc[:, 'rating'] = df.loc[:, 'rating'].replace(DU2EN)  # translate
            df.loc[:, 'rating'] = self.le.transform(df['rating'])

        # Handle doubles for session 1
        if return_doubles:
            dup_idx = df.index[df.index.duplicated()].unique()
            dups = df.loc[dup_idx, :].sort_index()
            return dups['rating']

        if strategy_doubles in ['hard', 'soft']:  # stupid slow code
            soft = True if strategy_doubles == 'soft' else False
            df = convert_doubles_to_single_labels(df['rating'], soft=soft, keepdims=False)

            if strategy_doubles == 'hard':
                # Isn't this redundant, because I can do convert_doubles ... with soft=False?
                df = pd.DataFrame(df.values.argmax(axis=1), columns=['rating'], index=df.index)

            """
            n_orig = df.shape[0]
            df = df.drop(dups.index, axis=0)  # remove duplicates

            filt_dups = []
            for i, idx in enumerate(dup_idx):
                tmp = dups.loc[idx, :]
                counts = tmp['rating'].value_counts()
                if strategy_doubles == 'argmax':
                    maxcount = counts.loc[counts == counts.max()].sample(n=1).index[0]
                    filt_dups.append(tmp.loc[tmp['rating'] == maxcount, :].sample(n=1))

            filt_dups = pd.concat(filt_dups, axis=0)
            df = pd.concat((df, filt_dups), axis=0)
            """
            self.log.info(f"Removed {n_orig - df.shape[0]} duplicates (N = {df.shape[0]}).")

        if strategy_doubles == 'soft':
            # WARNING: source of randomness
            df_tmp = pd.DataFrame(argmax_with_random_ties(df.values), columns=['rating'], index=df.index)

            df_tmp_train, df_test = train_test_split(df, test_size=0.15, random_state=self.rnd_seed, stratify=df_tmp['rating'])
            df = df.loc[df_tmp_train.index, :]
            self.y = df['rating']
            #self.intensity = df['intensity']
        else:
            df, df_test = train_test_split(df, test_size=0.15, random_state=self.rnd_seed, stratify=df['rating'])
            self.y = df['rating']
            #self.intensity = df['intensity']

        self.log.info(f"Split into train (N = {df.shape[0]}) and test (N = {df_test.shape[0]}).")
        self.rating_df = df

    def load_X(self, feature_set, n_comp=50):
        """ Loads in predictors/independent variables (X).

        Parameters
        ----------
        feature_set : str/list/tuple
            Name of one or more feature-sets to use as predictors
        n_comp : int
            Number of components
        """

        if not isinstance(feature_set, (list, tuple)):
            feature_set = (feature_set,)

        if self.y is None:
            raise ValueError("Call load_y before load_X!")

        X = []
        for fs in feature_set:  # load in 1 or more feature-sets
            self.log.info(f"Loading feature-set {fs}.")
            path = op.join(self.data_dir, f'featurespace-{fs}.tsv')
            df = pd.read_csv(path, sep='\t', index_col=0)

            if 'sub' in df.columns:
                df = df.query(f"sub == 'sub-{self.sub}'").drop('sub', axis=1)

            """
            if 'rep' in df.columns:
                tmp_index = self.rating_df.index + '_' + self.rating_df['rep'].astype(str)
                df.index = df.index + '_' + df['rep'].astype(str)
            else:
                tmp_index = self.rating_df.index
            """
            #df = df.loc[tmp_index, :].set_index(self.rating_df.index)
            df = df.loc[self.rating_df.index, :].set_index(self.rating_df.index)

            if 'data_split' in df.columns:
                df = df.drop('data_split', axis=1)

            if 'pca' in fs or 'nmf' in fs:
                df = df.iloc[:, :n_comp]

            X.append(df)

        X = pd.concat(X, axis=1)

        self.X = X
        self.log.info(f"Shape X: {self.X.shape}.")

    def return_Xy(self):
        """ Returns the labels (y) and predictors (X). """
        return self.X, self.y


from __future__ import absolute_import, division, print_function, unicode_literals
__metaclass__ = type

import random
import os
import math

import h5py
import cv2
import numpy as np
import sqlite3

from .util import static_vars
from .task import DEBUG
from .google_storage import downloadIfAvailable

POSITIVE_IMAGE_DATABASE_FOLDER = 'aflw/data/'
POSITIVE_IMAGE_FOLDER = POSITIVE_IMAGE_DATABASE_FOLDER + 'flickr/'
POSITIVE_IMAGE_DATABASE_FILE = os.path.join(POSITIVE_IMAGE_DATABASE_FOLDER, 'aflw.sqlite')
OBJECT_DATABASE_PATHS = ('data/pos12.hdf', 'data/pos24.hdf', 'data/pos48.hdf')
TEST_IMAGE_DATABASE_FOLDER = 'Annotated Faces in the Wild/FDDB-folds/'
TEST_IMAGES_FOLDER = 'Annotated Faces in the Wild/originalPics'

DATASET_LABEL = 'data'
LABELS_LABEL = 'labels'
CHUNK_SIZE = 256

NEGATIVE_IMAGE_FOLDER = 'Negative Images/images/'
NEGATIVE_DATABASE_PATHS = ('data/neg12.hdf', 'data/neg24.hdf', 'data/neg48.hdf')
TARGET_NUM_NEGATIVES_PER_IMG = 40
TARGET_NUM_NEGATIVES = 300000
MIN_FACE_SCALE = 80
OFFSET = 4

SCALES = ((12,12), (24,24), (48,48))

CALIBRATION_DATABASE_PATHS = {SCALES[0][0]:'data/calib12.hdf', SCALES[1][0]:'data/calib24.hdf', SCALES[2][0]:'data/calib48.hdf'}
TARGET_NUM_CALIBRATION_SAMPLES = 337500
SN = (.83, .91, 1, 1.1, 1.21)
XN = (-.17, 0, .17)
YN = XN
CALIB_PATTERNS = [(sn, xn, yn) for sn in SN for xn in XN for yn in YN]
CALIB_PATTERNS_ARR = np.asarray(CALIB_PATTERNS)

RANDOM_SEED = 42
np.random.seed(RANDOM_SEED)

def numDetectionWindowsAlongAxis(size, stageIdx = 0):
    return (size-SCALES[stageIdx][0])//OFFSET+1

@static_vars(faces = [])
def getFaceAnnotations(dbPath = POSITIVE_IMAGE_DATABASE_FILE, posImgFolder = POSITIVE_IMAGE_FOLDER):
    if len(getFaceAnnotations.faces) == 0:
        with sqlite3.connect(dbPath) as conn:
            c = conn.cursor()

            select_string = "faceimages.filepath, facerect.x, facerect.y, facerect.w, facerect.h"
            from_string = "faceimages, faces, facepose, facerect"
            where_string = "faces.face_id = facepose.face_id and faces.file_id = faceimages.file_id and faces.face_id = facerect.face_id"
            query_string = "SELECT " + select_string + " FROM " + from_string + " WHERE " + where_string

            for row in c.execute(query_string):
                imgPath = os.path.join(posImgFolder, str(row[0]))

                if os.path.isfile(imgPath):
                    getFaceAnnotations.faces.append((imgPath,) + (row[1:]))

    return getFaceAnnotations.faces

def squashCoords(img, x, y, w, h):
    y = min(max(0, y), img.shape[0])
    x = min(max(0, x), img.shape[1])
    h = min(img.shape[0]-y, h)
    w = min(img.shape[1]-x, w)
    return (x, y, w, h)

def createPositiveDataset(stageIdx, getObjectAnnotations = getFaceAnnotations, posImgFolder = POSITIVE_IMAGE_FOLDER):
    fileName = OBJECT_DATABASE_PATHS[stageIdx]
    resizeTo = SCALES[stageIdx]
    objectAnnotations = getObjectAnnotations(posImgFolder = posImgFolder)
    images = np.zeros((len(objectAnnotations), resizeTo[1], resizeTo[0], 3), dtype = np.uint8)
    curImg = None
    prevImgPath = None

    for i, (imgPath, x, y, w, h) in enumerate(objectAnnotations):
        if imgPath != prevImgPath:
            curImg = cv2.imread(imgPath)

        x, y, w, h = squashCoords(curImg, x, y, w, h)
        images[i] = cv2.resize(curImg[y:y+h,x:x+w], resizeTo)
        prevImgPath = imgPath

    with h5py.File(fileName, 'w') as out:
        out.create_dataset(DATASET_LABEL, data = images, chunks = (CHUNK_SIZE,) + (images.shape[1:]))

def createNegativeDataset(stageIdx, negImgFolder = NEGATIVE_IMAGE_FOLDER, numNegatives = TARGET_NUM_NEGATIVES, numNegativesPerImg = TARGET_NUM_NEGATIVES_PER_IMG):
    fileName = NEGATIVE_DATABASE_PATHS[stageIdx]
    resizeTo = SCALES[stageIdx]
    negativeImagePaths = [os.path.join(negImgFolder, fileName) for fileName in os.listdir(negImgFolder)]
    images = np.zeros((numNegatives, resizeTo[1], resizeTo[0], 3), dtype = np.uint8)
    negIdx = 0
    numNegativesRetrievedFromImg = 0

    for i in np.random.permutation(len(negativeImagePaths)):
        if negIdx >= numNegatives: break
        img = cv2.resize(cv2.imread(negativeImagePaths[i]), None, fx = resizeTo[0]/MIN_FACE_SCALE, fy = resizeTo[1]/MIN_FACE_SCALE)

        for xOffset in np.random.permutation(numDetectionWindowsAlongAxis(img.shape[1], stageIdx)):
            if negIdx >= numNegatives or numNegativesRetrievedFromImg >= numNegativesPerImg: break

            for yOffset in np.random.permutation(numDetectionWindowsAlongAxis(img.shape[0], stageIdx)):
                if negIdx >= numNegatives or numNegativesRetrievedFromImg >= numNegativesPerImg: break
                x, y, w, h = squashCoords(img, resizeTo[0]*xOffset, resizeTo[1]*yOffset, *resizeTo)

                if (w == resizeTo[0] and h == resizeTo[1]):
                    images[negIdx] = img[y:y+h,x:x+w]
                    negIdx += 1
                    numNegativesRetrievedFromImg += 1

        numNegativesRetrievedFromImg = 0

    if negIdx < len(images)-1:
        images = np.delete(images, np.s_[negIdx:], 0)

    with h5py.File(fileName, 'w') as out:
        out.create_dataset(DATASET_LABEL, data = images, chunks = (CHUNK_SIZE,) + (images.shape[1:]))

def mineNegatives(stageIdx, negImgFolder = NEGATIVE_IMAGE_FOLDER, numNegatives = TARGET_NUM_NEGATIVES):
    from .detect import detectMultiscale

    fileName = NEGATIVE_DATABASE_PATHS[stageIdx]
    resizeTo = SCALES[stageIdx]
    negativeImagePaths = [os.path.join(negImgFolder, fileName) for fileName in os.listdir(negImgFolder)]
    images = np.zeros((numNegatives, resizeTo[1], resizeTo[0], 3), dtype = np.uint8)
    negIdx = 0

    for i in np.random.permutation(len(negativeImagePaths)):
        if negIdx >= numNegatives: break
        img = cv2.imread(negativeImagePaths[i])
        coords = detectMultiscale(img, stageIdx-1)

        for xMin, yMin, xMax, yMax in coords:
            if negIdx >= numNegatives: break
            xMin, yMin, w, h = squashCoords(img, xMin, yMin, xMax-xMin, yMax-yMin)
            images[negIdx] = cv2.resize(img[yMin:yMin+h, xMin:xMin+w], resizeTo)
            negIdx += 1

    if negIdx < len(images)-1:
        images = np.delete(images, np.s_[negIdx:], 0)

    with h5py.File(fileName, 'w') as out:
        out.create_dataset(DATASET_LABEL, data = images, chunks = (CHUNK_SIZE,) + (images.shape[1:]))

def createCalibrationDataset(stageIdx, getObjectAnnotations = getFaceAnnotations, posImgFolder = POSITIVE_IMAGE_FOLDER, numCalibrationSamples = TARGET_NUM_CALIBRATION_SAMPLES, calibPatterns = CALIB_PATTERNS):
    numCalibrationSamples = math.inf if numCalibrationSamples is None else numCalibrationSamples
    objectAnnotations = getObjectAnnotations(posImgFolder = posImgFolder)

    resizeTo = SCALES[stageIdx]
    datasetLen = len(objectAnnotations)*len(calibPatterns) if numCalibrationSamples == math.inf else numCalibrationSamples
    dataset = np.zeros((datasetLen, resizeTo[1], resizeTo[0], 3), np.uint8)
    labels = np.zeros((datasetLen, 1))
    sampleIdx = 0

    fileName = CALIBRATION_DATABASE_PATHS.get(resizeTo[0])

    curImg = None
    prevImgPath = None

    for i, (imgPath, x, y, w, h) in enumerate(objectAnnotations):
        if sampleIdx >= numCalibrationSamples: break

        if imgPath != prevImgPath:
            curImg = cv2.imread(imgPath)

        for n, (sn, xn, yn) in enumerate(CALIB_PATTERNS):
            if sampleIdx >= numCalibrationSamples: break
            box = squashCoords(curImg, x + int(xn*w), y + int(yn*h), int(w*sn), int(h*sn))

            if box[2] > 0 and box[3] > 0:
                (box_x, box_y, box_w, box_h) = box
                dataset[sampleIdx] = cv2.resize(curImg[box_y:box_y+box_h,box_x:box_x+box_w], resizeTo)
                labels[sampleIdx] = n
                sampleIdx += 1

    if sampleIdx < datasetLen:
        labels = np.delete(labels, np.s_[sampleIdx:], 0)
        dataset = np.delete(dataset, np.s_[sampleIdx:], 0)

    with h5py.File(fileName, 'w') as out:
        out.create_dataset(LABELS_LABEL, data = labels, chunks = (CHUNK_SIZE, 1))
        out.create_dataset(DATASET_LABEL, data = dataset, chunks = (CHUNK_SIZE, resizeTo[1], resizeTo[0], 3))

def getTestImagePaths(testImgDbFolder = TEST_IMAGE_DATABASE_FOLDER, testImgsFolder = TEST_IMAGES_FOLDER):
    imgPaths = []

    for fileName in os.listdir(testImgDbFolder):
        if 'ellipse' not in fileName:
            with open(os.path.join(testImgDbFolder, fileName)) as inFile:
                imgPaths.extend(inFile.read().splitlines())

    return [os.path.join(testImgsFolder, imgPath) + '.jpg' for imgPath in imgPaths]

class DatasetManager():
    normalizers = {}

    def __init__(self, model, getObjectAnnotations = getFaceAnnotations, posImgFolder = POSITIVE_IMAGE_FOLDER, negImgFolder = NEGATIVE_IMAGE_FOLDER):
        from .model import ObjectCalibrator
        self.model = model
        self.isCalib = isinstance(self.model, ObjectCalibrator)
        self.stageIdx = model.getStageIdx()
        self.posDatasetFilePath = OBJECT_DATABASE_PATHS[self.stageIdx]
        self.negDatasetFilePath = NEGATIVE_DATABASE_PATHS[self.stageIdx]
        self.calibDatasetFilePath = CALIBRATION_DATABASE_PATHS[SCALES[self.stageIdx][0]]
        self.posImgFolder = posImgFolder
        self.negImgFolder = negImgFolder

        self.getObjectAnnotations = getObjectAnnotations

    def getParams(self):
        return {'getObjectAnnotations': self.getObjectAnnotations, 'posImgFolder': self.posImgFolder, 'negImgFolder': self.negImgFolder}

    def _createFile(self, fileName, **kwargs):
        if not os.path.isfile(fileName) and not downloadIfAvailable(fileName):
            if fileName in OBJECT_DATABASE_PATHS:
                createPositiveDataset(self.stageIdx, **kwargs)
            elif fileName in  NEGATIVE_DATABASE_PATHS:
                (createNegativeDataset if self.stageIdx == 0 else mineNegatives)(self.stageIdx, **kwargs)
            elif fileName in CALIBRATION_DATABASE_PATHS.values():
                createCalibrationDataset(self.stageIdx, **kwargs)

    def getPosDatasetFilePath(self):
        self._createFile(self.posDatasetFilePath, posImgFolder = self.posImgFolder, getObjectAnnotations = self.getObjectAnnotations)
        return self.posDatasetFilePath

    def getNegDatasetFilePath(self):
        self._createFile(self.negDatasetFilePath, negImgFolder = self.negImgFolder)
        return self.negDatasetFilePath

    def getCalibDatasetFilePath(self):
        if self.isCalib: self._createFile(self.calibDatasetFilePath, posImgFolder = self.posImgFolder, getObjectAnnotations = self.getObjectAnnotations)
        return self.calibDatasetFilePath

    def getLabels(self):
        labels = None

        if self.isCalib:
            calibDatasetFilePath = self.getCalibDatasetFilePath()

            with h5py.File(calibDatasetFilePath, 'r') as calibDatasetFile:
                labels = calibDatasetFile[LABELS_LABEL][:]

        return labels

    def getNormalizer(self):
        from .preprocess import ImageNormalizer

        if DatasetManager.normalizers.get(self.model) is None:
            normalizer = ImageNormalizer(self.getPosDatasetFilePath(), self.getNegDatasetFilePath(), self.model.getNormalizationMethod())
            normalizer.addDataAugmentationParams(self.model.getNormalizationParams())
            DatasetManager.normalizers[self.model] = normalizer

        return DatasetManager.normalizers[self.model]

    def getPaths(self):
        return [self.getPosDatasetFilePath(), self.getNegDatasetFilePath()] if not self.isCalib else [self.getCalibDatasetFilePath(), None]