text:
#!/usr/bin/env python
#
# Copyright 2019 DFKI GmbH.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the
# following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
# NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
# USE OR OTHER DEALINGS IN THE SOFTWARE.
import os
import math
import numpy as np
from copy import copy
from PySide2.QtWidgets import  QDialog, QListWidgetItem, QTableWidgetItem, QTableWidget, QFileDialog
from PySide2.QtCore import QTimer, Qt
from PySide2.QtGui import QColor
from OpenGL.GL import *
from motion_analysis.gui.layout.skeleton_editor_dialog_ui import Ui_Dialog
from .utils import get_animation_controllers
from transformations import quaternion_matrix, quaternion_multiply, quaternion_about_axis
from motion_analysis.gui.widgets.scene_viewer import SceneViewerWidget
from vis_utils.animation import load_motion_from_bvh
from vis_utils.scene.editor_scene import EditorScene
from vis_utils.animation.animation_editor import AnimationEditor
from vis_utils.io import save_json_file
from anim_utils.animation_data import BVHReader, BVHWriter, MotionVector, SkeletonBuilder
from anim_utils.retargeting.analytical import create_local_cos_map_from_skeleton_axes_with_map, find_rotation_between_vectors, OPENGL_UP_AXIS
from anim_utils.animation_data.skeleton_models import STANDARD_MIRROR_MAP, STANDARD_MIRROR_MAP_LEFT, STANDARD_MIRROR_MAP_RIGHT, JOINT_CONSTRAINTS

DEFAULT_TARGET_CS_UP = [1,0,0]

def normalize(v):
    return v / np.linalg.norm(v)

def quaternion_from_vector_to_vector(a, b):
    """src: http://stackoverflow.com/questions/1171849/finding-quaternion-representing-the-rotation-from-one-vector-to-another
    http://wiki.ogre3d.org/Quaternion+and+Rotation+Primer"""

    v = np.cross(a, b)
    w = np.sqrt((np.linalg.norm(a) ** 2) * (np.linalg.norm(b) ** 2)) + np.dot(a, b)
    q = np.array([w, v[0], v[1], v[2]])
    if np.dot(q,q) != 0:
        return q/ np.linalg.norm(q)
    else:
        idx = np.nonzero(a)[0]
        q = np.array([0, 0, 0, 0])
        q[1 + ((idx + 1) % 2)] = 1 # [0, 0, 1, 0] for a rotation of 180 around y axis
        return q

def rotate_vector(q, v):
    m = quaternion_matrix(q)[:3, :3]
    v = np.dot(m, v)
    return v

X = np.array([1,0,0])
Y = np.array([0,1,0])

STANDARD_JOINTS = ["","root","pelvis", "spine", "spine_1", "spine_2", "neck", "left_clavicle", "head", "left_shoulder", "left_elbow", "left_wrist", "right_clavicle", "right_shoulder",
                    "left_hip", "left_knee", "left_ankle", "left_toe", "right_elbow", "right_wrist", "right_hip", "right_knee", "right_ankle", "right_toe"]

STANDARD_JOINTS += ["left_thumb_base","left_thumb_mid", "left_thumb_tip","left_thumb_end",
                    "left_index_finger_root","left_index_finger_base","left_index_finger_mid", "left_index_finger_tip","left_index_finger_end",
                    "left_middle_finger_root","left_middle_finger_base","left_middle_finger_mid","left_middle_finger_tip","left_middle_finger_end",
                    "left_ring_finger_root","left_ring_finger_base","left_ring_finger_mid","left_ring_finger_tip", "left_ring_finger_end",
                    "left_pinky_finger_root","left_pinky_finger_base","left_pinky_finger_mid","left_pinky_finger_tip", "left_pinky_finger_end",

                    "right_thumb_base","right_thumb_mid","right_thumb_tip","right_thumb_end",
                    "right_index_finger_root","right_index_finger_base","right_index_finger_mid","right_index_finger_tip","right_index_finger_end",
                    "right_middle_finger_root","right_middle_finger_base","right_middle_finger_mid","right_middle_finger_tip","right_middle_finger_end",
                    "right_ring_finger_root","right_ring_finger_base","right_ring_finger_mid","right_ring_finger_tip","right_ring_finger_end",
                    "right_pinky_finger_root","right_pinky_finger_base","right_pinky_finger_mid","right_pinky_finger_tip","right_pinky_finger_end"
                    ]
def find_key(joint_map,value):
    for key,v in joint_map.items():
        if value == v:
            return key
    return None

def axes_to_q(g_twist, g_swing, flip=False):
    q = [1, 0, 0, 0]
    q_y = quaternion_from_vector_to_vector(Y, g_twist)
    q_y = normalize(q_y)
    q = quaternion_multiply(q_y, q)
    X_prime = rotate_vector(q_y, X)
    X_prime = normalize(X_prime)
    q_x = quaternion_from_vector_to_vector(X_prime, g_swing)
    q_x = normalize(q_x)
    q = quaternion_multiply(q_x, q)
    q = normalize(q)
    Y_prime = rotate_vector(q, Y)
    dot = np.dot(Y_prime, g_twist)
    #dot = min(dot,1)
    dot = max(dot,-1)
    if dot == -1:
        q180 = quaternion_about_axis(np.deg2rad(180), g_swing)
        q180 = normalize(q180)
        q = quaternion_multiply(q180, q)
        q = normalize(q)
    elif abs(dot) != 1.0:
        q_y = quaternion_from_vector_to_vector(Y_prime, g_twist)
        q = quaternion_multiply(q_y, q)
        q = normalize(q)
    return q

OPENGL_UP_AXIS = [0,1,0]
OPENGL_FORWARD_AXIS = [0,0,1]
def get_axis_correction(skeleton, joint_name, up_vector, target_vector=OPENGL_UP_AXIS):
    node = skeleton.nodes[joint_name]
    t_pose_global_m = node.get_global_matrix(skeleton.reference_frame)
    global_original = np.dot(t_pose_global_m[:3, :3], up_vector)
    global_original = normalize(global_original)
    qoffset = find_rotation_between_vectors(global_original, target_vector)
    return qoffset

def get_child_index(skeleton, joint_name):
    for idx, c in enumerate(skeleton.nodes[joint_name].parent.children):
        if c.node_name == joint_name:
            return idx
    return -1

def get_traversal_map_from_parent(skeleton, joint_name, target_parent):
    traversal_map = []
    while joint_name != target_parent and skeleton.nodes[joint_name].parent is not None:
        #print("get traversal map", joint_name, target_parent, )
        idx = get_child_index(skeleton, joint_name)
        if idx == -1:
            print("Error index is none", joint_name, skeleton.nodes[joint_name].parent.children)
            break
        traversal_map.append(idx)
        joint_name = skeleton.nodes[joint_name].parent.node_name
    return list(reversed(traversal_map))

def get_joint_from_traversal_map(skeleton, traversal_map, joint_name):
    for idx in traversal_map:
        joint_name = skeleton.nodes[joint_name].children[idx].node_name
    return joint_name

def generate_sequence_from_root_to_joint(skeleton, skeleton_model, joint_key_list, joint_val_list, joint_name, known_joints, partial_mirror_map):
    joint_list = []
    while joint_name != "pelvis":
        #print("add", joint_name)
        joint_list.append(joint_name)
        skel_j = skeleton_model["joints"][joint_name]
        parent = skeleton.nodes[skel_j].parent
        while parent is not None and parent.node_name not in joint_val_list:
            parent = parent.parent
        if parent is None:
            print("Error: parent is None")
            break
        idx = joint_val_list.index(parent.node_name)
        joint_name = joint_key_list[idx]
        skel_j = skeleton_model["joints"][joint_name]
        if (len(joint_list)> 1 and joint_name in known_joints) or  joint_name not in partial_mirror_map:
            break
    return list(reversed(joint_list))

def mirror_sequence(skeleton, src_joint_list, skeleton_model):
    """ mirror sequence by corresponding joint hierarchy"""
    standard_mirror_map = STANDARD_MIRROR_MAP
    joint_map = skeleton_model["joints"]
    parent_src = src_joint_list[0]
    parent_dst = standard_mirror_map[parent_src]
    if parent_dst not in joint_map:
        return skeleton_model
    for src in src_joint_list[1:]:
        src_j = joint_map[src]
        dst = standard_mirror_map[src]
        src_parent_j = joint_map[parent_src]
        parent_dst_j = joint_map[parent_dst]
        traversal_map = get_traversal_map_from_parent(skeleton, src_j, src_parent_j)
        dst_j = get_joint_from_traversal_map(skeleton, traversal_map, parent_dst_j)
        skeleton_model["joints"][dst] = dst_j
        print("set", dst, dst_j)
        #src_cos_map = skeleton_model["cost_map"][src_j]
        #skeleton_model["cost_map"][dst_j] = src_cos_map
        parent_src = src
        parent_dst = dst

    return skeleton_model

def mirror_join_map(skeleton, skeleton_model, partial_mirror_map):
    """ for each end effector generate a sequence to the root or the last mirrored joint and then try to mirror the sequence"""
    known_joints = set()
    joint_key_list = list(skeleton_model["joints"].keys())
    joint_val_list = list(skeleton_model["joints"].values())
    for key in partial_mirror_map:
        print("try to mirror  ", key)
        if key in known_joints:
            print("seen joint already")
            continue
        if key not in skeleton_model["joints"] :
            print(key, "not in joint map")
            continue
        if skeleton_model["joints"][key] is None:
            print("joint map set to None")
            continue
        #skel_j = skeleton_model["joints"][key]
        #if len(skeleton.nodes[skel_j].children) > 0:
        #    print("too many children")
        #    continue
        sequence = generate_sequence_from_root_to_joint(skeleton, skeleton_model, joint_key_list, joint_val_list, key, known_joints, partial_mirror_map)
        print("generate ", sequence)
        for k in sequence:
            known_joints.add(k)
        skeleton_model = mirror_sequence(skeleton, sequence, skeleton_model)
    return skeleton_model

class SkeletonEditorDialog(QDialog, Ui_Dialog):
    def __init__(self, name, skeleton, share_widget, parent=None, enable_line_edit=False, skeleton_model=None):
        self.initialized = False
        QDialog.__init__(self, parent)
        Ui_Dialog.setupUi(self, self)
        self.view = SceneViewerWidget(parent, share_widget, size=(400,400))
        self.view.setObjectName("left")
        self.view.setMinimumSize(400,400)
        self.view.initializeGL()
        self.nameLineEdit.setText(name)
        self.nameLineEdit.setEnabled(enable_line_edit)
        self.name = name
        self.view.enable_mouse_interaction = True
        self.view.mouse_click.connect(self.on_mouse_click)
        self.viewerLayout.addWidget(self.view)

        self.radius = 1.5
        self.fps = 60
        self.dt = 1/60
        self.timer = QTimer()
        self.timer.timeout.connect(self.draw)
        self.timer.start(0)
        self.timer.setInterval(1000.0/self.fps)
        self.skeleton = skeleton
        self.view.makeCurrent()
        self.scene = EditorScene(True)
        self.scene.enable_scene_edit_widget = True

        if skeleton_model is not None:
            self.skeleton_model = skeleton_model
        elif skeleton.skeleton_model is not None:
            self.skeleton_model = skeleton.skeleton_model
        else:
            self.skeleton_model = dict()
            print("create new skeleton model")
        if "cos_map" not in self.skeleton_model:
            self.skeleton_model["cos_map"] = dict()
        if "joints" not in self.skeleton_model:
            self.skeleton_model["joints"] = dict()
        if "joint_constraints" not in self.skeleton_model:
            self.skeleton_model["joint_constraints"] = dict()

        motion_vector = MotionVector()
        self.reference_frame = skeleton.reference_frame
        print(self.reference_frame[:3])
        motion_vector.frames = [skeleton.reference_frame]
        motion_vector.n_frames = 1
        o = self.scene.object_builder.create_object("animation_controller", "skeleton", skeleton, motion_vector, skeleton.frame_time)
        self.controller = o._components["animation_controller"]
        self.skeleton = self.controller.get_skeleton()
        self.init_joints(self.controller)
        self.fill_joint_map()

        self.selectButton.clicked.connect(self.slot_accept)
        self.cancelButton.clicked.connect(self.slot_reject)
        self.applyTwistRotationButton.clicked.connect(self.slot_set_twist)
        self.applySwingRotationButton.clicked.connect(self.slot_set_swing)

        self.setOrthogonalTwistButton.clicked.connect(self.slot_set_orthogonal_twist)
        self.setOrthogonalSwingButton.clicked.connect(self.slot_set_orthogonal_swing)
        self.rotateTwistButton.clicked.connect(self.slot_rotate_twist)
        self.rotateSwingButton.clicked.connect(self.slot_rotate_swing)

        self.flipTwistButton.clicked.connect(self.slot_flip_twist)
        self.flipSwingButton.clicked.connect(self.slot_flip_swing)

        self.flipZAxisButton.clicked.connect(self.slot_flip_z_axis)
        self.alignToUpAxisButton.clicked.connect(self.slot_align_to_up_axis)
        self.alignToForwardAxisButton.clicked.connect(self.slot_align_to_forward_axis)

        self.guessSelectedButton.clicked.connect(self.slot_guess_selected_cos_map)
        self.resetSelectedCosButton.clicked.connect(self.slot_reset_selected_cos_map)
        self.guessAllCosButton.clicked.connect(self.slot_guess_cos_map)
        self.resetAllCosButton.clicked.connect(self.slot_reset_cos_map)
        self.loadDefaultPoseButton.clicked.connect(self.slot_load_default_pose)
        self.applyScaleButton.clicked.connect(self.slot_apply_scale)
        self.jointMapComboBox.currentIndexChanged.connect(self.slot_update_joint_map)
        self.aligningRootComboBox.currentIndexChanged.connect(self.slot_update_aligning_root_joint)

        self.mirrorLeftButton.clicked.connect(self.slot_mirror_left_to_right)
        self.mirrorRightButton.clicked.connect(self.slot_mirror_right_to_left)

        self.is_updating_joint_info = False
        self.success = False
        self.initialized = False
        self.skeleton_data = None
        self.precision = 3
        self.aligning_root_node = self.skeleton.aligning_root_node
        self.fill_root_combobox()
        self.init_aligning_root_node()

    def init_aligning_root_node(self):
        print("init",self.skeleton.root, self.skeleton.aligning_root_node)
        if self.aligning_root_node is None:
            self.aligning_root_node = self.skeleton.root
        if self.aligning_root_node is not None:
            index = self.aligningRootComboBox.findText(self.aligning_root_node, Qt.MatchFixedString)
            print("found index", index, self.aligning_root_node)
            if index >= 0:
                self.aligningRootComboBox.setCurrentIndex(index)

    def closeEvent(self, e):
         self.timer.stop()
         self.view.makeCurrent()
         del self.view

    def on_mouse_click(self, event, ray_start, ray_dir, pos, node_id):
        if event.button() == Qt.LeftButton:
            self.scene.select_object(node_id)
            joint_knob = self.get_selected_joint()
            self.update_joint_info(joint_knob)

    def update_joint_info(self, joint_knob):
        self.is_updating_joint_info = True
        self.scene.scene_edit_widget.reset_rotation()
        self.jointMapComboBox.setCurrentIndex(0)
        label = "Selected Joint: "
        if joint_knob is None:
            label += "None"
            self.jointLabel.setText(label)
            self.is_updating_joint_info = False
            return

        label += joint_knob.joint_name
        joint_name = joint_knob.joint_name
        if "joints" in self.skeleton_model:
            key = find_key(self.skeleton_model["joints"], joint_name)
            print("key", joint_name, key)
            if key is not None:
                index = self.jointMapComboBox.findText(key, Qt.MatchFixedString)
                if index >= 0:
                    self.jointMapComboBox.setCurrentIndex(index)

        if "cos_map" in self.skeleton_model and joint_name in self.skeleton_model["cos_map"]:
            x_vector = self.skeleton_model["cos_map"][joint_name]["x"]
            if x_vector is None:
                x_vector = [1,0,0]
                self.skeleton_model["cos_map"][joint_name]["x"] = x_vector
            y_vector = self.skeleton_model["cos_map"][joint_name]["y"]
            if y_vector is None:
                y_vector = [0,1,0]
                self.skeleton_model["cos_map"][joint_name]["y"] = y_vector
            swing = np.round(x_vector, self.precision)
            twist = np.round(y_vector, self.precision)
            self.set_swing_text(swing)
            self.set_twist_text(twist)
            m = self.skeleton.nodes[joint_name].get_global_matrix(self.reference_frame)[:3,:3]
            g_swing = np.dot(m, swing)
            g_swing = normalize(g_swing)
            g_twist = np.dot(m, twist)
            g_twist = normalize(g_twist)
            q = axes_to_q(g_twist, g_swing)
            m = quaternion_matrix(q)
            #print("g_twist", g_twist, twist)
            #print("g_swing", g_swing, swing)
            self.scene.scene_edit_widget.rotation = m[:3,:3].T
        else:
            print("no cos map", self.skeleton_model.keys())

        self.jointLabel.setText(label)
        self.is_updating_joint_info = False

    def set_swing_text(self, swing):
        self.swingXLineEdit.setText(str(swing[0]))
        self.swingYLineEdit.setText(str(swing[1]))
        self.swingZLineEdit.setText(str(swing[2]))

    def set_twist_text(self, twist):
        self.twistXLineEdit.setText(str(twist[0]))
        self.twistYLineEdit.setText(str(twist[1]))
        self.twistZLineEdit.setText(str(twist[2]))

    def fill_joint_map(self):
        self.jointMapComboBox.clear()
        for idx, joint in enumerate(STANDARD_JOINTS):
            print("add", joint)
            self.jointMapComboBox.addItem(joint, idx)

    def fill_root_combobox(self):
        self.aligningRootComboBox.clear()
        for idx, joint in enumerate(self.controller.get_animated_joints()):
            self.aligningRootComboBox.addItem(joint, idx)

    def init_joints(self, controller):
        for joint_name in controller.get_animated_joints():
            if len(self.skeleton.nodes[joint_name].children) > 0: # filter out end site joints
                child_node = self.skeleton.nodes[joint_name]#.children[0]
                if np.linalg.norm(child_node.offset)> 0:
                    self.scene.object_builder.create_object("joint_control_knob", controller, joint_name, self.radius)

    def draw(self):
        """ draw current scene on the given view
        (note before calling this function the context of the view has to be set as current using makeCurrent() and afterwards the doubble buffer has to swapped to display the current frame swapBuffers())
        """
        if not self.initialized:
            if self.view.graphics_context is not None:
                self.view.resize(400,400)
                self.initialized = True
        self.scene.update(self.dt)
        self.view.makeCurrent()
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        self.view.graphics_context.render(self.scene)
        self.view.swapBuffers()

    def left_display_changed(self, frame_idx):
        if self.controller is not None:
            self.controller.setCurrentFrameNumber(frame_idx)
            self.leftDisplayFrameSpinBox.setValue(frame_idx)

    def left_spinbox_frame_changed(self, frame):
        self.leftStartFrameSlider.setValue(self.leftStartFrameSpinBox.value())
        self.leftEndFrameSlider.setValue(self.leftEndFrameSpinBox.value())

    def left_slider_frame_changed(self, frame):
        self.leftStartFrameSpinBox.setValue(self.leftStartFrameSlider.value())
        self.leftEndFrameSpinBox.setValue(self.leftEndFrameSlider.value())

    def slot_accept(self):
        self.name = str(self.nameLineEdit.text())
        if self.name != "":
            print("accept")
            self.success = True
            self.skeleton = self.controller.get_skeleton()
            self.skeleton.set_reference_frame(self.reference_frame)
            if self.aligning_root_node is not None:
                self.skeleton.aligning_root_node = self.aligning_root_node
            self.skeleton_data = self.skeleton.to_unity_format()
            if "cos_map" in self.skeleton_model:
                for k in self.skeleton_model["cos_map"]:
                    for l in self.skeleton_model["cos_map"][k]:
                        if type(self.skeleton_model["cos_map"][k][l]) == np.ndarray:
                            self.skeleton_model["cos_map"][k][l] = self.skeleton_model["cos_map"][k][l].tolist()
                        else:
                            self.skeleton_model["cos_map"][k][l] = self.skeleton_model["cos_map"][k][l]
            self.close()
        else:
            print("Please provide a name")

    def slot_reject(self):
        self.close()

    def get_selected_joint(self):
        joint_knob = None
        o = self.scene.selected_scene_object
        if o is not None and "joint_control_knob" in o._components:
            joint_knob = o._components["joint_control_knob"]
        return joint_knob

    def set_twist(self, joint_name):
        x = round(float(self.twistXLineEdit.text()),self.precision)
        y = round(float(self.twistYLineEdit.text()),self.precision)
        z = round(float(self.twistZLineEdit.text()),self.precision)
        #set twist axis
        twist = np.array([x,y,z])
        magnitude = np.linalg.norm(twist)
        if magnitude > 0:
            twist /= magnitude
        self.skeleton_model["cos_map"][joint_name]["y"] = twist

    def set_swing(self, joint_name):
        x = round(float(self.swingXLineEdit.text()),self.precision)
        y = round(float(self.swingYLineEdit.text()),self.precision)
        z = round(float(self.swingZLineEdit.text()),self.precision)
        #set swing axis
        swing = np.array([x,y,z])
        magnitude = np.linalg.norm(swing)
        if magnitude > 0:
            swing /= magnitude
        self.skeleton_model["cos_map"][joint_name]["x"] = swing.tolist()

    def slot_set_twist(self):
        plot = False
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        self.set_swing(joint_knob.joint_name)
        self.set_twist(joint_knob.joint_name)
        self.update_joint_info(joint_knob)

    def slot_set_swing(self):
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        self.set_swing(joint_knob.joint_name)
        self.set_twist(joint_knob.joint_name)
        self.update_joint_info(joint_knob)

    def slot_set_orthogonal_twist(self):
        """ https://stackoverflow.com/questions/33658620/generating-two-orthogonal-vectors-that-are-orthogonal-to-a-particular-direction """
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        joint_name = joint_knob.joint_name
        #get swing axis
        swing = np.array(self.skeleton_model["cos_map"][joint_name]["x"])
        # find orthogonal vector
        y = np.array(self.skeleton_model["cos_map"][joint_name]["y"])

        #y = np.random.randn(3)  # take a random vector
        y -= y.dot(swing) * swing      # make it orthogonal to twist
        y /= np.linalg.norm(y)  # normalize it

        #replace twist axis
        self.set_twist_text(y)
        self.skeleton_model["cos_map"][joint_name]["y"] = y
        self.update_joint_info(joint_knob)

    def slot_set_orthogonal_swing(self):
        """ https://stackoverflow.com/questions/33658620/generating-two-orthogonal-vectors-that-are-orthogonal-to-a-particular-direction """
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        joint_name = joint_knob.joint_name
        #get twist axis
        twist = np.array(self.skeleton_model["cos_map"][joint_name]["y"] )
        x = np.array(self.skeleton_model["cos_map"][joint_name]["x"] )
        x -= x.dot(twist) * twist      # make it orthogonal to twist
        x /= np.linalg.norm(x)  # normalize it
        #replace twist axis
        self.set_swing_text(x)
        self.skeleton_model["cos_map"][joint_name]["x"] = x
        self.update_joint_info(joint_knob)

    def slot_flip_twist(self):
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        joint_name = joint_knob.joint_name
        #get twist axis
        twist = np.array(self.skeleton_model["cos_map"][joint_name]["y"])
        twist *= -1
        self.skeleton_model["cos_map"][joint_name]["y"] = twist
        self.set_twist_text(twist)
        self.update_joint_info(joint_knob)

    def slot_flip_swing(self):
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        joint_name = joint_knob.joint_name
        swing = np.array(self.skeleton_model["cos_map"][joint_name]["x"])
        swing *= -1
        self.skeleton_model["cos_map"][joint_name]["x"] = swing
        self.set_swing_text(swing)
        self.update_joint_info(joint_knob)

    def slot_rotate_twist(self):
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        joint_name = joint_knob.joint_name
        #get twist axis
        angle = round(float(self.twistRotationLineEdit.text()),self.precision)
        rotation_axis = np.array(self.skeleton_model["cos_map"][joint_name]["x"])
        q = quaternion_about_axis(np.deg2rad(angle), rotation_axis)
        q = normalize(q)
        twist = np.array(self.skeleton_model["cos_map"][joint_name]["y"])
        twist = rotate_vector(q, twist)
        twist = normalize(twist)
        self.skeleton_model["cos_map"][joint_name]["y"] = twist
        self.set_twist_text(twist)
        self.update_joint_info(joint_knob)

    def slot_rotate_swing(self):
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        joint_name = joint_knob.joint_name

        angle = round(float(self.swingRotationLineEdit.text()),self.precision)
        rotation_axis = np.array(self.skeleton_model["cos_map"][joint_name]["y"])
        q = quaternion_about_axis(np.deg2rad(angle), rotation_axis)
        q = normalize(q)
        swing = np.array(self.skeleton_model["cos_map"][joint_name]["x"])
        swing = rotate_vector(q, swing)
        swing = normalize(swing)

        self.skeleton_model["cos_map"][joint_name]["x"] = swing
        self.set_swing_text(swing)
        self.update_joint_info(joint_knob)

    def slot_flip_z_axis(self):
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        joint_name = joint_knob.joint_name
        twist = np.array(self.skeleton_model["cos_map"][joint_name]["y"])
        swing = np.array(self.skeleton_model["cos_map"][joint_name]["x"])
        new_swing = twist
        new_twist = swing
        #print("new swing", new_swing, swing)
        self.skeleton_model["cos_map"][joint_name]["y"] = new_twist
        self.skeleton_model["cos_map"][joint_name]["x"] = new_swing
        self.set_swing_text(new_swing)
        self.set_twist_text(new_twist)
        self.update_joint_info(joint_knob)

    def slot_guess_cos_map(self):
        """ creates a guess for the coordinate system for all joints"""
        temp_skeleton = copy(self.skeleton)
        temp_skeleton.skeleton_model = self.skeleton_model
        cos_map = create_local_cos_map_from_skeleton_axes_with_map(temp_skeleton)
        self.skeleton_model["cos_map"] = cos_map
        joint_knob = self.get_selected_joint()
        if joint_knob is not None:
            self.update_joint_info(joint_knob)

    def slot_reset_cos_map(self):
        """ resets the coordinate systems for all joints"""
        for joint_name in self.skeleton_model["cos_map"]:
            up_vector = self.skeleton_model["cos_map"][joint_name]["y"]
            x_vector = self.skeleton_model["cos_map"][joint_name]["x"]
            if up_vector is not None and x_vector is not None:
                new_up_vector, new_x_vector = self.reset_joint_cos(joint_name, up_vector, x_vector)
                self.skeleton_model["cos_map"][joint_name]["y"] = new_up_vector
                self.skeleton_model["cos_map"][joint_name]["x"] = new_x_vector
        joint_knob = self.get_selected_joint()
        if joint_knob is not None:
            self.update_joint_info(joint_knob)

    def slot_guess_selected_cos_map(self):
        """ creates a guess for the for the selected joint"""
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        joint_name = joint_knob.joint_name
        temp_skeleton = copy(self.skeleton)
        temp_skeleton.skeleton_model = self.skeleton_model
        cos_map = create_local_cos_map_from_skeleton_axes_with_map(temp_skeleton)
        self.skeleton_model["cos_map"][joint_name] = cos_map[joint_name]
        self.update_joint_info(joint_knob)

    def slot_reset_selected_cos_map(self):
        """ creates resetrs the coordinate system for the selected joint"""
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        joint_name = joint_knob.joint_name
        if joint_name in self.skeleton_model["cos_map"]:
            up_vector = self.skeleton_model["cos_map"][joint_name]["y"]
            x_vector = self.skeleton_model["cos_map"][joint_name]["x"]
            new_up_vector, new_x_vector = self.reset_joint_cos(joint_name, up_vector, x_vector)
            self.skeleton_model["cos_map"][joint_name]["y"] = new_up_vector
            self.skeleton_model["cos_map"][joint_name]["x"] = new_x_vector
            self.update_joint_info(joint_knob)

    def slot_update_joint_map(self):
        if not self.is_updating_joint_info and "joints" in self.skeleton_model:
            joint_knob = self.get_selected_joint()
            if joint_knob is not None:
                new_joint_key = str(self.jointMapComboBox.currentText())
                old_joint_key = find_key(self.skeleton_model["joints"], joint_knob.joint_name)
                if old_joint_key in self.skeleton_model["joints"]:
                    self.skeleton_model["joints"][old_joint_key] = None
                self.skeleton_model["joints"][new_joint_key] = joint_knob.joint_name
                print("update joint mapping", joint_knob.joint_name, new_joint_key)
            else:
                print("is updating joint info")

    def reset_joint_cos(self, joint_name, up_vector, x_vector, target_up_vector=DEFAULT_TARGET_CS_UP):
        """ rotates the up_vector to look towards target_up_vector and rotates the x_vector with the same rotation """
        m = self.skeleton.nodes[joint_name].get_global_matrix(self.skeleton.reference_frame)[:3,:3]
        m_inv = np.linalg.inv(m)
        target_up_vector = normalize(target_up_vector)
        local_target = np.dot(m_inv, target_up_vector)
        local_target = normalize(local_target)
        q = quaternion_from_vector_to_vector(up_vector, local_target)
        x_vector = rotate_vector(q, x_vector)

        x_vector -= x_vector.dot(local_target) * local_target      # make it orthogonal to twist
        x_vector /= np.linalg.norm(x_vector)  # normalize it
        x_vector = normalize(x_vector)
        return local_target, x_vector

    def slot_update_aligning_root_joint(self):
        if not self.is_updating_joint_info and "joints" in self.skeleton_model:
            self.aligning_root_node = str(self.aligningRootComboBox.currentText())

    def slot_load_default_pose(self):
        filename = QFileDialog.getOpenFileName(self, 'Load From File', '.')[0]
        filename = str(filename)
        if os.path.isfile(filename):
            motion = load_motion_from_bvh(filename)
            if len(motion.frames):
                self.reference_frame = motion.frames[0]
                frames = [self.reference_frame]
                self.controller.replace_frames(frames)
                self.controller.set_reference_frame(0)
                self.controller.updateTransformation()
                print("replaced frames")

    def slot_apply_scale(self):
        scale = float(self.scaleLineEdit.text())
        if scale > 0:
            self.controller.set_scale(scale)
            frames = [self.reference_frame]
            self.controller.replace_frames(frames)
            self.controller.currentFrameNumber = 0
            self.controller.updateTransformation()

    def slot_align_to_up_axis(self):
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        joint_name = joint_knob.joint_name
        if joint_name in self.skeleton_model["cos_map"]:
            up_vector = self.skeleton_model["cos_map"][joint_name]["y"]
            x_vector = self.skeleton_model["cos_map"][joint_name]["x"]
            q_offset = get_axis_correction(self.skeleton, joint_name, up_vector, OPENGL_UP_AXIS)
            up_vector = rotate_vector(q_offset, up_vector)
            x_vector = rotate_vector(q_offset, x_vector)
            self.skeleton_model["cos_map"][joint_name]["x"] = normalize(x_vector)
            self.skeleton_model["cos_map"][joint_name]["y"] = normalize(up_vector)
            self.update_joint_info(joint_knob)

    def slot_align_to_forward_axis(self):
        joint_knob = self.get_selected_joint()
        if joint_knob is None:
            return
        joint_name = joint_knob.joint_name
        if joint_name in self.skeleton_model["cos_map"]:
            up_vector = self.skeleton_model["cos_map"][joint_name]["y"]

            m = self.skeleton.nodes[joint_name].get_global_matrix(self.skeleton.reference_frame)[:3,:3]
            m_inv = np.linalg.inv(m)
            target_vector = np.dot(m, up_vector)
            target_vector[1] = 0
            target_vector = normalize(target_vector)
            local_up = np.dot(m_inv, target_vector)
            local_up = normalize(local_up)
            self.skeleton_model["cos_map"][joint_name]["y"] = local_up

            x_vector = self.skeleton_model["cos_map"][joint_name]["x"]
            q = quaternion_from_vector_to_vector(up_vector, local_up)
            x_vector = rotate_vector(q, x_vector)

            x_vector -= x_vector.dot(local_up) * local_up      # make it orthogonal to twist
            x_vector /= np.linalg.norm(x_vector)  # normalize it
            self.skeleton_model["cos_map"][joint_name]["x"] = normalize(x_vector)
            self.update_joint_info(joint_knob)

    def slot_mirror_left_to_right(self):
        self.skeleton_model = mirror_join_map(self.skeleton, self.skeleton_model, STANDARD_MIRROR_MAP_LEFT)
        print("mirrored left to right")
        print(self.skeleton_model["joints"])

    def slot_mirror_right_to_left(self):
        self.skeleton_model = mirror_join_map(self.skeleton, self.skeleton_model, STANDARD_MIRROR_MAP_RIGHT)
        print("mirrored right to left")


text:
import numpy as np
import itertools

def build_full(action_space):
    actions = []
    for i in range(action_space.n):
        action_array = np.zeros(action_space.n)
        action_array[i] = 1
        actions.append(action_array)

    return np.asarray(actions)

def build_simplified(action_space):
    """
    Return all combinations of actions 0, 6, 7, 8 (no-np, left, right, jump)
    """
    actions = [4, 5, 6, 7, 8]

    multi_binary_actions = []
    for combination in itertools.chain.from_iterable(itertools.combinations(actions, i) for i in range(len(actions) + 1)):
        multi_binary_action = np.zeros(action_space.n)
        multi_binary_action[list(combination)] = 1
        multi_binary_actions.append(multi_binary_action)

    multi_binary_actions.append(np.zeros(action_space.n))

    return np.array(multi_binary_actions)

def build(action_space):
    return build_simplified(action_space)

ext: py

max_stars_repo_name: hermesdt/retro-contest

max_issues_repo_name: hermesdt/retro-contest

max_forks_repo_name: hermesdt/retro-contest

reason: import numpy


text:
# vim: set et sw=4 sts=4 fileencoding=utf-8:
#
# Python Motion Tracker module of the pyCAMTracker package
# Copyright (c) 2017-2018 Axel Barnitzke <barney@xkontor.org>
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the copyright holder nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

from __future__ import (
    unicode_literals,
    print_function,
    division,
    absolute_import,
    )

# Make Py2's str equivalent to Py3's
str = type('')
try:
    range = xrange
except NameError:
    pass

import sys
import threading
from collections import deque
import numpy as np
import cv2
from math import atan2,hypot,degrees,acos,pi,sqrt
from time import sleep,time
import prctl

#- globals
MAX_TRACKS     = 16

#- sorting functions
def distance(t,p):
    if t.updates == 0:
        return 99999
    # BUG: upper corner as reference point is more stable than center
    px = p[0] + p[2] / 2.0
    py = p[1] + p[3] / 2.0
    #px = p[0]
    #py = p[1]
    return abs(t.cx - px) + abs(t.cy - py)
    #return np.hypot(t.cx - px, t.cy - py)

# for same distance prefer the more mature track
# remember: we sort from smallest to biggest
def weighted_distance(t,p):
    dist = float(distance(t,p))
    val = 0.0
    divisor = 10.0
    if t.updates > 50:
        divisor = 100.0
    if dist > 0.0:
        val = 100.0 * (dist - t.updates / divisor)
    #print("   >%s(%d) %4.1f %4.1f"%(t.name,t.updates,dist,val))
    return(val)

#- sort tracks by number of updates
def by_updates(t):
    return t.updates

#- move angle into range [-pi, pi]
def normalize_angle(x):
    x %= 2 * np.pi
    if x > np.pi:
        x -= 2 * np.pi
    return(x)

class Tracker(threading.Thread):
    """
    Track manager: assigns macro blocks to tracks
    """
    #--------------------------------------------------------------------
    #-- constructor
    #--------------------------------------------------------------------
    def __init__(self, camera, greenLed=None, redLed=None, config=None):
        super(Tracker,self).__init__()
        self.lock = threading.Lock()
        self.config = config
        self.camera = camera
        self.resx = camera.resolution[0]
        self.resy = camera.resolution[1]
        self.greenLEDThread = greenLed
        self.redLEDThread   = redLed
        self.updates  = 0
        self.frame  = 0
        self.noise  = 0.0
        self.active_tracks = 0
        self.motion = None
        self.locked = False
        self.maxDist = 5
        self.trackLifeTime = 17
        self.trackMaturity = 7
        self.debug = False
        self.fobj = None
        self.cols = 0
        self.rows = 0
        self.direction = 0
        self.detectionDelay = -0.999
        prctl.set_name('ptrk.Tracker')

        #- initialize a fixed number of threads (less garbarge collection)
        self.track_pool = []

        #- do things according configuration
        if config is not None:
            Track.xCross = config.conf['xCross']
            Track.yCross = config.conf['yCross']
            Track.maxDist = self.maxDist = config.conf['maxDist']
            Track.minCosDelta = config.conf['minCosDelta']
            self.trackLifeTime = config.conf['trackLifeTime']
            self.trackMaturity = config.conf['trackMaturity']
            self.debug = config.conf['debug']

        #- create track instances
        for i in range(0,config.conf['maxTracks']):
            self.track_pool.append(Track(self))

        #- thread initialisation stuff
        self.event = threading.Event()
        self.event.clear()
        self.q = deque()
        self.terminated = False
        self.daemon = True
        self.start()

    #--------------------------------------------------------------------
    #-- callback for 'trackMaturity' command
    #--------------------------------------------------------------------
    def set_trackMaturity(self, value):
        if value > 3 and value <= 10:
            self.trackMaturity = value
            print("MotionTracker::trackMaturity: %d" % value)
            if self.config:
                self.config.conf['trackMaturity'] = value

    #--------------------------------------------------------------------
    #-- callback for 'maxDist' command
    #--------------------------------------------------------------------
    def set_maxDist(self, value):
        if value > 0 and value < 25:
            Track.maxDist = value
            print("MotionTracker::maxDist: %d" % value)
            if self.config:
                self.config.conf['maxDist'] = value

    #--------------------------------------------------------------------
    #-- called by picamera after sizes are known
    #--------------------------------------------------------------------
    def setup_sizes(self, rows, cols):
        Track.maxX = self.cols = cols
        Track.maxY = self.rows = rows
        Track.maxArea = cols * rows

    def testCrossing(self, value):
        if value > 0:
            xm = int(self.cols / 2)
            ym = int(self.rows / 2)
            #         [target_rect    ]         [bounding_rect          ]
            motion = [[xm-5,ym-5,10,10],[10,10],[xm-10,ym-10,xm+10,ym+10]]

            frame = self.camera.frame.index
            self.crossed(99,time(),frame,motion)

    #--------------------------------------------------------------------
    #-- release lock
    #--------------------------------------------------------------------
    def releaseLock(self):
        with self.lock:
            self.locked = False
            frame  = self.frame
            motion = self.motion
        return (frame, motion)

    #--------------------------------------------------------------------
    #-- get crossing status
    #--------------------------------------------------------------------
    def getStatus(self):
        with self.lock:
            frame  = self.frame
            motion = self.motion
            delay = self.detectionDelay
            self.frame = 0
            self.direction = 0
            self.detectionDelay = -0.999
        return (delay, frame, motion)

    #--------------------------------------------------------------------
    #-- callback for crossing event
    #--------------------------------------------------------------------
    def crossed(self, updates, timestamp, frame, motion, positive_direction=False):
        if self.locked:
            print("blocked")
            return False

        with self.lock:
            self.locked = True
            self.camera.request_key_frame()
            if self.greenLEDThread:
                self.greenLEDThread.event.set()
            self.updates  = updates
            self.frame  = frame
            self.detectionDelay = time() - timestamp
            self.motion = motion
            self.direction = 1 if positive_direction else -1

        return True

    #--------------------------------------------------------------------
    #-- callback for cut event
    #--------------------------------------------------------------------
    def turned(self, updates, timestamp, frame, motion, positive_direction=False):
        if self.locked:
            print("blocked")
            return False

        with self.lock:
            self.locked = True
            self.camera.request_key_frame()
            if self.redLEDThread:
                self.redLEDThread.event.set()
            self.updates  = updates
            self.frame  = -frame
            self.detectionDelay = time() - timestamp
            self.motion = motion
            self.direction = 0

        return True

    #--------------------------------------------------------------------
    #-- without thread (don't start the thread otherwise index error raised!)
    #--------------------------------------------------------------------
    def deb_update_tracks(self, frame, motion):
        self.update_track_pool(frame, motion)

    #--------------------------------------------------------------------
    #- create debug output
    #--------------------------------------------------------------------
    def debug_out(self, frame, motion):
        if self.fobj is None:
            try:
                self.fobj = open("debug_tracker.csv", "w")
            except:
                raise
        for rr,vv in motion:
            self.fobj.write("%4d,%0d,%0d,%0d,%0d,%4.2f,%4.2f\n" % (frame,rr[0],rr[1],rr[2],rr[3],vv[0],vv[1]))

    #--------------------------------------------------------------------
    #-- queue new points and feed worker
    #--------------------------------------------------------------------
    def update_tracks(self, timestamp, frame, motion):
        self.q.append([timestamp,frame,motion])
        self.event.set()

    #--------------------------------------------------------------------
    #-- stop all threading stuff
    #--------------------------------------------------------------------
    def stop(self):
        if self.fobj:
            self.fobj.close()
        self.terminated = True
        self.q.append([0, [[[0,0,0,0],[0,0]]]])

    #--------------------------------------------------------------------
    #-- Thread run function (calls update_track_pool)
    #-- (needs dqueue to be fats enough)
    #--------------------------------------------------------------------
    def run(self):
        while not self.terminated:
            try:
                timestamp,frame,motion = self.q.popleft()
                self.update_track_pool(timestamp,frame,motion)
            except IndexError:
                self.event.clear()
                self.event.wait(1)

    #--------------------------------------------------------------------
    #-- offer new points to existing tracks
    #--------------------------------------------------------------------
    def update_track_pool(self, timestamp, frame, motion):
        # walk through all changes
        self.updated = False
        has_been_tracked = 0x00000000
        min_dist = 0

        for rn,vn in motion:
            #-- search a track for this coordinate
            tracked = 0x00000000
            # >>> debug
            #cx = rn[0] + rn[2] / 2
            #cy = rn[1] + rn[3] / 2
            #print( "frame:%d %d,%d ===============================" % ( frame,rn[0],rn[1] ))
            # <<< debug
            #-- sorting by distance really makes sence here
            for track in sorted(self.track_pool, key=lambda t: weighted_distance(t,rn)):
                #-- skip tracks already updated
                if has_been_tracked & track.id:
                    continue

                #-- the rest of the coordinates can be ignored
                dist = distance(track,rn)
                maxDist = max(self.maxDist,2*(rn[2]+rn[3]))
                if dist > maxDist:
                    #print("   [%s] precheck: dist too big! %d > %d" % (track.name, dist, maxDist))
                    break

                # >>> debug
                #print("   check: [%s](%d): @%d,%d dist:%d" % (track.name,track.updates,track.re[0],track.re[1],dist))
                # <<< debug

                #-- check if track takes coordinates
                tracked = track.update(timestamp,frame,rn,vn)
                if tracked:
                    has_been_tracked |= tracked
                    #print( "   [%s](%d) updated with: %d,%d" % (track.name, track.updates, rn[0],rn[1]))
                    self.updated = True
                    break

            #-- not yet tracked -> find a free slot
            if not tracked:
                for track in self.track_pool:
                    if track.updates == 0:
                        #print("   [%s] new %d/%d" % (track.name, rn[0],rn[1]))
                        track.new_track(timestamp,frame,rn,vn)
                        self.updated = True
                        break

        #-- remove aged tracks
        noise = 0
        active = 0
        for track in self.track_pool:
            updates = track.updates
            if updates > 0:
                active += 1
                if updates < 3:
                    noise += 1
            if frame - track.lastFrame > 2:
                track.reset()
                continue
            if updates and frame - track.lastFrame > self.trackLifeTime:
                track.reset()

        #self.noise = float(noise / len(self.track_pool))
        self.noise = float(active / len(self.track_pool))

        self.active_tracks = active

    def showTracks(self, frame, vis):
        """
        show graphical interpretation of all tracks
        """
        for track in self.track_pool:
            track.showTrack(vis,frame)

class Track:
    """
    data container for tracks
    """
    track_names = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
    numtracks   = 0
    #minCosDelta = 0.707 #cos(2*22.5)
    #minCosDelta = 0.9
    minCosDelta = 0.5 #cos(2*30.0)
    maxDist     = 20.0
    maxLifeTime = 17
    maxX        = 99999
    maxY        = 99999
    maxArea     = 1
    xCross      = -1
    yCross      = -1

    #--------------------------------------------------------------------
    #--
    #--------------------------------------------------------------------
    def __init__(self, parent=None):
        # track identification
        index = Track.numtracks % 32
        Track.numtracks += 1
        self.id = 1 << index
        self.name = Track.track_names[index]
        self.parent = parent
        # reset track data
        self.updates = 9999
        self.reset()

    def reset(self):
        """
        reset/reuse this object to to avoid too much garbage collection
        """
        if self.updates < 1:
            return

        #print("[%s](%d) reset" % (self.name,self.updates))

        self.updates = 0
        self.tr   = []
        self.re   = [0,0,0,0]
        self.vv   = np.array([0.0,0.0])
        self.cx   = 0
        self.cy   = 0
        self.distance = [0,0]
        self.dirxOK = False
        self.diryOK = False
        self.distxOK = False
        self.distyOK = False
        self.old_dir  = None
        self.old_dist = None
        self.old_area = None
        self.maxx = 0
        self.maxy = 0
        self.minx = 99999
        self.miny = 99999
        self.maturity = self.parent.trackMaturity
        self.progressx = False
        self.progressy = False
        self.deltaX = 0
        self.deltaY = 0
        self.noprogressx  = 0
        self.noprogressy  = 0
        self.lastFrame = 0
        self.timestamp = 0
        self.isGrowing = True
        self.cleanCrossings()

    #--------------------------------------------------------------------
    #--
    #--------------------------------------------------------------------
    def cleanCrossings(self):
        self.turnedX   = False
        self.turnedY   = False
        self.crossedX  = False
        self.crossedY  = False

    def clean(self,frame):
        """
        check if track can be resetted
        """
        if self.updates > 0:
            # TODO: keep this status alive for a couple of frames
            #if self.updates > 20 and self.progressx == 0 and self.progressy == 0:
            #    print "[%s](%d) no motion!" % (self.name, self.updates)
            #    self.reset()
            #    return

            if frame - self.lastFrame > Track.maxLifeTime:
                #print "[%s] (%d) clean" % (self.name, self.updates)
                self.reset()

    def new_track(self,timestamp,frame,rn,vn):
        """
        start a new track
        """

        #if rn[0] + vn[0] < 0 or rn[1] + vn[1] < 0:
        #    #print("[%s]: %d/%d %3.1f/%3.1f new track rejected (too low)" % (self.name,rn[0],rn[1],vn[0],vn[1]))
        #    return 0
        #if rn[0] + rn[2] + vn[0] > Track.maxX or rn[1] + rn[3] + vn[1] > Track.maxY:
        #    #print("[%s]: %d/%d %3.1f/%3.1f new track rejected (too high)" % (self.name, rn[0]+rn[2],rn[1]+rn[3],vn[0],vn[1]))
        #    return 0

        cxn  = rn[0]+rn[2]/2 #xn+wn/2
        cyn  = rn[1]+rn[3]/2 #yn+hn/2

        #cxn  = rn[0]
        #cyn  = rn[1]

        self.re  = rn
        self.vv  = np.array(vn)
        self.cx  = cxn
        self.cy  = cyn
        self.tr.append([cxn,cyn])
        self.updates = 1
        self.lastFrame = frame
        self.timestamp = timestamp

        return self.id

    def checkStartingConditions(self):
        """
        check conditions at point #3
        """

        if self.updates >= 3:

            if Track.xCross > 0:
                x0 = int(self.tr[0][0])
                x1 = int(self.tr[2][0])
                xc = int(Track.maxX / 2)
                #- if the first occurence is right from center movement must be negative
                if x0 >= xc and (x1-x0) <= 0:
                    return True
                #- if the first occurence is left from center movement must be positive
                if x0 < xc and (x1-x0) >= 0:
                    return True

            if Track.yCross > 0:
                y0 = int(self.tr[0][1])
                y1 = int(self.tr[2][1])
                yc = int(Track.maxY / 2)
                #- if the first occurence is above center movement must be negative
                if y0 >= yc and (y1-y0) <= 0:
                    return True
                #- if the first occurence is below center movement must be positive
                if y0 < yc and (y1-y0) >= 0:
                    return True

                #print("[%s](%d) start failed: yc: %d dy: %d" % (self.name, self.updates, y0, (y1-y0)))

        return False

    #--------------------------------------------------------------------
    #-- update growing status
    #--------------------------------------------------------------------
    def updateGrowingStatus(self, rn):
        neg_count = 6
        # does the track expand into any direction?
        maxx = max(self.maxx, rn[0]+rn[2])
        maxy = max(self.maxy, rn[1]+rn[3])
        minx = min(self.minx, rn[0])
        miny = min(self.miny, rn[1])

        # update progress indicators in y direction
        if maxy > self.maxy or miny < self.miny:
            self.progressy = True
            self.deltaY = maxy - miny
            self.noprogressy = 0
        else:
            self.progressy = False

        # update progress indicators in x direction
        if maxx > self.maxx or minx < self.minx:
            self.progressx = True
            self.deltaX = maxx - minx
            self.noprogressx = 0
        else:
            self.progressx = False

        # update moving area
        self.maxx = maxx
        self.maxy = maxy
        self.minx = minx
        self.miny = miny

        self.isGrowing =  self.progressx or self.progressy

    #--------------------------------------------------------------------
    #-- check if object is turnung before crossing plane
    #--------------------------------------------------------------------
    def detectTurn(self, dx, dy, rn):
        backward_maturiy = 5  #self.maturity
        dmin = 10
        checkDist = 5

        #- check for Y-Turn
        if Track.yCross > 0 and not (self.progressy or self.deltaY<dmin or self.turnedY or self.crossedY):
            # develope validity near by the turn
            if self.noprogressy == 1:
                # develope some criteria for a valid detection
                dy0 = self.distance[1]
                rel = -0.67 * (rn[3] /  self.deltaY) + 0.9
                self.diryOK = (dy0 < 0 and self.miny > Track.yCross and abs(self.miny-Track.yCross) < checkDist) or \
                              (dy0 > 0 and self.maxy < Track.yCross and abs(Track.yCross-self.maxy) < checkDist)
                self.distyOK = abs(dy0) > rel * self.deltaY
                #dist = (Track.yCross-self.maxy) if dy > 0 else (self.miny-Track.yCross)
                #print("[%s](%d) dy: %d, md: %d rel: %4.2f<%4.2f? dist: %d" % (self.name, self.updates, dy0, self.deltaY, rel, abs(dy0)/self.deltaY, dist))

            # track needs some maturity to have a turn detected
            self.noprogressy += 1
            if self.noprogressy > backward_maturiy and self.diryOK and self.distyOK:
                self.turned()
                self.turnedY = True
                print("[%s](%02d) y:%d/%d Y-TURN"
                      % (self.name,self.updates,rn[1],rn[0]))

        #- check for X-Turn
        if Track.xCross > 0 and not (self.progressx or self.deltaX<dmin or self.turnedX or self.crossedX):
            # develope validity near by the turn
            if self.noprogressx == 1:
                # develope some criteria for a valid detection
                dx0 = self.distance[0]
                rel = -0.67 * (rn[2] /  self.deltaX) + 0.9
                self.dirxOK = (dx0 < 0 and self.minx > Track.xCross and abs(self.minx-Track.xCross) < checkDist) or \
                              (dx0 > 0 and self.maxx < Track.xCross and abs(Track.xCross-self.maxx) < checkDist)
                self.distxOK = abs(dx0) > rel * self.deltaX
                #print("[%s](%d) dy: %d, md: %d rel: %4.2f<%4.2f?" % (self.name, self.updates, dx0, self.deltaX, rel, abs(dx0)/self.deltaX))

            # track needs some maturity to have a turn detected
            self.noprogressx += 1
            if self.noprogressx > backward_maturiy and self.dirxOK and self.distxOK:
                self.turned()
                self.turnedX = True
                print("[%s](%02d) x:%d/%d X-TURN"
                      % (self.name,self.updates,rn[0],rn[1]))

    #--------------------------------------------------------------------
    #-- raise crossing handler in parent class
    #--------------------------------------------------------------------
    def crossed(self, positive=False):
        if self.parent:
            self.parent.crossed(self.updates, self.timestamp, self.lastFrame, [self.re, self.vv, [self.minx, self.miny, self.maxx, self.maxy]], positive)

    def turned(self, positive=False):
        if self.parent:
            self.parent.turned(self.updates, self.timestamp, self.lastFrame, [self.re, self.vv, [self.minx, self.miny, self.maxx, self.maxy]], positive)

    #--------------------------------------------------------------------
    #-- main target: is the object crossing the crossing line?
    #-- TODO: make the same for x direction
    #--------------------------------------------------------------------
    def detectCrossing(self, dx, dy, r):
        delta = 3

        if Track.yCross > 0:
            # track is crossing target line in Y direction
            #  x0,y0 +--------------+
            #        |              |
            #        +--------------+ x1,y1
            #
            #  v > 0 |          --->|
            #        |              |
            #  v < 0 |<---          |
            #print("   [%s](%02d) x0/y0:%d/%d x1/y1:%d/%d vy:%3.1f/%3.1f dy:%d/%d maturity: %d"
            #        % (self.name,self.updates,r[0],r[1],r[0]+r[2],r[1]+r[3],-self.vv[0],-self.vv[1],dy,dx, self.maturity))
            #print("   %4.2f > 0.4?" % (self.distance[1] / self.deltaY))
            if self.updates >= self.maturity and self.progressy == True and self.crossedY == False:
                # develope indicators
                vx = -self.vv[0]; vy = -self.vv[1] # remember the velocity has wrong direction!
                x0 = r[0]; y0 = r[1]
                y1 = r[1] + r[3]

                # for low speeds take distance as indicator
                vy_ = abs(vy)
                if vy_ < 0.1:
                    vy = float(dy)
                if vy_ > 5:
                    delta = int(vy_/2) + 1

                # this model uses a simple >= limit to detect a crossing event
                crossedYPositive =  vy >  0.1 and y1 >= Track.yCross and (y1 - delta) < Track.yCross and self.miny < Track.yCross and self.deltaY > self.maturity #and self.distance[1] / self.deltaY > 0.4
                crossedYNegative =  vy < -0.1 and y0 <= Track.yCross and (y0 + delta) > Track.yCross and self.maxy > Track.yCross and self.deltaY > self.maturity #and -self.distance[1] / self.deltaY > 0.4

                if crossedYPositive:
                    delay = (time() - self.timestamp) * 1000.0
                    print("[%s](%02d/%3.0f) y1:%d/%d vy:%3.1f/%3.1f dy:%d/%d delta:%d dist:%d Y-CROSSED++++++++++++++++++++"
                        % (self.name,self.updates,delay,y1,x0,vy,vx,dy,dx,self.deltaY,self.distance[1]))
                    self.crossedY = True
                    self.crossed(positive=True)

                if crossedYNegative:
                    delay = (time() - self.timestamp) * 1000.0
                    print("[%s](%02d/%3.0f) y0:%d/%d vy:%3.1f/%3.1f dy:%d/%d delta:%d dist:%d Y-CROSSED--------------------"
                        % (self.name,self.updates,delay,y0,x0,vy,vx,dy,dx,self.deltaY,self.distance[1]))
                    self.crossedY = True
                    self.crossed(positive=False)

        if Track.xCross > 0:
            # track is crossing target line in X direction
            #  x0,y0 +--------------+
            #        |              |
            #        +--------------+ x1,y1
            #
            #  v > 0 |          --->|
            #        |              |
            #  v < 0 |<---          |

            if self.updates >= self.maturity and self.progressx == True and self.crossedX == False:
                # develope indicators
                vx = -self.vv[0]; vy = -self.vv[1] # remember the velocity has wrong direction!
                x0 = r[0]; y0 = r[1]
                x1 = r[0] + r[2]

                # for low speeds take distance as indicator
                vx_ = abs(vx)
                if vx_ < 0.1:
                    vx = float(dx)
                if vx_ > 5:
                    delta = int(vx_/2) + 1

                # this model uses a simple >= limit to detect a crossing event
                crossedXPositive =  vx >  0.1 and x1 >= Track.xCross and (x1 - delta) < Track.xCross and self.minx < Track.xCross and self.deltaX > self.maturity
                crossedXNegative =  vx < -0.1 and x0 <= Track.xCross and (x0 + delta) > Track.xCross and self.maxx > Track.xCross and self.deltaX > self.maturity

                if crossedXPositive:
                    delay = (time() - self.timestamp) * 1000.0
                    print("[%s](%02d/%3.0f) x1:%d/%d vx:%3.1f/%3.1f dx:%d/%d X-CROSSED++++++++++++++++++++"
                        % (self.name,self.updates,delay,x1,y0,vx,vy,dx,dy))
                    self.crossedX = True
                    self.crossed(positive=True)

                if crossedXNegative:
                    delay = (time() - self.timestamp) * 1000.0
                    print("[%s](%02d/%3.0f) x0:%d/%d vx:%3.1f/%3.1f dx:%d/%d X-CROSSED--------------------"
                        % (self.name,self.updates,delay,x0,y0,vx,vy,dx,dy))
                    self.crossedX = True
                    self.crossed(positive=False)

    #--------------------------------------------------------------------
    #-- does the track leave the tracking area?
    #--------------------------------------------------------------------
    def isLeaving(self, dx, dy):
        leavex = self.maxx+dx > Track.maxX or self.minx+dx < 0
        leavey = self.maxy+dy > Track.maxY or self.miny+dy < 0
        return leavex and leavey

    #--------------------------------------------------------------------
    #-- update track data
    #--------------------------------------------------------------------
    def update(self,timestamp,frame,rn,vn):

        # this is not the place for new tracks
        if self.updates < 1:
            return 0x00000000

        # remove expired tracks
        if frame - self.lastFrame > Track.maxLifeTime:
            self.reset()
            return 0x00000000

        self.lastFrame = frame
        self.timestamp = timestamp

        # PiMotionAnalysis.analyse may be called more than once per frame -> double hit?
        cxn  = rn[0]+rn[2]/2.0
        cyn  = rn[1]+rn[3]/2.0

        # estimating via the object center produces very much noise. (??? I didn't think much about that)
        # estimation is done via the upper left corner (until I have something better)

        #cxn  = rn[0]
        #cyn  = rn[1]

        # accept slow moving objects
        if self.cx == cxn and self.cy == cyn and self.re[2] == rn[2] and self.re[3] == rn[3]:
            self.updates += 1
            self.updateGrowingStatus(rn)
            #print "[%s] double hit" % self.name
            return self.id

        # try to append new coordinates to track
        # TODO: use vx/vy to improve accuracy
        dx    = cxn - self.cx
        dy    = cyn - self.cy
        vx    = -vn[0]
        vy    = -vn[1]
        dist  = hypot(dx,dy)
        area  = rn[2] * rn[3]

        if self.old_area is None:
            self.old_area = area

        # reject objects which change area too much from frame to frame
        if area > self.old_area:
            delta_area = area/self.old_area
        else:
            delta_area = self.old_area/area

        # reject too big area changes
        if delta_area > 15.0:
            return 0x00000000

        self.old_area = area
        found = 0x00000000

        # poor mans perspective
        # big nearby objects may move fast --- far away objects may move slow
        # max_dist = m*x + b
        max_dist = Track.maxDist
        fill_grade = area / Track.maxArea

        if(fill_grade > 0.2):
            m_factor = 10. if (fill_grade > 0.9) else 1.0 / (1.0 - fill_grade)
            #- enhace distance for big objects (just an empiric estimation)
            max_dist *= int(2.0 * m_factor)
            #- reduce maturity for big objects
            self.maturity = max(3,int(self.parent.trackMaturity / m_factor))

        vlength = 0.0
        oodir = self.old_dir
        # >>> debug
        #print("[%s](%d) xn/yn: %2d/%2d, vx/vy: %2d/%2d dist: %4.2f, delta_area: %4.2f" %
        #      (self.name,self.updates,rn[0],rn[1],vx,vy,dist,delta_area))
        # <<< debug

        # 1. is the new point in range?
        if dist >= 0.0 and dist < max_dist:
            # 2. is the new point in the right direction?
            # wait track to become mature and then check for angle
            if self.updates >= 3:
                # save first old status
                if self.old_dir is None:
                    dxo = self.cx - self.tr[-1][0]
                    dyo = self.cy - self.tr[-1][1]
                    self.old_dist = hypot(dxo,dyo)
                    self.old_dir  = np.array([dyo,dxo])

                # multiply vector magnitudes
                vlength = self.old_dist * dist
                #new_dir = np.array([vy ,vx])
                new_dir = np.array([dy ,dx])
                # accept all directions if one movement vector is zero
                if vlength <= 2.0:
                    cos_delta = 1.0
                else:
                    cos_delta = np.dot(self.old_dir, new_dir) / vlength

                self.old_dir  = new_dir
                self.old_dist = dist
                self.old_area = area
            else:
                cos_delta = 1.0

            # >>> debug
            #print("%s,%5d,%2d,%2d,%2d,%2d,%2d,%4.2f,%4.2f" %
            #      (self.name,frame,self.updates,rn[0],rn[1],rn[2],rn[3],cos_delta,dist))
            # <<< debug
            # reject all tracks out of direction
            #if self.updates < 3 or dist <= 1.1 or abs(cos_delta) > Track.minCosDelta:
            if self.updates < 3 or abs(cos_delta) > Track.minCosDelta:
                #- update base data
                found = self.id
                self.distance[0] += dx
                self.distance[1] += dy
                self.re = rn
                self.vv = np.array(vn)
                self.cx = cxn
                self.cy = cyn
                self.updates += 1
                self.tr.append([cxn,cyn])
                if(len(self.tr) > 24):
                    del self.tr[0]

                #- stop investigation if object is moving in the wrong direction
                if self.updates == 3:
                    if not self.checkStartingConditions():
                        self.reset()
                        return self.id

                #- is the coverered area still growing?
                self.updateGrowingStatus(rn)

                #- does the track leave the playground?
                if self.isLeaving(dx,dy):
                    self.reset()
                    return self.id

                # crossing status
                self.detectCrossing(dx,dy,rn)

                # turning status
                self.detectTurn(dx,dy,rn)

                # >>> debug
                #print("%s[%d] append %2d/%2d" %
                #      (self.name,self.updates-1,rn[0],rn[1]))
                # <<< debug
            else:
                # >>> debug
                #if oodir is not None:
                #    dxo = oodir[0]
                #    dyo = oodir[1]
                #    print("[%s](%d) delta-: (%4.2f) dx/dy: %4.2f/%4.2f dxo/dyo %4.2f/%4.2f dist: %4.2f, vlength: %4.2f" %
                #        (self.name, self.updates,degrees(acos(cos_delta)),dx,dy,dxo,dyo,dist,vlength))
                #print("     vx:%3d->%3d, vy:%3d->%3d dist: %4.2f" % (self.vv[0],vn[0], self.vv[1],vn[1],dist))
                # <<< debug
                ii = 0
        else:
            #if delta_area < 10.0:
            #print("[%s] dist: %4.2f > %4.2f delta_area: %4.2f" % (self.name,dist,max_dist,delta_area))
            ii = 0

        return found

    #--------------------------------------------------------------------
    #--
    #--------------------------------------------------------------------
    def showTrack(self, vis, frame=0, color=(220,0,0)):
        tsize = 2
        if self.updates > 3:
            ci = ord(self.name) % 3
            if ci == 0:
                color = (220,20,20)
            if ci == 1:
                color = (20,220,20)
            if ci == 2:
                color = (20,20,220)
        else:
            color = (80,80,80)
            tsize = 1

        r = self.re
        x = 8 * r[0]
        y = 8 * r[1]
        w = 8 * r[2]
        h = 8 * r[3]

        pts=np.int32(self.tr[-25:]) * 8
        cv2.polylines(vis, [pts], False, color,2)

        text = "%s(%d)" % (self.name, self.updates)
        if x > 3 and y > 3:
            cv2.putText(vis,text,(x-3,y-3),cv2.FONT_HERSHEY_SIMPLEX,0.5,color,tsize)
        else:
            cv2.putText(vis,text,(x+w+3,y+h+3),cv2.FONT_HERSHEY_SIMPLEX,0.5,color,tsize)
        hold = self.parent.greenLEDThread and self.parent.greenLEDThread.event.isSet()
        #if self.crossedY and hold:
        #    cv2.rectangle(vis,(x,y),(x+w,y+h),color,-4)
        #else:
        cv2.rectangle(vis,(x,y),(x+w,y+h),color,2)
        cv2.rectangle(vis,(8*self.minx,8*self.miny),(8*self.maxx,8*self.maxy),color,1)
        ###
        xm = int(x+w/2)
        ym = int(y+h/2)
        dx = -4 * self.vv[0]
        dy = -4 * self.vv[1]
        xe = int(xm+dx)
        ye = int(ym+dy)
        cv2.arrowedLine(vis,(xm,ym),(xe,ye),color,3)
        cv2.circle(vis,(xm,ym),3,color,1)

    #--------------------------------------------------------------------
    #--
    #--------------------------------------------------------------------
    def printTrack(self, frame=0):
        if self.progressx != 0 or self.progressy != 0:
            sys.stdout.write("[%s]:" %(self.name))
            for x,y in self.tr[-4:]:
                sys.stdout.write("  %02d,%02d -> " %(x,y))
            print("(#%2d vx:%02d vy:%02d) age:%d" % (self.updates,int(-self.vv[0]),int(-self.vv[1]),frame-self.lastFrame))

if __name__ == '__main__':

    from time import sleep
    #t = Track()
    #t1 = Track()
    #print("mask: %08x" % t.id)
    #t.new_track(1,[8,8,14,4],[0.0,-4.0])
    #t1.new_track(1,[1,1,4,4],[3.0,4.0])
    #t.update(2,[0,0,5,5],[3.0,4.0])
    #t.update(0x50, 3,3,5,5,3.0,4.0)
    #t.update(0xAA, 0,0,5,5,3.0,4.0)
    #t.update(0xAA, 1,1,5,4,2.0,4.0)
    #t.printAll()
    #t.reset(0x3f)
    #t.printAll()
    tracker = Tracker()
    tracker.setup_sizes(10,10)
    aa = []
    tracker.update_tracks(1, [[[5,5,2,2],[1.,1.]]])
    tracker.update_tracks(2, [[[6,6,2,2],[1.,1.]]])
    tracker.update_tracks(3, [[[7,7,2,2],[1.,1.]]])
    sleep(1)
    tracker.stop()
    tracker.join()


text:
import numpy as np
import cv2

x = np.load('output/scarlett-johansson_0.npy')
x = x * 255
cv2.imwrite('out.jpg', x)


text:
import sys
import glob
import os
import numpy as np

import matplotlib.pyplot as plt
import matplotlib.image as mpimg

import scipy
import imageio

import cv2
import pathlib
import skimage

import PIL
from PIL import Image

from interactive_crop import interactive_crop

def crop_datastack(exp_folder, alignment_channel):

    # Assert alignment_channel is correct
    assert alignment_channel in [561, 488]

    path = exp_folder + "\\rigid_align\\"
    storm_path = path + "\\storm_merged\\"
    conv_path = path + "\\conv_merged\\"
    conv_path_ds = path + "\\conv_merged_ds\\"
    align_path = path + "\\for_align\\"
    wga_path = path + "\\conv_{}\\".format(str(alignment_channel))

    tif_ext = r"*.tif"
    png_ext = r"*.png"

    wga_files = glob.glob(wga_path + tif_ext) + glob.glob(wga_path + png_ext)
    storm_files = glob.glob(storm_path + tif_ext) + glob.glob(storm_path + png_ext)
    conv_files = glob.glob(conv_path + tif_ext) + glob.glob(conv_path + png_ext)
    conv_files_ds = glob.glob(conv_path_ds + tif_ext) + glob.glob(conv_path_ds + png_ext)
    align_files = glob.glob(align_path + tif_ext) + glob.glob(align_path + png_ext)

    num_images = len(storm_files)

    C1 = []

    print("normalizing conv-1")

    # normalize intensity of conv_merged images
    for k in range(num_images):
        A = imageio.imread(conv_files_ds[k])
        C1.append(A)
    C1 = np.stack(C1, -1) # h x w x 3 x n_images
    C2 = C1.max(axis=3)

    # determine angle of rotation
    ang = -8
    C3 = skimage.transform.rotate(C2,ang)

    # interactive crop
    crop_reg = interactive_crop(C3)
    x_start, y_start, x_end, y_end = crop_reg
    C4 = C3[x_start:x_end, y_start:y_end, :]
    print('Cropping Coordinates: {}, Angle: {}'.format(crop_reg, ang))

	# Change here!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    # crop_storm = np.ceil(crop_reg*10)
    crop_storm = [0,0,0,0]
    crop_storm[0] = int(np.ceil(crop_reg[0]*10))
    crop_storm[1] = int(np.ceil(crop_reg[1]*10))
    crop_storm[2] = int(np.ceil(crop_reg[2]*10))
    crop_storm[3] = int(np.ceil(crop_reg[3]*10))
    #crop_storm_x_range = list(range(crop_storm[0], crop_storm[2]))
    #crop_storm_y_range = list(range(crop_storm[1], crop_storm[3]))
    crop_storm_x_range = range(crop_storm[1], crop_storm[3])
    crop_storm_y_range = range(crop_storm[0], crop_storm[2])
    #crop_reg = np.ceil(crop_reg)

    print('Making directories!')
    if not os.path.exists(exp_folder + "\\cropped\\"):
        os.mkdir(exp_folder + "\\cropped\\")
        os.mkdir(exp_folder + "\\cropped\\storm_merged\\")
        os.mkdir(exp_folder + "\\cropped\\storm_merged_ds\\")
        os.mkdir(exp_folder + "\\cropped\\conv_merged\\")
        os.mkdir(exp_folder + "\\cropped\\conv_merged_ds\\")
        os.mkdir(exp_folder + "\\cropped\\conv_{}\\".format(alignment_channel))
        os.mkdir(exp_folder + "\\cropped\\conv_{}_ds\\".format(alignment_channel))
        os.mkdir(exp_folder + "\\cropped\\for_align\\")
        os.mkdir(exp_folder + "\\cropped\\for_align_ds\\")

    print('Applying angle and crop to full size images and then saving out!')

	# Change from i to k here !!!!!!!!!!!!!!!!!!!!!!!!!!!
	# Delete 'storm_path+'!!!!!!!!!!!!!!!!!!!!!
	# cropping was incorrect.
    for k in range(num_images):
        print(k)
        A1 = cv2.imread(storm_files[k])
        A2 = skimage.transform.rotate(A1,ang)
        #A3 = (A2[crop_storm_x_range, crop_storm_y_range]).astype(np.float32)
        A3_1 = (A2[crop_storm_x_range]).astype(np.float32)
        A3 = (A3_1[:,crop_storm_y_range]).astype(np.float32)
        R_small = skimage.transform.rescale(A3[:,:,0] , 0.1)
        A3_small = np.zeros([R_small.shape[0],R_small.shape[1],3])
        A3_small[:,:,0] = R_small
        A3_small[:,:,1] = skimage.transform.rescale(A3[:,:,1], 0.1)
        A3_small[:,:,2] = skimage.transform.rescale(A3[:,:,2], 0.1)
        A3 *= 255
        A3 = A3.astype(np.uint8)
        cv2.imwrite(exp_folder+"/cropped/storm_merged/"+str(k).zfill(3) +'.tif', A3)
        A3_small *= 255
        A3_small = A3_small.astype(np.uint8)
        cv2.imwrite(exp_folder+"/cropped/storm_merged_ds/"+str(k).zfill(3) +'.tif', A3_small)

        A1c = cv2.imread(conv_files[k])
        A2c = skimage.transform.rotate(A1c,ang)
        A3_1c = (A2c[crop_storm_x_range]).astype(np.float32)
        A3c = (A3_1c[:,crop_storm_y_range]).astype(np.float32)
        Rc_small = skimage.transform.rescale(A3c[:,:,0] , 0.1)
        A3c_small = np.zeros([Rc_small.shape[0],Rc_small.shape[1],3])
        A3c_small[:,:,0] = Rc_small
        A3c_small[:,:,1] = skimage.transform.rescale(A3c[:,:,1], 0.1)
        A3c_small[:,:,2] = skimage.transform.rescale(A3c[:,:,2], 0.1)
        A3c *= 255
        A3c = A3c.astype(np.uint8)
        cv2.imwrite(exp_folder+"/cropped/conv_merged/"+str(k).zfill(3) +'.tif', A3c)
        A3c_small *= 255
        A3c_small = A3c_small.astype(np.uint8)
        cv2.imwrite(exp_folder+"/cropped/conv_merged_ds/"+str(k).zfill(3) +'.tif', A3c_small)

        A1w1 = cv2.imread(align_files[k])
        A2w1 = skimage.transform.rotate(A1w1, ang)
        A3_1w1 = (A2w1[crop_storm_x_range]).astype(np.float32)
        A3w1 = (A3_1w1[:,crop_storm_y_range]).astype(np.float32)
        Rw1_small = skimage.transform.rescale(A3w1[:,:,0] , 0.1)
        A3w1_small = np.zeros([Rw1_small.shape[0],Rw1_small.shape[1],3])
        A3w1_small[:,:,0] = Rw1_small
        A3w1_small[:,:,1] = skimage.transform.rescale(A3w1[:,:,1], 0.1)
        A3w1_small[:,:,2] = skimage.transform.rescale(A3w1[:,:,2], 0.1)
        A3w1 *= 255
        A3w1 = A3w1.astype(np.uint8)
        cv2.imwrite(exp_folder+"/cropped/for_align/"+str(k).zfill(3) +'.tif', A3w1)
        A3w1_small *= 255
        A3w1_small = A3w1_small.astype(np.uint8)
        cv2.imwrite(exp_folder+"/cropped/for_align_ds/"+str(k).zfill(3) +'.tif', A3w1_small)

        A1w = cv2.imread(wga_files[k])
        A2w = skimage.transform.rotate(A1w,ang)
        A3_1w = (A2w[crop_storm_x_range]).astype(np.float32)
        A3w = (A3_1w[:,crop_storm_y_range]).astype(np.float32)
        Rw_small = skimage.transform.rescale(A3w[:,:,0] , 0.1)
        A3w_small = np.zeros([Rw_small.shape[0],Rw_small.shape[1],3])
        A3w_small[:,:,0] = Rw_small
        A3w_small[:,:,1] = skimage.transform.rescale(A3w[:,:,1], 0.1)
        A3w_small[:,:,2] = skimage.transform.rescale(A3w[:,:,2], 0.1)
        A3w *= 255
        A3w = A3w.astype(np.uint8)
        cv2.imwrite(exp_folder+"/cropped/conv_{}/".format(alignment_channel)+str(k).zfill(3) +'.tif', A3w)
        A3w_small *= 255
        A3w_small = A3w_small.astype(np.uint8)
        cv2.imwrite(exp_folder+"/cropped/conv_{}_ds/".format(alignment_channel)+str(k).zfill(3) +'.tif', A3w_small)

##    # Load in new ds images and check max project
##
##    path = exp_folder + "\\cropped\\"
##    storm_path = exp_folder + "\\cropped\\storm_merged_ds\\"
##    conv_path = exp_folder + "\\cropped\\conv_merged_ds\\"
##    conv_path_ds = exp_folder + "\\cropped\\conv_{}_ds\\".format(alignment_channel)
##    wga_path = exp_folder + "\\cropped\\for_align_ds\\"
##
##    wga_files = glob.glob(wga_path + tif_ext) + glob.glob(wga_path + png_ext)
##    storm_files = glob.glob(storm_path + tif_ext) + glob.glob(storm_path + png_ext)
##    conv_files = glob.glob(conv_path + tif_ext) + glob.glob(conv_path + png_ext)
##    conv_files_ds = glob.glob(conv_path_ds + tif_ext) + glob.glob(conv_path_ds + png_ext)
##
##    num_images = len(storm_files)
##
##    C1 = []
##
##    print("normalizing conv-1 2.0")
##
##    # normalize intensity of conv_merged images
##    for k in range(num_images):
##        A = imageio.imread(conv_files_ds[k])
##        C1.append(A)
##    C1 = np.stack(C1, -1) # h x w x 3 x n_images
##
##    cv2.imwrite(exp_folder + '/cropped/conv_xy_projection.tif',
##                C1.max(axis=3)) # h x w x 3
##    cv2.imwrite(exp_folder + '/cropped/conv_xz_projection.tif',
##                C1.max(axis=1).squeeze().transpose(0,2,1))  # h x 3 x n_images -> h x n_images x 3
##    cv2.imwrite(exp_folder + '/cropped/conv_yz_projection.tif',
##                C1.max(axis=0).squeeze().transpose(0,2,1))  # w x 3 x n_images -> w x n_images x 3
##
##
##    print("normalizing storm-1 2.0")
##
##    # normalize intensity of storm_merged images
##    for k in range(num_images):
##        A = imageio.imread(storm_files[k])
##        C1.append(A)
##    C1 = np.stack(C1, -1) # h x w x 3 x n_images
##
##    cv2.imwrite(exp_folder + '/cropped/storm_xy_projection.tif',
##                C1.max(axis=3)) # h x w x 3
##    cv2.imwrite(exp_folder + '/cropped/storm_xz_projection.tif',
##                C1.max(axis=1).squeeze().transpose(0,2,1))  # h x 3 x n_images -> h x n_images x 3
##    cv2.imwrite(exp_folder + '/cropped/storm_yz_projection.tif',
##                C1.max(axis=0).squeeze().transpose(0,2,1))  # w x 3 x n_images -> w x n_images x 3
##
##    print("normalizing wga-1 2.0")
##    # normalize intensity of aligned images
##    for k in range(num_images):
##        A = imageio.imread(wga_files[k])
##        C1.append(A)
##    C1 = np.stack(C1, -1) # h x w x 3 x n_images
##
##    cv2.imwrite(exp_folder + '/cropped/wga_xy_projection.tif',
##                C1.max(axis=3)) # h x w x 3
##    cv2.imwrite(exp_folder + '/cropped/wga_xz_projection.tif',
##                C1.max(axis=1).squeeze().transpose(0,2,1))  # h x 3 x n_images -> h x n_images x 3
##    cv2.imwrite(exp_folder + '/cropped/wga_yz_projection.tif',
##                C1.max(axis=0).squeeze().transpose(0,2,1))  # w x 3 x n_images -> w x n_images x 3

    print('Done with cropping!')

    return True

max_stars_repo_name: SpeerLab/STORM-UI-

max_issues_repo_name: SpeerLab/STORM-UI

max_forks_repo_name: SpeerLab/STORM-UI

reason: import numpy,import scipy


text:
import numpy as np
import torch
from nqs_playground import *
from nqs_playground import _C

try:
    from typing_extensions import Final
except ImportError:
    # If you don't have `typing_extensions` installed, you can use a
    # polyfill from `torch.jit`.
    from torch.jit import Final

def test_exact_diagonalisation():
    n = 10
    basis = SpinBasis([], number_spins=n, hamming_weight=n // 2)
    hamiltonian = Heisenberg([(1.0, i, (i + 1) % n) for i in range(n)], basis)
    energy, eigenvector = diagonalise(hamiltonian)
    assert np.isclose(energy, -18.06178542)
    basis = SpinBasis(
        make_group(
            [
                Symmetry(list(range(1, n)) + [0], sector=n // 2),
                Symmetry(list(range(n))[::-1], sector=1),
            ]
        ),
        number_spins=n,
        hamming_weight=n // 2,
    )
    hamiltonian = Heisenberg([(1.0, i, (i + 1) % n) for i in range(n)], basis)
    energy, eigenvector = diagonalise(hamiltonian)
    assert np.isclose(energy, -18.06178542)

def test_sparse_diagonalisation():
    import scipy.sparse
    import scipy.sparse.linalg

    n = 10
    basis = SpinBasis([], number_spins=n, hamming_weight=n // 2)
    basis.build()
    hamiltonian = Heisenberg([(1.0, i, (i + 1) % n) for i in range(n)], basis)
    m = hamiltonian.to_csr()
    energy, _ = scipy.sparse.linalg.eigsh(m, k=1)
    assert np.isclose(energy, -18.06178542)

def reference_apply(spins, hamiltonian, psi):
    with torch.no_grad():
        basis = hamiltonian.basis
        states = torch.from_numpy(basis.states.view(np.int64))
        psi = psi(states)
        scale = torch.max(psi[:, 0]).item()
        psi[:, 0] -= scale
        psi = np.exp(psi.numpy().view(np.complex64))
        hamiltonian = hamiltonian.to_csr()
        out = np.empty(len(spins), dtype=np.complex64)
        for i, s in enumerate(spins):
            state = np.zeros(basis.number_states, dtype=np.complex64)
            state[basis.index(s)] = 1.0
            state = hamiltonian @ state
            out[i] = scale + np.log(np.dot(state.conj(), psi))
        return out

def reference_diag(spins, hamiltonian):
    basis = hamiltonian.basis
    hamiltonian = hamiltonian.to_csr()
    out = np.empty(len(spins), dtype=np.complex64)
    for i, s in enumerate(spins):
        state = np.zeros(basis.number_states, dtype=np.complex64)
        state[basis.index(s)] = 1.0
        state = hamiltonian @ state
        out[i] = state[basis.index(s)]
    return out

class Unpack(torch.nn.Module):
    n: Final[int]

    def __init__(self, n: int):
        super().__init__()
        self.n = n

    def forward(self, x):
        return unpack(x, self.n)

def test_apply():
    n = 20
    basis = SpinBasis([], number_spins=n, hamming_weight=n // 2)
    basis.build()
    hamiltonian = Heisenberg([(1.0, i, (i + 1) % n) for i in range(n)], basis)
    for i in range(5):
        spins = torch.from_numpy(basis.states.view(np.int64))[
            torch.randperm(basis.number_states)[:500]
        ]
        psi = torch.jit.script(
            torch.nn.Sequential(
                Unpack(n),
                torch.nn.Linear(n, 5),
                torch.nn.Tanh(),
                torch.nn.Linear(5, 2, bias=False),
            )
        )
        expected = reference_apply(spins, hamiltonian, psi)
        spins_512 = torch.cat(
            [spins.unsqueeze(dim=1), torch.zeros(spins.size(0), 7, dtype=torch.int64)],
            dim=1,
        )
        predicted = (
            _C.apply(spins_512, hamiltonian, psi._c._get_method("forward"))
            .numpy()
            .view(np.complex64)
            .squeeze()
        )
        assert np.allclose(predicted, expected)

def test_diag():
    n = 20
    basis = SpinBasis([], number_spins=n, hamming_weight=n // 2)
    basis.build()
    hamiltonian = Heisenberg([(1.0, i, (i + 1) % n) for i in range(n)], basis)
    spins = torch.from_numpy(basis.states.view(np.int64))[
        torch.randperm(basis.number_states)[:1000]
    ]
    expected = reference_diag(spins, hamiltonian)
    spins_512 = torch.cat(
        [spins.unsqueeze(dim=1), torch.zeros(spins.size(0), 7, dtype=torch.int64)],
        dim=1,
    )
    predicted = _C.diag(spins_512, hamiltonian).numpy().view(np.complex64).squeeze()
    assert np.allclose(predicted, expected)

max_stars_repo_name: twesterhout/nqs-playground

max_issues_repo_name: twesterhout/nqs-playground

max_forks_repo_name: twesterhout/nqs-playground


text:
import numpy as np
from pathos.multiprocessing import Pool
import os

import sys
sys.path.append(os.path.join(os.getcwd(), '../..'))

import src.sim_tool.evaluate as evaluate
from src.sim_tool.Env import Env
from src.model.MAP import MAP
import src.utility.plot_util as plot_util

def run_MAP(dict_env, env, tau, tot_user=100, num_thread=1):

    map = MAP(tau=tau, tot_chunk=env.tot_chunk, tot_EN=dict_env['tot_EN'],
              mu_xi=env.chunk_mu, sigma_xi=env.chunk_sigma, EN_start=0)
    pof = map.get_price_of_fog()

    # plot_util.fast_plot_MRI(map)

    def one_user_fn(user_id):
        quest_dict = evaluate.user_quests(size_file=dict_env['size_file'], size_chunk=dict_env['size_chunk'],
                                          bw_channel=dict_env['bw_channel'], trace_len=dict_env['trace_len'],
                                          speed_mu=dict_env['speed_mu'], speed_sigma=dict_env['speed_sigma'],
                                          tot_EN=dict_env['tot_EN'])
        hit_arr, \
        tot_hit, \
        tot_miss,\
        PoF, = evaluate.test_quests_on_cache(quest_dict=quest_dict, cache_dict=map.cache_dict,
                                                 tot_EN=dict_env['tot_EN'],
                                                 size_file=dict_env['size_file'], size_chunk=dict_env['size_chunk'])
        avg_hit_elem = tot_hit
        avg_miss_elem = tot_miss
        avg_PoF_elem = PoF

        return [avg_hit_elem, avg_miss_elem, avg_PoF_elem]

    pool = Pool(num_thread)
    outputs = np.asarray(list(pool.map(one_user_fn, range(tot_user))))
    pool.close()
    pool.join()
    avg_hit = np.average(outputs[:, 0])
    avg_miss = np.average(outputs[:, 1])
    avg_PoF = np.average(outputs[:, 2])

    return avg_hit, avg_miss, avg_PoF

if __name__ == '__main__':
    log_path = "../../results/log.csv"

    tau = 0.9

    for idx_test, speed_mu in enumerate(range(15, 34)):  # 15-33
        dict_env = {'speed_mu': speed_mu,
                    'speed_sigma': speed_mu / 8,
                    'bw_channel': 100,
                    'trace_len': 300,
                    'tot_EN': 50,
                    'size_chunk': 8 * 8,
                    'size_file': 5000 * 8,
                    'tot_user': 100}
        env = Env(size_file=dict_env['size_file'],
                  size_chunk=dict_env['size_chunk'],
                  bw_channel=dict_env['bw_channel'],
                  trace_len=dict_env['trace_len'],
                  speed_mu=dict_env['speed_mu'],
                  speed_sigma=dict_env['speed_sigma'])

        avg_hit, avg_miss, avg_PoF = run_MAP(dict_env=dict_env, env=env,
                                             tau=tau, tot_user=dict_env['tot_user'], num_thread=2)

        result_dict = {'algorithm': 'MAP'}
        for key in dict_env.keys():
            result_dict[key] = [dict_env[key]]
        result_dict['tau'] = [tau]
        result_dict['avg_hit'] = [avg_hit]
        result_dict['avg_miss'] = [avg_miss]
        result_dict['avg_PoF'] = [avg_PoF]

        plot_util.log_results(os.path.join(os.getcwd(),
                                           log_path), result_dict)

max_stars_repo_name: zhaoyanlyu/ACBC

max_issues_repo_name: zhaoyanlyu/ACBC

max_forks_repo_name: zhaoyanlyu/ACBC