import gc
import hydra
import logging
import numpy as np
import os
import pandas as pd
import sys
import torch
import torchvision.transforms as T
import traceback

from functools import partial
from hydra.utils import instantiate
from nvidia.dali.plugin.pytorch import DALIGenericIterator
from omegaconf import DictConfig
from torch import nn, Tensor
from typing import Any, Dict, List, Tuple

# TODO: make proper setup.py
sys.path.insert(0, f'/home/{os.environ["USER"]}/projects/dfdc/vendors/Pytorch_Retinaface')
from data import cfg_mnet, cfg_re50
from layers.functions.prior_box import PriorBox

from image import crop_square_torch
from dataset.utils import pad_torch
from detection_utils import max_num_faces
from detectors.retinaface import init_detector, prepare_imgs, postproc_detections_gpu
from video import VideoPipe, parse_meta, read_frames_cv2

def write_file_list(files: List[str], path: str) -> None:
    with open(path, mode='w') as h:
        for i, f in enumerate(files):
            if os.path.isfile(f):
                h.write(f'{f} {i}\n')

def split_files_by_res(files_meta: np.ndarray, min_freq: int
                       ) -> Tuple[List[np.ndarray], List[float], np.ndarray]:
    px_count = files_meta[:, 1] * files_meta[:, 2]
    clusters, freq = np.unique(px_count, return_counts=True)
    cluster_idxs, px_mults, heap = [], [], []

    for i, px_c in enumerate(clusters):
        idxs = (px_count == px_c).nonzero()[0]
        if freq[i] >= min_freq:
            cluster_idxs.append(idxs)
            px_m = float(px_c / (1920 * 1080))
            px_mults.append(px_m)
        else:
            heap.append(idxs)

    if len(heap) > 0:
        heap = np.hstack(heap)
        heap.sort()
    else:
        heap = np.empty(0, dtype=np.int64)
    return cluster_idxs, px_mults, heap

def merge_detector_cfg(conf: DictConfig) -> Dict[str, any]:
    cfg = cfg_mnet if conf.encoder == 'mnet' else cfg_re50
    for key in "batch_size, score_thresh, nms_thresh, top_k, keep_top_k".split(", "):
        if key not in conf or conf[key] is None:
            raise AttributeError("Missing {} in detector config".format(key))
    cfg = {**cfg, **conf}
    return cfg

def find_faces(frames: Tensor, model: torch.nn.Module,
               device: torch.device, conf: Dict[str, Any]) -> List[Tensor]:
    D, H, W, C = frames.shape
    frames_orig = frames.permute(0, 3, 1, 2)
    frames, scale = prepare_imgs(frames)
    prior_box = PriorBox(conf, image_size=(H, W))
    priors = prior_box.forward().to(device)
    scale = scale.to(device)

    with torch.no_grad():
        locations, confidence, landmarks = model(frames)
        detections = postproc_detections_gpu(
            locations, confidence, priors, scale, conf)

    num_faces = np.array(list(map(len, detections)), dtype=np.uint8)
    while (num_faces.mean() < conf['min_positive_rate'] and
           conf['score_thresh'] >= conf['score_thresh_min']):
        conf = dict(conf)
        conf['score_thresh'] -= conf['score_thresh_step']
        detections = postproc_detections_gpu(
            locations, confidence, priors, scale, conf)
        num_faces = np.array(list(map(len, detections)), dtype=np.uint8)
        logging.debug(
            "Rerun detection postprocessing with score_thresh={:.02f}, "
            "avg_pos_rate={:.02f}".format(conf['score_thresh'], num_faces.mean()))

    max_faces = max_num_faces(num_faces, conf['max_face_num_thresh'])
    del locations, confidence, landmarks, priors, prior_box, scale, frames

    faces = []
    for f in range(D):
        for bbox in detections[f][:max_faces]:
            face = crop_square_torch(frames_orig[f], bbox[:4])
            if face is not None:
                faces.append(face)
    del detections, frames_orig
    return faces

def load_model(conf: DictConfig):
    model = instantiate(conf.model)
    state = torch.load(conf.model.weights)
    if 'model' in state.keys():
        state = state['model']
    model.load_state_dict(state)
    for p in model.parameters():
        p.requires_grad_(False)
    model.eval()
    return model

@hydra.main(config_path="../config/predict.yaml")
def main(conf: DictConfig):
    print(conf.pretty())

    if 'gpu' in conf.general.keys():
        torch.cuda.set_device(conf.general.gpu)
    device = torch.device('cuda')
    torch.manual_seed(conf.general.seed)

    model = load_model(conf).to(device)

    face_det_conf = merge_detector_cfg(conf['face-detection'])
    detector = init_detector(face_det_conf, face_det_conf['weights'], device).to(device)
    crop_faces = partial(find_faces, model=detector, device=device, conf=face_det_conf)

    data_conf = conf.data.test
    base_dir = data_conf.dir
    file_names = [file for file in os.listdir(base_dir) if file.endswith('mp4')]
    files = [os.path.join(base_dir, file) for file in file_names]
    if not len(files):
        raise RuntimeError("No files was found in {}".format(base_dir))
    else:
        logging.info("Total number of files: {}".format(len(files)))
    logging.info("Parsing video metadata...")
    files_meta = parse_meta(files)
    splits, px_mults, leftover = split_files_by_res(files_meta, 5)
    px_mults_str = ' '.join(map(lambda f: '%.03f' % f, iter(px_mults)))
    logging.debug("Found {} clusters, px_mults are: {}. {} items are left over".format(
        len(splits), px_mults_str, len(leftover)))
    transforms = T.Compose([instantiate(val['transform']) for val in data_conf.transforms])
    logging.debug("Using transforms: {}".format(transforms))

    def _predict(images: List[Tensor]) -> float:
        last_idx = len(images) - 1
        num_samples = min(data_conf.sample.max_samples, len(images))
        idxs = np.linspace(0, last_idx, num_samples, dtype=int, endpoint=True)
        images = list(map(images.__getitem__, idxs))

        x = torch.stack(list(map(transforms, images)))
        pad_amount = reader_conf.frames - x.size(0)
        if pad_amount > 0:
            x = pad_torch(x, pad_amount, 'start')
        # D, C, H, W -> C, D, H, W
        x = x.transpose(0, 1).unsqueeze_(0)
        with torch.no_grad():
            out = model(x, None)
        y_pred = model.to_y(*out).cpu().item()
        return y_pred

    def _handle(images: List[Tensor], split_idx: int, pipe_name: str):
        glob_idx = split_idxs[split_idx]
        file_path = split_files[split_idx]
        if len(images) > 0:
            y_hat = _predict(images)
            df.loc[glob_idx, 'label'] = y_hat
            logging.info("{} | {} | faces: {} | y: {:.03f} | {}".format(
                pipe_name, glob_idx, len(images), y_hat, file_path))
        else:
            ignore_split_files[split_idx] = True
            logging.warning("No faces have found in ({}): {}".format(glob_idx, file_path))

    def _on_exception(e: Exception):
        nonlocal video_batch, frames, faces
        logging.error(traceback.format_exc())
        video_batch, frames = None, None
        faces = []
        gc.collect()

    def _save():
        save_dir = conf.get('general.save_dir', os.getcwd())
        save_path = os.path.join(save_dir, conf.general.csv_name)
        logging.info("Saving predictions to {}".format(save_path))
        for _ in range(5):
            try:
                df.to_csv(save_path, index=False)
                break
            except Exception as e:
                continue

    reader_conf = data_conf.sample
    df = pd.DataFrame(file_names, columns=['filename'])
    df['label'] = 0.5
    del file_names, files_meta

    for s, split_idxs in enumerate(splits):
        split_files = list(map(lambda i: files[i], split_idxs))
        logging.debug("Num files in cluster {}: {}".format(s, len(split_files)))
        handled_split_files = np.zeros(len(split_files), dtype=np.bool)
        ignore_split_files = np.zeros(len(split_files), dtype=np.bool)
        seq_len = int(reader_conf.frames / reader_conf.num_pass / max(px_mults[s], 1.0))
        stride = reader_conf.stride
        logging.debug(
            "Running predictions on cluster {} (size_factor={:.03f})".format(s, px_mults[s]))
        step = data_conf.max_open_files

        for offset in range(0, len(split_files), step):
            last = min(offset + step, len(split_files))
            logging.debug("Running chunk [{}:{}] from {} cluster".format(offset, last, s))
            write_file_list(split_files[offset:last], path='temp')
            logging.debug("Creating new pipe with seq_len={}, stride={}".format(seq_len, stride))
            pipe = VideoPipe('temp', seq_len=seq_len, stride=stride, device_id=0)
            pipe.build()
            num_samples_read = pipe.epoch_size('reader')
            logging.debug("Pipe length: {}".format(num_samples_read))

            if num_samples_read > 0:
                logging.debug("Creating DALI iterator...")
                data_iter = DALIGenericIterator(
                    [pipe], ['frames', 'label'], num_samples_read, dynamic_shape=True)
                video_batch, frames, prev_idx = None, None, None
                faces = []
                while True:
                    try:
                        video_batch = next(data_iter)
                        frames = video_batch[0]['frames'].squeeze(0)
                        read_idx = video_batch[0]['label'].item()
                        new_faces = crop_faces(frames)
                        del video_batch, frames

                        if prev_idx is None or prev_idx == read_idx:
                            faces += new_faces
                        else:
                            abs_idx = offset + prev_idx
                            _handle(faces, abs_idx, 'dali')
                            handled_split_files[abs_idx] = True
                            faces = new_faces
                        prev_idx = read_idx
                    except StopIteration:
                        break
                    except Exception as e:
                        _on_exception(e)
                # last sample in pipe
                try:
                    abs_idx = offset + prev_idx
                    _handle(faces, abs_idx, 'dali')
                    handled_split_files[abs_idx] = True
                except Exception as e:
                    _on_exception(e)

                del pipe, data_iter
                gc.collect()
                _save()
            logging.debug("Finished chunk [{}:{}] from {} cluster".format(offset, last, s))

        unhandled_files = (~handled_split_files).nonzero()[0]
        num_bad_samples = len(unhandled_files)
        if num_bad_samples > 0:
            logging.warning("Unable to read %d videos with DALI, running fallback "
                            "decoding through OpenCV..." % num_bad_samples)
            for idx in unhandled_files:
                abs_idx = split_idxs[idx]
                path = split_files[idx]
                if ignore_split_files[idx]:
                    logging.info("Ignoring file ({}): {}".format(abs_idx, path))
                    continue
                try:
                    frames = read_frames_cv2(path, reader_conf.frames)
                    if frames is not None and len(frames) > 0:
                        frames = torch.from_numpy(frames).to(device)
                        faces = crop_faces(frames)
                        del frames
                        _handle(faces, idx, 'cv2')
                        del faces
                except Exception as e:
                    _on_exception(e)
        del handled_split_files, unhandled_files, ignore_split_files
        _save()
        logging.debug("Finished cluster {} (size_factor={:.03f})".format(s, px_mults[s]))

    logging.info("Running predictions on leftover ({} samples)".format(len(leftover)))
    for idx in leftover:
        path = files[idx]
        try:
            frames = read_frames_cv2(path, reader_conf.frames)
            if frames is not None and len(frames) > 0:
                frames = torch.from_numpy(frames).to(device)
                faces = crop_faces(frames)
                del frames
                if len(faces) > 0:
                    y_hat = _predict(faces)
                    df.loc[idx, 'label'] = y_hat
                    logging.info("cv2 | {} | faces: {} | y: {:.03f} | {}".format(
                        idx, len(faces), y_hat, path))
                else:
                    logging.warning("No faces have found in ({}): {}".format(idx, path))
                del faces
        except Exception as e:
            _on_exception(e)
    logging.info("Finished leftover")

    _save()
    logging.info("DONE")

if __name__ == '__main__':
    assert torch.cuda.is_available()
    torch.backends.cudnn.benchmark = True
    main()


#!/usr/bin/env python3
""" Facial landmarks extractor for faceswap.py
    Code adapted and modified from:
    https://github.com/1adrianb/face-alignment
"""
import cv2
import numpy as np

from lib.model.session import KSession
from ._base import Aligner, logger

class Align(Aligner):
    """ Perform transformation to align and get landmarks """
    def __init__(self, **kwargs):
        git_model_id = 13
        model_filename = "face-alignment-network_2d4_keras_v2.h5"
        super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs)
        self.name = "FAN"
        self.input_size = 256
        self.color_format = "RGB"
        self.vram = 2240
        self.vram_warnings = 512  # Will run at this with warnings
        self.vram_per_batch = 64
        self.batchsize = self.config["batch-size"]
        self.reference_scale = 200. / 195.

    def init_model(self):
        """ Initialize FAN model """
        self.model = KSession(self.name,
                              self.model_path,
                              allow_growth=self.config["allow_growth"],
                              exclude_gpus=self._exclude_gpus)
        self.model.load_model()
        # Feed a placeholder so Aligner is primed for Manual tool
        placeholder_shape = (self.batchsize, self.input_size, self.input_size, 3)
        placeholder = np.zeros(placeholder_shape, dtype="float32")
        self.model.predict(placeholder)

    def process_input(self, batch):
        """ Compile the detected faces for prediction """
        logger.debug("Aligning faces around center")
        batch["center_scale"] = self.get_center_scale(batch["detected_faces"])
        faces = self.crop(batch)
        logger.trace("Aligned image around center")
        faces = self._normalize_faces(faces)
        batch["feed"] = np.array(faces, dtype="float32")[..., :3] / 255.0
        return batch

    def get_center_scale(self, detected_faces):
        """ Get the center and set scale of bounding box """
        logger.debug("Calculating center and scale")
        center_scale = np.empty((len(detected_faces), 68, 3), dtype='float32')
        for index, face in enumerate(detected_faces):
            x_center = (face.left + face.right) / 2.0
            y_center = (face.top + face.bottom) / 2.0 - face.h * 0.12
            scale = (face.w + face.h) * self.reference_scale
            center_scale[index, :, 0] = np.full(68, x_center, dtype='float32')
            center_scale[index, :, 1] = np.full(68, y_center, dtype='float32')
            center_scale[index, :, 2] = np.full(68, scale, dtype='float32')
        logger.trace("Calculated center and scale: %s", center_scale)
        return center_scale

    def crop(self, batch):  # pylint:disable=too-many-locals
        """ Crop image around the center point """
        logger.debug("Cropping images")
        sizes = (self.input_size, self.input_size)
        batch_shape = batch["center_scale"].shape[:2]
        resolutions = np.full(batch_shape, self.input_size, dtype='float32')
        matrix_ones = np.ones(batch_shape + (3,), dtype='float32')
        matrix_size = np.full(batch_shape + (3,), self.input_size, dtype='float32')
        matrix_size[..., 2] = 1.0
        upper_left = self.transform(matrix_ones, batch["center_scale"], resolutions)
        bot_right = self.transform(matrix_size, batch["center_scale"], resolutions)

        # TODO second pass .. convert to matrix
        new_images = []
        for image, top_left, bottom_right in zip(batch["image"], upper_left, bot_right):
            height, width = image.shape[:2]
            channels = 3 if image.ndim > 2 else 1
            bottom_right_width, bottom_right_height = bottom_right[0].astype('int32')
            top_left_width, top_left_height = top_left[0].astype('int32')
            new_dim = (bottom_right_height - top_left_height,
                       bottom_right_width - top_left_width,
                       channels)
            new_img = np.empty(new_dim, dtype=np.uint8)

            new_x = slice(max(0, -top_left_width), min(bottom_right_width, width) - top_left_width)
            new_y = slice(max(0, -top_left_height),
                          min(bottom_right_height, height) - top_left_height)
            old_x = slice(max(0, top_left_width), min(bottom_right_width, width))
            old_y = slice(max(0, top_left_height), min(bottom_right_height, height))
            new_img[new_y, new_x] = image[old_y, old_x]

            interp = cv2.INTER_CUBIC if new_dim[0] < self.input_size else cv2.INTER_AREA
            new_images.append(cv2.resize(new_img, dsize=sizes, interpolation=interp))
        logger.trace("Cropped images")
        return new_images

    @staticmethod
    def transform(points, center_scales, resolutions):
        """ Transform Image """
        logger.debug("Transforming Points")
        num_images, num_landmarks = points.shape[:2]
        transform_matrix = np.eye(3, dtype='float32')
        transform_matrix = np.repeat(transform_matrix[None, :], num_landmarks, axis=0)
        transform_matrix = np.repeat(transform_matrix[None, :, :], num_images, axis=0)
        scales = center_scales[:, :, 2] / resolutions
        translations = center_scales[..., 2:3] * -0.5 + center_scales[..., :2]
        transform_matrix[:, :, 0, 0] = scales  # x scale
        transform_matrix[:, :, 1, 1] = scales  # y scale
        transform_matrix[:, :, 0, 2] = translations[:, :, 0]  # x translation
        transform_matrix[:, :, 1, 2] = translations[:, :, 1]  # y translation
        new_points = np.einsum('abij, abj -> abi', transform_matrix, points, optimize='greedy')
        retval = new_points[:, :, :2].astype('float32')
        logger.trace("Transformed Points: %s", retval)
        return retval

    def predict(self, batch):
        """ Predict the 68 point landmarks """
        logger.debug("Predicting Landmarks")
        # TODO Remove lazy transpose and change points from predict to use the correct
        # order
        batch["prediction"] = self.model.predict(batch["feed"])[-1].transpose(0, 3, 1, 2)
        logger.trace(batch["prediction"].shape)
        return batch

    def process_output(self, batch):
        """ Process the output from the model """
        self.get_pts_from_predict(batch)
        return batch

    def get_pts_from_predict(self, batch):
        """ Get points from predictor """
        logger.debug("Obtain points from prediction")
        num_images, num_landmarks, height, width = batch["prediction"].shape
        image_slice = np.repeat(np.arange(num_images)[:, None], num_landmarks, axis=1)
        landmark_slice = np.repeat(np.arange(num_landmarks)[None, :], num_images, axis=0)
        resolution = np.full((num_images, num_landmarks), 64, dtype='int32')
        subpixel_landmarks = np.ones((num_images, num_landmarks, 3), dtype='float32')

        flat_indices = batch["prediction"].reshape(num_images, num_landmarks, -1).argmax(-1)
        indices = np.array(np.unravel_index(flat_indices, (height, width)))
        min_clipped = np.minimum(indices + 1, height - 1)
        max_clipped = np.maximum(indices - 1, 0)
        offsets = [(image_slice, landmark_slice, indices[0], min_clipped[1]),
                   (image_slice, landmark_slice, indices[0], max_clipped[1]),
                   (image_slice, landmark_slice, min_clipped[0], indices[1]),
                   (image_slice, landmark_slice, max_clipped[0], indices[1])]
        x_subpixel_shift = batch["prediction"][offsets[0]] - batch["prediction"][offsets[1]]
        y_subpixel_shift = batch["prediction"][offsets[2]] - batch["prediction"][offsets[3]]
        # TODO improve rudimentary sub-pixel logic to centroid of 3x3 window algorithm
        subpixel_landmarks[:, :, 0] = indices[1] + np.sign(x_subpixel_shift) * 0.25 + 0.5
        subpixel_landmarks[:, :, 1] = indices[0] + np.sign(y_subpixel_shift) * 0.25 + 0.5

        batch["landmarks"] = self.transform(subpixel_landmarks, batch["center_scale"], resolution)
        logger.trace("Obtained points from prediction: %s", batch["landmarks"])


import numpy as np
from skrobot.coordinates import Coordinates
from skrobot.coordinates.math import rotate_vector
from skrobot.coordinates.math import rotation_distance

import actionlib
import actionlib_msgs.msg
import control_msgs.msg
import dynamic_reconfigure.msg
import dynamic_reconfigure.srv
import geometry_msgs.msg
import move_base_msgs.msg
import nav_msgs.msg
import rospy
import std_srvs.srv
import trajectory_msgs.msg

from .base import ROSRobotInterfaceBase
from .tf_utils import coords_to_geometry_pose
from .tf_utils import geometry_pose_to_coords
from .tf_utils import tf_pose_to_coords
from .transform_listener import TransformListener

class ROSRobotMoveBaseInterface(ROSRobotInterfaceBase):

    def __init__(self, *args, **kwargs):
        self.move_base_action = None
        self.move_base_trajectory_action = None
        self.move_base_goal_msg = None
        self.move_base_goal_coords = None
        self.move_base_goal_map_to_frame = None
        self.odom_topic = None
        self.go_pos_unsafe_goal_msg = None
        self.current_goal_coords = None
        self.map_frame_id = kwargs.pop(
            'map_frame_id', 'map')
        self.move_base_action_name = kwargs.pop(
            'move_base_action_name', 'move_base')
        self.base_frame_id = kwargs.pop(
            'base_frame_id',
            'base_footprint')
        self.base_controller_action_name = kwargs.pop(
            'base_controller_action_name',
            "/base_controller/follow_joint_trajectory")
        self.move_base_trajectory_joint_names = kwargs.pop(
            'base_controller_joint_names',
            ["base_link_x", "base_link_y", "base_link_pan"])
        self.move_base_simple_name = kwargs.pop(
            'move_base_simple_name',
            'move_base_simple')
        self.odom_topic = kwargs.pop('odom_topic', '/base_odometry/odom')
        self.use_tf2 = kwargs.pop('use_tf2', False)

        super(ROSRobotMoveBaseInterface, self).__init__(
            *args, **kwargs)

        self.tf_listener = TransformListener(
            use_tf2=self.use_tf2)

        self.move_base_action = actionlib.SimpleActionClient(
            self.move_base_action_name,
            move_base_msgs.msg.MoveBaseAction)

        if self.base_controller_action_name:
            self.move_base_trajectory_action = actionlib.SimpleActionClient(
                self.base_controller_action_name,
                control_msgs.msg.FollowJointTrajectoryAction)
            if self.move_base_trajectory_action.wait_for_server(
                    rospy.Duration(3)) is False:
                rospy.logwarn('{} is not found'.format(
                    self.base_controller_action_name))
                self.move_base_trajectory_action = None

        self.go_pos_unsafe_goal_msg = None

        self.move_base_simple_publisher = rospy.Publisher(
            "{}/goal".format(self.move_base_simple_name),
            geometry_msgs.msg.PoseStamped,
            queue_size=1)

        self.odom_msg = None
        self.odom_subscriber = rospy.Subscriber(
            self.odom_topic, nav_msgs.msg.Odometry,
            callback=self.odom_callback,
            queue_size=1)

    @property
    def odom(self):
        """Return Coordinates of this odom

        Returns
        -------
        odom_coords : skrobot.coordinates.Coordinates
            coordinates of odom.
        """
        if self.odom_msg is None:
            raise RuntimeError(
                'odom is not set. Please check odom topic {}'
                ' is published.'.format(self.odom_topic))
        pos = (self.odom_msg.pose.pose.position.x,
               self.odom_msg.pose.pose.position.y,
               self.odom_msg.pose.pose.position.z)
        q_wxyz = (self.odom_msg.pose.pose.orientation.w,
                  self.odom_msg.pose.pose.orientation.x,
                  self.odom_msg.pose.pose.orientation.y,
                  self.odom_msg.pose.pose.orientation.z)
        return Coordinates(pos=pos, rot=q_wxyz)

    def odom_callback(self, msg):
        """ROS's subscriber callback for odom.

        Parameters
        ----------
        msg : nav_msgs.msg.Odometry
            odometry message.
        """
        self.odom_msg = msg

    def go_stop(self, force_stop=True):
        """Cancel move_base.

        Parameters
        ----------
        force_stop : bool
            if force_stop is True, send go_velocity(0, 0, 0)
        """
        if self.joint_action_enable is True:
            ret = self.move_base_action.cancel_all_goals()
        if force_stop is True:
            self.go_velocity(0.0, 0.0, 0.0)
        return ret

    def move_to(self, coords, wait=True,
                frame_id=None):
        """Move Robot to target coords.

        Parameters
        ----------
        coords : str or skrobot.coordinates.Coordinates
            target tf name or target coords.
        wait : bool
            if True, wait until move end.
        frame_id : None or str

        Returns
        -------
        result : bool
            if move_base is succeeded, return True.
        """
        if frame_id is None:
            frame_id = self.base_frame_id
        if isinstance(coords, str):
            base_to_target = self.tf_listener.lookup_transform(
                self.base_frame_id, coords,
                rospy.Time(0),
                rospy.Duration(3))
            if base_to_target is False:
                rospy.logwarn('Could not lookup transform {} to {}'.
                              format(self.base_frame_id, coords))
                return False
            coords = tf_pose_to_coords(base_to_target)
        if self.move_to_send(coords, frame_id=frame_id) is False:
            return False
        if wait:
            return self.move_to_wait(frame_id=frame_id)
        return True

    def move_to_send(self,
                     coords,
                     frame_id=None,
                     wait_for_server_timeout=5.0):
        """Send MoveBaseAction

        Parameters
        ----------
        coords : skrobot.coordinates.Coordinates

        Return
        ------
        result : bool
            False or True. If False, could not send MoveBaseAction

        """
        if frame_id is None:
            frame_id = self.base_frame_id
        self.move_base_goal_msg = move_base_msgs.msg.MoveBaseActionGoal()
        self.move_base_goal_coords = coords
        count = 0
        ros_time = rospy.Time.now()
        map_to_frame = self.tf_listener.lookup_transform(
            self.map_frame_id, frame_id, rospy.Time(0))
        if map_to_frame is False:
            rospy.logwarn("Could not lookup transform {} to {}".format(
                self.map_frame_id, frame_id))
            return False
        map_to_frame_transform = tf_pose_to_coords(map_to_frame)
        # store in member variable for self.move_to_wait()
        self.move_base_goal_map_to_frame = map_to_frame_transform

        if self.move_base_action.wait_for_server(
                rospy.Duration(wait_for_server_timeout)) is False:
            return False

        self.move_base_goal_msg.header.stamp = ros_time
        self.move_base_goal_msg.goal.target_pose.header.stamp = ros_time
        if map_to_frame:
            self.move_base_goal_msg.goal.target_pose.header.frame_id = \
                self.map_frame_id
            self.move_base_goal_msg.goal.target_pose.pose = \
                coords_to_geometry_pose(
                    coords.copy_worldcoords().transform(
                        map_to_frame_transform, 'world'))
        else:
            self.move_base_goal_msg.goal.target_pose.header.frame_id = frame_id
            self.move_base_goal_msg.goal.target_pose.pose = \
                coords_to_geometry_pose(coords)
        self.move_base_goal_msg.header.seq = count
        self.move_base_action.send_goal(self.move_base_goal_msg.goal)
        return self.move_base_goal_msg

    def move_to_wait(self,
                     retry=10,
                     frame_id='world'):
        count = 0
        if self.move_base_goal_msg is None:
            return False

        ret = False
        while ret is False and count < retry:
            if count > 0:
                self.clear_costmap()
                self.move_base_goal_msg.header.seq = count
                self.move_base_goal_msg.goal.target_pose.header.seq = count
                self.move_base_action.send_goal(self.move_base_goal_msg.goal)
            self.move_base_action.wait_for_result()
            if self.move_base_action.get_state() == \
               actionlib_msgs.msg.GoalStatus.PREEMPTED:
                ret = False
                continue
            if self.move_base_action.get_state() == \
               actionlib_msgs.msg.GoalStatus.SUCCEEDED:
                ret = True
            count += 1
        return ret

    def _calc_move_diff_coords(self, frame_id):
        if frame_id == self.base_frame_id:
            map_goal_coords = self.move_base_goal_map_to_frame.\
                copy_worldcoords().transform(
                    self.move_base_goal_coords.copy_worldcoords())
        else:
            map_to_frame = self.tf_listener.lookup_transform(
                self.map_frame_id, frame_id,
                rospy.Time(0))
            if map_to_frame is False:
                rospy.logwarn("Could not lookup transform {} to {}".format(
                    self.map_frame_id, frame_id))
                return False
            map_to_frame_transform = tf_pose_to_coords(
                map_to_frame)
            map_goal_coords = map_to_frame_transform.transform(
                self.move_base_goal_coords.copy_worldcoords())

        current_coords = self.tf_listener.lookup_transform(
            self.map_frame_id, self.base_frame_id,
            rospy.Time(0))
        if current_coords is False:
            rospy.logwarn("Could not lookup transform {} to {}".format(
                self.map_frame_id, self.base_frame_id))
            return False
        current_coords = tf_pose_to_coords(current_coords)
        diff_coords = current_coords.transformation(
            map_goal_coords).copy_worldcoords()
        rospy.logwarn("move_to error translation: {}, rotation: {}".
                      format(diff_coords.translation,
                             diff_coords.rpy_angle()[0]))
        return diff_coords

    def go_pos(self, x=0.0, y=0.0, yaw=0.0, wait=True):
        """Move Robot using MoveBase

        Parameters
        ----------
        x : float
            move distance with respect to x axis. unit is [m].
        y : float
            move distance with respect to y axis. unit is [m].
        yaw : float
            rotate angle. unit is [rad].
        wait : bool
            if wait is True, wait until move base done.
        """
        c = Coordinates(pos=(x, y, 0)).rotate(yaw, 'z')
        return self.move_to(c, wait=wait, frame_id=self.base_frame_id)

    def go_pos_unsafe(self, x=0.0, y=0.0, yaw=0.0, wait=False):
        """Move Robot using MoveBase

        Parameters
        ----------
        x : float
            move distance with respect to x axis. unit is [m].
        y : float
            move distance with respect to y axis. unit is [m].
        yaw : float
            rotate angle. unit is [rad].
        wait : bool
            if wait is True, wait until stop go_pos_unsafe
        """
        self.go_pos_unsafe_no_wait(x=x, y=y, yaw=yaw)
        if wait is True:
            return self.go_pos_unsafe_wait()
        else:
            return True

    def go_pos_unsafe_no_wait(self, x=0.0, y=0.0, yaw=0.0, sec=None):
        """Move Robot using MoveBase

        Parameters
        ----------
        x : float
            move distance with respect to x axis. unit is [m].
        y : float
            move distance with respect to y axis. unit is [m].
        yaw : float
            rotate angle. unit is [rad].
        """
        if self.move_base_trajectory_action is None:
            rospy.logwarn("go_pos_unsafe is disabled. "
                          'move_base_trajectory_action is not found')
            return True
        maxvel = 0.295
        maxrad = 0.495
        ratio = 0.8
        if sec is None:
            sec = max(np.linalg.norm([x, y]) / (maxvel * ratio),
                      abs(yaw) / (maxrad * ratio),
                      1.0)
        self.go_pos_unsafe_goal_msg = self.move_trajectory(
            x, y, yaw,
            sec, stop=True)
        return self.move_base_trajectory_action.send_goal(
            self.go_pos_unsafe_goal_msg.goal)

    def go_pos_unsafe_wait(self, wait_counts=3):
        maxvel = 0.295
        maxrad = 0.495
        ratio = 0.8
        counter = 0
        if self.go_pos_unsafe_goal_msg is None:
            return False
        if self.move_base_trajectory_action is None:
            rospy.logwarn("go_pos_unsafe_wait is disabled. "
                          'move_base_trajectory_action is not found')
            return True
        while counter < wait_counts:
            action_ret = self.move_base_trajectory_action.wait_for_result()
            if action_ret is False:
                return False

            goal_position = np.array(
                self.go_pos_unsafe_goal_msg.goal.
                trajectory.points[1].positions,
                dtype=np.float32)
            odom = self.odom
            odom_pos = odom.translation
            odom_angle = odom.rpy_angle()[0][0]
            diff_position = goal_position - (
                odom_pos + np.array((0, 0, odom_angle)))
            v = rotate_vector(
                np.array((diff_position[0], diff_position[1], 0.0)),
                -odom_angle, 'z') - np.array((0, 0, odom_angle))
            x = v[0]
            y = v[1]
            yaw = diff_position[2]
            if yaw > 2 * np.pi:
                yaw = yaw - 2 * np.pi
            if yaw < - 2 * np.pi:
                yaw = yaw + 2 * np.pi

            sec = max(np.linalg.norm([x, y]) / (maxvel * ratio),
                      abs(yaw) / (maxrad * ratio))
            sec = max(sec, 1.0)
            step = 1.0 / sec

            rospy.loginfo("                diff-pos {} {}, diff-angle {}".
                          format(x, y, yaw))
            if np.sqrt(x * x + y * y) <= 0.025 and \
               abs(yaw) <= np.deg2rad(2.5) and \
               counter > 0:  # try at least 1 time
                self.go_pos_unsafe_goal_msg = None
                return True

            self.go_pos_unsafe_goal_msg = self.move_trajectory(
                x * step, y * step, yaw * step, sec, stop=True)
            self.move_base_trajectory_action.send_goal(
                self.go_pos_unsafe_goal_msg.goal)
            counter += 1
        self.go_pos_unsafe_goal_msg = None
        return True

    def go_velocity(self, x=0.0, y=0.0, yaw=0.0,
                    sec=1.0, stop=True, wait=False):
        """Move Robot using MoveBase

        Parameters
        ----------
        x : float
            move velocity with respect to x axis. unit is [m/s].
        y : float
            move velocity with respect to y axis. unit is [m/s].
        yaw : float
            rotate angle. unit is [rad/s].
        sec : float
            time.
        wait : bool
            if wait is True, wait until move base done.
        """
        if self.move_base_trajectory_action is None:
            rospy.logwarn("go_velocity is disabled. "
                          'move_base_trajectory_action is not found')
            return True
        x = x * sec
        y = y * sec
        yaw = yaw * sec
        goal = self.move_trajectory(x, y, yaw, sec, stop=stop)
        ret = self.move_base_trajectory_action.send_goal(goal.goal)
        if wait:
            self.move_base_trajectory_action.wait_for_result()
        return ret

    def move_trajectory_sequence(self, trajectory_points, time_list, stop=True,
                                 start_time=None, send_action=None):
        """Move base following the trajectory points at each time points

        trajectory-points [ list of #f(x y yaw) ([m] for x, y; [rad] for yaw) ]
        time-list [list of time span [sec] ]
        stop [ stop after msec moveing ]
        start-time [ robot will move at start-time [sec or ros::Time] ]
        send-action [ send message to action server, it means robot will move ]

        """
        if len(trajectory_points) != len(time_list):
            raise ValueError
        while self.odom_msg is None:
            rospy.sleep(0.01)
        odom_coords = geometry_pose_to_coords(self.odom_msg.pose.pose)
        goal = control_msgs.msg.FollowJointTrajectoryActionGoal()
        msg = trajectory_msgs.msg.JointTrajectory()
        msg.joint_names = self.move_base_trajectory_joint_names
        if start_time is not None:
            msg.header.stamp = start_time
        else:
            msg.header.stamp = rospy.Time.now()
        coords_list = [odom_coords]
        for traj_point in trajectory_points:
            cds = Coordinates(pos=odom_coords.translation,
                              rot=odom_coords.rotation)
            cds.translate((traj_point[0], traj_point[1], 0.0))
            cds.rotate(traj_point[2], 'z')
            coords_list.append(cds)
        cur_cds = odom_coords
        cur_yaw = cur_cds.rpy_angle()[0][0]
        cur_time = 0
        pts_msg_list = []
        for i, (cur_cds, next_cds) in enumerate(zip(
                coords_list[:-1],
                coords_list[1:])):
            next_time = time_list[i]
            tra = cur_cds.transformation(next_cds)
            rot = cur_cds.rotate_vector(tra.translation)
            diff_yaw = rotation_distance(cur_cds.rotation, next_cds.rotation)
            pts_msg_list.append(
                trajectory_msgs.msg.JointTrajectoryPoint(
                    positions=[cur_cds.translation[0],
                               cur_cds.translation[1],
                               cur_yaw],
                    velocities=[rot[0] / next_time,
                                rot[1] / next_time,
                                tra.rpy_angle()[0][0] / next_time],
                    time_from_start=rospy.Time(cur_time)))
            cur_time += next_time
            cur_cds = next_cds
            if tra.rpy_angle()[0][0] > 0:
                cur_yaw += abs(diff_yaw)
            else:
                cur_yaw -= abs(diff_yaw)

        # append last point
        if stop:
            velocities = [0, 0, 0]
        else:
            velocities = [rot[0] / next_time,
                          rot[1] / next_time,
                          tra.rpy_angle()[0][0] / next_time]
        pts_msg_list.append(
            trajectory_msgs.msg.JointTrajectoryPoint(
                positions=[cur_cds.translation[0],
                           cur_cds.translation[1],
                           cur_yaw],
                velocities=velocities,
                time_from_start=rospy.Time(cur_time)))

        msg.points = pts_msg_list
        goal.goal.trajectory = msg
        if send_action:
            if self.move_base_trajectory_action is None:
                rospy.logerror(
                    'send_action is True, '
                    'but move_base_trajectory_action is not found')
                return False
            self.move_base_trajectory_action.send_goal(goal.goal)
            """
            if self.move_base_trajectory_action.wait_for_result():
                return self.move_base_trajectory_action.get_result()
            else:
                return False
            """
        return goal

    def move_trajectory(self, x, y, yaw, sec=1.0, stop=True,
                        start_time=None,
                        send_action=None):
        """Move trajectory.

        This function call move_trajectory_sequence internally.

        x : float
            move distance with respect to x axis. unit is [m].
        y : float
            move distance with respect to y axis. unit is [m].
        yaw : float
            rotate angle. unit is [rad].
        sec : float
            time. unit is [sec].
        """
        return self.move_trajectory_sequence([[x, y, yaw]], [sec],
                                             stop=stop,
                                             start_time=start_time,
                                             send_action=send_action)

    def clear_costmap(self):
        """Send signal to clear costmap for obstacle avoidance to move_base.

        """
        service_name = "{}/clear_costmaps".format(self.move_base_action_name)
        rospy.wait_for_service(service_name)
        empty_service = rospy.ServiceProxy(service_name, std_srvs.srv.Empty)
        return empty_service()

    def change_inflation_range(self,
                               inflation_range=0.2,
                               node_name='move_base_node',
                               costmap_name='local_costmap',
                               inflation_name='inflation'):
        """Changes inflation range of local costmap for obstacle avoidance.

        Parameters
        ----------
        inflation_range : float
            range of inflation
        node_name : str
            name of move_base_node
        costmap_name : str
            name of costmap
        inflation_name : str
            name of inflation
        """
        service_name = "{}/{}/{}/set_parameters".format(
            node_name, costmap_name, inflation_name)
        request = dynamic_reconfigure.srv.ReconfigureRequest()
        request.config.doubles.append(
            dynamic_reconfigure.msg.DoubleParameter(
                name='inflation_radius',
                value=inflation_range))
        rospy.wait_for_service(service_name)
        return rospy.ServiceProxy(
            service_name,
            dynamic_reconfigure.srv.Reconfigure)()


""" Shared stuff between producer and consumer
 Author: Tobi Delbruck
 Source: https://github.com/SensorsINI/joker-network
 """
import logging
import math
import os
import sys
import time
from pathlib import Path
from subprocess import TimeoutExpired
import os
from typing import Optional

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0' # all TF messages

import tensorflow as tf

import numpy as np
import atexit
from engineering_notation import EngNumber  as eng  # only from pip
from matplotlib import pyplot as plt
import numpy as np
# https://stackoverflow.com/questions/35851281/python-finding-the-users-downloads-folder
import os
if os.name == 'nt':
    import ctypes
    from ctypes import windll, wintypes
    from uuid import UUID

    # ctypes GUID copied from MSDN sample code
    class GUID(ctypes.Structure):
        _fields_ = [
            ("Data1", wintypes.DWORD),
            ("Data2", wintypes.WORD),
            ("Data3", wintypes.WORD),
            ("Data4", wintypes.BYTE * 8)
        ]

        def __init__(self, uuidstr):
            uuid = UUID(uuidstr)
            ctypes.Structure.__init__(self)
            self.Data1, self.Data2, self.Data3, \
            self.Data4[0], self.Data4[1], rest = uuid.fields
            for i in range(2, 8):
                self.Data4[i] = rest>>(8-i-1)*8 & 0xff

    SHGetKnownFolderPath = windll.shell32.SHGetKnownFolderPath
    SHGetKnownFolderPath.argtypes = [
        ctypes.POINTER(GUID), wintypes.DWORD,
        wintypes.HANDLE, ctypes.POINTER(ctypes.c_wchar_p)
    ]

    def _get_known_folder_path(uuidstr):
        pathptr = ctypes.c_wchar_p()
        guid = GUID(uuidstr)
        if SHGetKnownFolderPath(ctypes.byref(guid), 0, 0, ctypes.byref(pathptr)):
            raise ctypes.WinError()
        return pathptr.value

    FOLDERID_Download = '{374DE290-123F-4565-9164-39C4925E467B}'

    def get_download_folder():
        return _get_known_folder_path(FOLDERID_Download)
else:
    def get_download_folder():
        home = os.path.expanduser("~")
        return os.path.join(home, "Downloads")

LOGGING_LEVEL = logging.INFO
PORT = 12000  # UDP port used to send frames from producer to consumer
IMSIZE = 224  # input image size, must match model
UDP_BUFFER_SIZE = int(math.pow(2, math.ceil(math.log(IMSIZE * IMSIZE + 1000) / math.log(2))))

EVENT_COUNT_PER_FRAME = 2300  # events per frame
EVENT_COUNT_CLIP_VALUE = 3  # full count value for colleting histograms of DVS events
SHOW_DVS_OUTPUT = True # producer shows the accumulated DVS frames as aid for focus and alignment
MIN_PRODUCER_FRAME_INTERVAL_MS=7.0 # inference takes about 3ms and normalization takes 1ms, hence at least 2ms
        # limit rate that we send frames to about what the GPU can manage for inference time
        # after we collect sufficient events, we don't bother to normalize and send them unless this time has
        # passed since last frame was sent. That way, we make sure not to flood the consumer
MAX_SHOWN_DVS_FRAME_RATE_HZ=15 # limits cv2 rendering of DVS frames to reduce loop latency for the producer
FINGER_OUT_TIME_S = 2  # time to hold out finger when joker is detected
ROOT_DATA_FOLDER= os.path.join(get_download_folder(),'trixsyDataset') # does not properly find the Downloads folder under Windows if not on same disk as Windows

DATA_FOLDER = os.path.join(ROOT_DATA_FOLDER,'data') #/home/tobi/Downloads/trixsyDataset/data' #'data'  # new samples stored here
NUM_NON_JOKER_IMAGES_TO_SAVE_PER_JOKER = 3 # when joker detected by consumer, this many random previous nonjoker frames are also saved
JOKERS_FOLDER = DATA_FOLDER + '/jokers'  # where samples are saved during runtime of consumer
NONJOKERS_FOLDER = DATA_FOLDER + '/nonjokers'
SERIAL_PORT = "/dev/ttyUSB0"  # port to talk to arduino finger controller

LOG_DIR='logs'
SRC_DATA_FOLDER = os.path.join(ROOT_DATA_FOLDER,'source_data') #'/home/tobi/Downloads/trixsyDataset/source_data'
TRAIN_DATA_FOLDER=os.path.join(ROOT_DATA_FOLDER,'training_dataset') #'/home/tobi/Downloads/trixsyDataset/training_dataset' # the actual training data that is produced by split from dataset_utils/make_train_valid_test()

MODEL_DIR='models' # where models stored
JOKER_NET_BASE_NAME='joker_net' # base name
USE_TFLITE = True  # set true to use TFLITE model, false to use full TF model for inference
TFLITE_FILE_NAME=JOKER_NET_BASE_NAME+'.tflite' # tflite model is stored in same folder as full-blown TF2 model
CLASS_DICT={'nonjoker':1, 'joker':2} # class1 and class2 for classifier
JOKER_DETECT_THRESHOLD_SCORE=.95 # minimum 'probability' threshold on joker output of CNN to trigger detection

import signal
def alarm_handler(signum, frame):
    raise TimeoutError
def input_with_timeout(prompt, timeout=30):
    """ get input with timeout
    :param prompt: the prompt to print
    :param timeout: timeout in seconds, or None to disable
    :returns: the input
    :raises: TimeoutError if times out
    """
    # set signal handler
    if timeout is not None:
        signal.signal(signal.SIGALRM, alarm_handler)
        signal.alarm(timeout) # produce SIGALRM in `timeout` seconds
    try:
        time.sleep(.5) # get input to be printed after logging
        return input(prompt)
    except TimeoutError as to:
        raise to
    finally:
        if timeout is not None:
            signal.alarm(0) # cancel alarm

def yes_or_no(question, default='y', timeout=None):
    """ Get y/n answer with default choice and optional timeout
    :param question: prompt
    :param default: the default choice, i.e. 'y' or 'n'
    :param timeout: the timeout in seconds, default is None
    :returns: True or False
    """
    if default is not None and (default!='y' and default!='n'):
        log.error(f'bad option for default: {default}')
        quit(1)
    y='Y' if default=='y' else 'y'
    n='N' if default=='n' else 'n'
    while "the answer is invalid":
        try:
            to_str='' if timeout is None or os.name=='nt' else f'(Timeout {default} in {timeout}s)'
            if os.name=='nt':
                log.warning('cannot use timeout signal on windows')
                time.sleep(.1) # make the warning come out first
                reply=str(input(f'{question} {to_str} ({y}/{n}): ')).lower().strip()
            else:
                reply = str(input_with_timeout(f'{question} {to_str} ({y}/{n}): ',timeout=timeout)).lower().strip()
        except TimeoutError:
            log.warning(f'timeout expired, returning default={default} answer')
            reply=''
        if len(reply)==0 or reply=='':
            return True if default=='y' else False
        elif reply[0].lower() == 'y':
            return True
        if reply[0].lower() == 'n':
            return False

class CustomFormatter(logging.Formatter):
    """Logging Formatter to add colors and count warning / errors"""

    grey = "\x1b[38;21m"
    yellow = "\x1b[33;21m"
    red = "\x1b[31;21m"
    bold_red = "\x1b[31;1m"
    reset = "\x1b[0m"
    format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"

    FORMATS = {
        logging.DEBUG: grey + format + reset,
        logging.INFO: grey + format + reset,
        logging.WARNING: yellow + format + reset,
        logging.ERROR: red + format + reset,
        logging.CRITICAL: bold_red + format + reset
    }

    def format(self, record):
        log_fmt = self.FORMATS.get(record.levelno)
        formatter = logging.Formatter(log_fmt)
        return formatter.format(record)

def my_logger(name):
    # logging.basicConfig(stream=sys.stdout, level=logging.INFO)
    logger = logging.getLogger(name)
    logger.setLevel(LOGGING_LEVEL)
    # create console handler
    ch = logging.StreamHandler()
    ch.setFormatter(CustomFormatter())
    logger.addHandler(ch)
    return logger

log=my_logger(__name__)

timers = {}
times = {}
class Timer:
    def __init__(self, timer_name='', delay=None, show_hist=False, numpy_file=None):
        """ Make a Timer() in a _with_ statement for a block of code.
        The timer is started when the block is entered and stopped when exited.
        The Timer _must_ be used in a with statement.
        :param timer_name: the str by which this timer is repeatedly called and which it is named when summary is printed on exit
        :param delay: set this to a value to simply accumulate this externally determined interval
        :param show_hist: whether to plot a histogram with pyplot
        :param numpy_file: optional numpy file path
        """
        self.timer_name = timer_name
        self.show_hist = show_hist
        self.numpy_file = numpy_file
        self.delay=delay

        if self.timer_name not in timers.keys():
            timers[self.timer_name] = self
        if self.timer_name not in times.keys():
            times[self.timer_name]=[]

    def __enter__(self):
        if self.delay is None:
            self.start = time.time()
        return self

    def __exit__(self, *args):
        if self.delay is None:
            self.end = time.time()
            self.interval = self.end - self.start  # measured in seconds
        else:
            self.interval=self.delay
        times[self.timer_name].append(self.interval)

    def print_timing_info(self, logger=None):
        """ Prints the timing information accumulated for this Timer
        :param logger: write to the supplied logger, otherwise use the built-in logger
        """
        if len(times)==0:
            log.error(f'Timer {self.timer_name} has no statistics; was it used without a "with" statement?')
            return
        a = np.array(times[self.timer_name])
        timing_mean = np.mean(a) # todo use built in print method for timer
        timing_std = np.std(a)
        timing_median = np.median(a)
        timing_min = np.min(a)
        timing_max = np.max(a)
        s='{} n={}: {}s +/- {}s (median {}s, min {}s max {}s)'.format(self.timer_name, len(a),
                                                                      eng(timing_mean), eng(timing_std),
                                                                      eng(timing_median), eng(timing_min),
                                                                      eng(timing_max))

        if logger is not None:
            logger.info(s)
        else:
            log.info(s)

def print_timing_info():
    for k,v in times.items():  # k is the name, v is the list of times
        a = np.array(v)
        timing_mean = np.mean(a)
        timing_std = np.std(a)
        timing_median = np.median(a)
        timing_min = np.min(a)
        timing_max = np.max(a)
        log.info('== Timing statistics from all Timer ==\n{} n={}: {}s +/- {}s (median {}s, min {}s max {}s)'.format(k, len(a),
                                                                          eng(timing_mean), eng(timing_std),
                                                                          eng(timing_median), eng(timing_min),
                                                                          eng(timing_max)))
        if timers[k].numpy_file is not None:
            try:
                log.info(f'saving timing data for {k} in numpy file {timers[k].numpy_file}')
                log.info('there are {} times'.format(len(a)))
                np.save(timers[k].numpy_file, a)
            except Exception as e:
                log.error(f'could not save numpy file {timers[k].numpy_file}; caught {e}')

        if timers[k].show_hist:

            def plot_loghist(x, bins):
                hist, bins = np.histogram(x, bins=bins) # histogram x linearly
                if len(bins)<2 or bins[0]<=0:
                    log.error(f'cannot plot histogram since bins={bins}')
                    return
                logbins = np.logspace(np.log10(bins[0]), np.log10(bins[-1]), len(bins)) # use resulting bin ends to get log bins
                plt.hist(x, bins=logbins) # now again histogram x, but with the log-spaced bins, and plot this histogram
                plt.xscale('log')

            dt = np.clip(a,1e-6, None)
            # logbins = np.logspace(np.log10(bins[0]), np.log10(bins[-1]), len(bins))
            try:
                plot_loghist(dt,bins=100)
                plt.xlabel('interval[ms]')
                plt.ylabel('frequency')
                plt.title(k)
                plt.show()
            except Exception as e:
                log.error(f'could not plot histogram: got {e}')

# this will print all the timer values upon termination of any program that imported this file
atexit.register(print_timing_info)


import sys

import numpy

from collections import Mapping

from bx.misc.binary_file import BinaryFileReader, BinaryFileWriter
from six.moves import reduce

def cdbhash( s ):
    return reduce( lambda h, c: (((h << 5) + h) ^ ord(c)) & 0xffffffff, s, 5381 )

class FileCDBDict( Mapping ):
    """
    For accessing a CDB structure on disk. Read only. Currently only supports
    access by key (getitem).

    NOTE: The keys method could be implemented by scanning the main table.
    """
    def __init__( self, file, is_little_endian=True ):
        # TODO: Deal with endianess
        self.io = BinaryFileReader( file, is_little_endian=is_little_endian )
        self.header_offset = self.io.tell()
        # Read the whole header (only 2k)
        self.header = [( self.io.read_uint32(), self.io.read_uint32() ) for _ in range( 256 )]

    def __getitem__( self, key ):
        hash = cdbhash( key )
        # Find position of subtable using 8 LSBs of hash
        subtable_offset = self.header[ hash % 256 ][0]
        subtable_size = self.header[ hash % 256 ][1]
        if subtable_size == 0:
            raise KeyError
        # Seek into subtable and look for match
        start = ( hash >> 8 )
        for i in range( subtable_size ):
            offset = subtable_offset + ( ( start + i ) % subtable_size ) * 8
            self.io.seek( offset )
            h = self.io.read_uint32()
            p = self.io.read_uint32()
            # Hit an empty bin, no match for key
            if p == 0:
                raise KeyError
            # Hash matches, need to check full key
            if h == hash:
                self.io.seek( p )
                klen = self.io.read_uint32()
                vlen = self.io.read_uint32()
                k = self.io.read( klen ).decode()
                if k == key:
                    v = self.io.read( vlen ).decode()
                    return v
        else:
            # Visited every slot and no match (should never happen since
            # there are empty slots by contruction)
            raise KeyError

    def __iter__(self):
        raise NotImplementedError()

    def __len__(self):
        raise NotImplementedError()

    @classmethod
    def to_file( Class, dict, file, is_little_endian=True ):
        """
        For constructing a CDB structure in a file. Able to calculate size on
        disk and write to a file
        """
        io = BinaryFileWriter( file, is_little_endian=is_little_endian )
        start_offset = io.tell()
        # Header is of fixed length
        io.seek( start_offset + ( 8 * 256 ) )
        # For each item, key and value length (written as length prefixed
        # strings). We also calculate the subtables on this pass.
        # NOTE: This requires the key and value be byte strings, support for
        #       dealing with encoding specific value types should be
        #       added to this wrapper
        subtables = [ [] for i in range(256) ]
        for key, value in dict.items():
            pair_offset = io.tell()
            io.write_uint32( len( key ) )
            io.write_uint32( len( value ) )
            io.write( key )
            io.write( value )
            hash = cdbhash( key )
            subtables[ hash % 256 ].append( ( hash, pair_offset ) )
        # Save the offset where the subtables will start
        subtable_offset = io.tell()
        # Write subtables
        for subtable in subtables:
            if len( subtable ) > 0:
                # Construct hashtable to be twice the size of the number
                # of items in the subtable, and built it in memory
                ncells = len( subtable ) * 2
                cells = [ (0,0) for i in range( ncells ) ]
                for hash, pair_offset in subtable:
                    index = ( hash >> 8 ) % ncells
                    while cells[index][1] != 0:
                        index = ( index + 1 ) % ncells
                    # Guaranteed to find a non-empty cell
                    cells[index] = ( hash, pair_offset )
                # Write subtable
                for hash, pair_offset in cells:
                    io.write_uint32( hash )
                    io.write_uint32( pair_offset )
        # Go back and write the header
        end_offset = io.tell()
        io.seek( start_offset )
        index = subtable_offset
        for subtable in subtables:
            io.write_uint32( index )
            io.write_uint32( len( subtable * 2 ) )
            # For each cell in the subtable, a hash and a pointer to a value
            index += ( len( subtable ) * 2 ) * 8
        # Leave fp at end of cdb
        io.seek( end_offset )


#!/usr/bin/env python
# coding: utf-8

from keras import backend as K
from keras.models import load_model
from keras.preprocessing import image
from keras.optimizers import Adam
from imageio import imread
import numpy as np
from matplotlib import pyplot as plt

from models.keras_ssd300 import ssd_300
from keras_loss_function.keras_ssd_loss import SSDLoss
from keras_layers.keras_layer_AnchorBoxes import AnchorBoxes
from keras_layers.keras_layer_DecodeDetections import DecodeDetections
from keras_layers.keras_layer_DecodeDetectionsFast import DecodeDetectionsFast
from keras_layers.keras_layer_L2Normalization import L2Normalization

from ssd_encoder_decoder.ssd_output_decoder import decode_detections, decode_detections_fast

from data_generator.object_detection_2d_data_generator import DataGenerator
from data_generator.object_detection_2d_photometric_ops import ConvertTo3Channels
from data_generator.object_detection_2d_geometric_ops import Resize
from data_generator.object_detection_2d_misc_utils import apply_inverse_transforms
import os

def getfiles(path):
    filelist = []
    for root, dirs, files in os.walk(path):
        filelist.extend([ os.path.join(root, file) for file in files if file.endswith(('.jpg', '.png', '.JPG', '.PNG'))])
    return filelist

# Set the image size.
img_height = 300
img_width = 300

# ## 1. Load a trained SSD
# TODO: Set the path to the `.h5` file of the model to be loaded.
model_path = 'ssd_detect.h5'

# We need to create an SSDLoss object in order to pass that to the model loader.
ssd_loss = SSDLoss(neg_pos_ratio=3, n_neg_min=0, alpha=1.0)

K.clear_session() # Clear previous models from memory.

model = load_model(model_path, custom_objects={'AnchorBoxes': AnchorBoxes,
                                               'L2Normalization': L2Normalization,
                                               'DecodeDetections': DecodeDetections,
                                               'compute_loss': ssd_loss.compute_loss})

# ## 2. Load some images
outdir = './output1'
if not os.path.exists(outdir):
    os.makedirs(outdir)
imglists = getfiles('/data/vocGT/val_images/')
for img_path in imglists:

    orig_images = imread(img_path)
    img = image.load_img(img_path, target_size=(img_height, img_width))
    img = image.img_to_array(img)
    input_images = [img]
    input_images = np.array(input_images)

    # ## 3. Make predictions
    y_pred = model.predict(input_images)

    confidence_threshold = 0.5

    y_pred_thresh = [y_pred[k][y_pred[k,:,1] > confidence_threshold] for k in range(y_pred.shape[0])]
    '''
    np.set_printoptions(precision=2, suppress=True, linewidth=90)
    print("Predicted boxes:\n")
    print('   class   conf xmin   ymin   xmax   ymax')
    print(y_pred_thresh[0])
    '''
    # ## 4. Visualize the predictions

    # Display the image and draw the predicted boxes onto it.

    # Set the colors for the bounding boxes
    colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
    classes = ['background', 'papers_z', 'papers_f', 'card_z', 'card_f', 'papers_hk_z', 'papers_hk_f']

    plt.figure()
    plt.imshow(orig_images)

    current_axis = plt.gca()

    for box in y_pred_thresh[0]:
        # Transform the predicted bounding boxes for the 300x300 image to the original image dimensions.
        xmin = box[2] * orig_images.shape[1] / img_width
        ymin = box[3] * orig_images.shape[0] / img_height
        xmax = box[4] * orig_images.shape[1] / img_width
        ymax = box[5] * orig_images.shape[0] / img_height
        color = colors[int(box[0])]
        label = '{}: {:.2f}'.format(classes[int(box[0])], box[1])
        current_axis.add_patch(plt.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, color=color, fill=False, linewidth=2))
        current_axis.text(xmin, ymin, label, size='x-large', color='white', bbox={'facecolor':color, 'alpha':1.0})

    plt.savefig(os.path.join(outdir, os.path.basename(img_path)))


import os
import cv2
import torch
import numpy as np

class EfficientDet(object):

    def __init__(self):
        # load config
        self.detector = None
        if torch.cuda.is_available():
            self.detector = torch.load('../object_detection/efficientdet/trained_models/signatrix_efficientdet_coco.pth')
            self.detector = self.detector.cuda()
        else:
            self.detector = torch.load('../object_detection/efficientdet/trained_models/signatrix_efficientdet_coco.pth',map_location='cpu').module

    def detect(self, img, conf=0.4):
        """

        """
        image_size = 1024
        im_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        # boxes, confidences, classIds =  detector.detect(im_rgb)
        height, width = im_rgb.shape[:2]
        image = im_rgb.astype(np.float32) / 255
        image[:, :, 0] = (image[:, :, 0] - 0.485) / 0.229
        image[:, :, 1] = (image[:, :, 1] - 0.456) / 0.224
        image[:, :, 2] = (image[:, :, 2] - 0.406) / 0.225
        if height > width:
            scale = image_size / height
            resized_height = image_size
            resized_width = int(width * scale)
        else:
            scale = image_size / width
            resized_height = int(height * scale)
            resized_width = image_size

        image = cv2.resize(image, (resized_width, resized_height))

        new_image = np.zeros((image_size, image_size, 3))
        new_image[0:resized_height, 0:resized_width] = image
        new_image = np.transpose(new_image, (2, 0, 1))
        new_image = new_image[None, :, :, :]
        new_image = torch.Tensor(new_image)
        if torch.cuda.is_available():
            new_image = new_image.cuda()
        with torch.no_grad():
            confidences, classIds, boxes = self.detector(new_image) # boxes are xmin ymin xmax ymax
            boxes /= scale
        boxes[:,2]=boxes[:,2]-boxes[:,0] # we need x y w h
        boxes[:,3]=boxes[:,3]-boxes[:,1]
        boxes=boxes[confidences>conf].cpu()
        classIds=classIds[confidences>conf].cpu()
        confidences=confidences[confidences>conf].cpu()

        return confidences, classIds, boxes


import cv2
import numpy as np
import glob
import os
import configparser

from auto_pose.ae import factory, utils

import tensorflow as tf

class AePoseEstimator:
    """ """

    # Takes a configPath only!
    def __init__(self, test_configpath):

        test_args = configparser.ConfigParser()
        test_args.read(test_configpath)

        workspace_path = os.environ.get('AE_WORKSPACE_PATH')

        if workspace_path == None:
            print 'Please define a workspace path:\n'
            print 'export AE_WORKSPACE_PATH=/path/to/workspace\n'
            exit(-1)

        self._camPose = test_args.getboolean('CAMERA','camPose')
        self._camK = np.array(eval(test_args.get('CAMERA','K_test'))).reshape(3,3)
        self._width = test_args.getint('CAMERA','width')
        self._height = test_args.getint('CAMERA','height')

        self._upright = test_args.getboolean('AAE','upright')
        self.all_experiments = eval(test_args.get('AAE','experiments'))

        self.class_names = eval(test_args.get('DETECTOR','class_names'))
        self.det_threshold = eval(test_args.get('DETECTOR','det_threshold'))
        self.icp = test_args.getboolean('ICP','icp')

        if self.icp:
            self._depth_scale = test_args.getfloat('DATA','depth_scale')

        self.all_codebooks = []
        self.all_train_args = []
        self.pad_factors = []
        self.patch_sizes = []

        config = tf.ConfigProto(allow_soft_placement=True)
        config.gpu_options.allow_growth=True
        config.gpu_options.per_process_gpu_memory_fraction = test_args.getfloat('MODEL','gpu_memory_fraction')

        self.sess = tf.Session(config=config)
        # set_session(self.sess)

        for i, experiment in enumerate(self.all_experiments):
            full_name = experiment.split('/')
            experiment_name = full_name.pop()
            experiment_group = full_name.pop() if len(full_name) > 0 else ''
            log_dir = utils.get_log_dir(workspace_path,experiment_name,experiment_group)
            ckpt_dir = utils.get_checkpoint_dir(log_dir)
            train_cfg_file_path = utils.get_train_config_exp_file_path(log_dir, experiment_name)
            print train_cfg_file_path
            # train_cfg_file_path = utils.get_config_file_path(workspace_path, experiment_name, experiment_group)
            train_args = configparser.ConfigParser()
            train_args.read(train_cfg_file_path)
            print(experiment)
            self.all_train_args.append(train_args)
            self.pad_factors.append(train_args.getfloat('Dataset','PAD_FACTOR'))
            self.patch_sizes.append((train_args.getint('Dataset','W'), train_args.getint('Dataset','H')))

            self.all_codebooks.append(factory.build_codebook_from_name(experiment_name, experiment_group, return_dataset=False))
            saver = tf.train.Saver(var_list=tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=experiment_name))
            factory.restore_checkpoint(self.sess, saver, ckpt_dir)

            # if self.icp:
            #     assert len(self.all_experiments) == 1, 'icp currently only works for one object'
            #     # currently works only for one object
            #     from auto_pose.icp import icp
            #     self.icp_handle = icp.ICP(train_args)
        if test_args.getboolean('ICP','icp'):
            from auto_pose.icp import icp
            self.icp_handle = icp.ICP(test_args, self.all_train_args)

    def extract_square_patch(self, scene_img, bb_xywh, pad_factor,resize=(128,128),interpolation=cv2.INTER_NEAREST,black_borders=False):

        x, y, w, h = np.array(bb_xywh).astype(np.int32)
        size = int(np.maximum(h, w) * pad_factor)

        left = np.maximum(x+w//2-size//2, 0)
        right = x+w//2+size/2
        top = np.maximum(y+h//2-size//2, 0)
        bottom = y+h//2+size//2

        scene_crop = scene_img[top:bottom, left:right].copy()

        if black_borders:
            scene_crop[:(y-top),:] = 0
            scene_crop[(y+h-top):,:] = 0
            scene_crop[:,:(x-left)] = 0
            scene_crop[:,(x+w-left):] = 0

        scene_crop = cv2.resize(scene_crop, resize, interpolation = interpolation)

        return scene_crop

    def process_detection(self, color_img):

        H, W = color_img.shape[:2]

        pre_image = preprocess_image(color_img)
        res_image, scale = resize_image(pre_image)

        batch_image = np.expand_dims(res_image, axis=0)
        print batch_image.shape
        print batch_image.dtype
        boxes, scores, labels = self.detector.predict_on_batch(batch_image)

        valid_dets = np.where(scores[0] >= self.det_threshold)

        boxes /= scale

        scores = scores[0][valid_dets]
        boxes = boxes[0][valid_dets]
        labels = labels[0][valid_dets]

        filtered_boxes = []
        filtered_scores = []
        filtered_labels = []

        for box,score,label in zip(boxes, scores, labels):

            box[0] = np.minimum(np.maximum(box[0],0),W)
            box[1] = np.minimum(np.maximum(box[1],0),H)
            box[2] = np.minimum(np.maximum(box[2],0),W)
            box[3] = np.minimum(np.maximum(box[3],0),H)

            bb_xywh = np.array([box[0],box[1],box[2]-box[0],box[3]-box[1]])
            if bb_xywh[2] < 0 or bb_xywh[3] < 0:
                continue

            filtered_boxes.append(bb_xywh)
            filtered_scores.append(score)
            filtered_labels.append(label)
        return (filtered_boxes, filtered_scores, filtered_labels)

    def process_pose(self, filtered_boxes, filtered_labels, color_img, depth_img=None, camPose=None):

        H, W = color_img.shape[:2]

        all_pose_estimates = []
        all_class_idcs = []

        for j,(box_xywh,label) in enumerate(zip(filtered_boxes,filtered_labels)):
            H_est = np.eye(4)

            try:
                clas_idx = self.class_names.index(label)
            except:
                print('%s not contained in config class_names %s', (label, self.class_names))
                continue

            det_img = self.extract_square_patch(color_img,
                                                box_xywh,
                                                self.pad_factors[clas_idx],
                                                resize=self.patch_sizes[clas_idx],
                                                interpolation=cv2.INTER_LINEAR,
                                                black_borders=True)

            Rs_est, ts_est = self.all_codebooks[clas_idx].auto_pose6d(self.sess,
                                                                        det_img,
                                                                        box_xywh,
                                                                        self._camK,
                                                                        1,
                                                                        self.all_train_args[clas_idx],
                                                                        upright=self._upright)

            R_est = Rs_est.squeeze()
            t_est = ts_est.squeeze()

            if self.icp:
                assert H == depth_img.shape[0]

                depth_crop = self.extract_square_patch(depth_img,
                                                    box_xywh,
                                                    self.pad_factors[clas_idx],
                                                    resize=self.patch_sizes[clas_idx],
                                                    interpolation=cv2.INTER_NEAREST) * self._depth_scale
                R_est_auto = R_est.copy()
                t_est_auto = t_est.copy()

                R_est, t_est = self.icp_handle.icp_refinement(depth_crop, R_est, t_est, self._camK, (W,H), clas_idx=clas_idx, depth_only=True)
                _, ts_est, _ = self.all_codebooks[clas_idx].auto_pose6d(self.sess,
                                                                            det_img,
                                                                            box_xywh,
                                                                            self._camK,
                                                                            1,
                                                                            self.all_train_args[clas_idx],
                                                                            upright=self._upright,
                                                                            depth_pred=t_est[2])
                t_est = ts_est.squeeze()
                R_est, _ = self.icp_handle.icp_refinement(depth_crop, R_est, ts_est.squeeze(), self._camK, (W,H), clas_idx=clas_idx, no_depth=True)
                # depth_crop = self.extract_square_patch(depth_img,
                #                                     box_xywh,
                #                                     self.pad_factors[clas_idx],
                #                                     resize=self.patch_sizes[clas_idx],
                #                                     interpolation=cv2.INTER_NEAREST)
                # R_est, t_est = self.icp_handle.icp_refinement(depth_crop, R_est, t_est, self._camK, (W,H))

                H_est[:3,3] = t_est / self._depth_scale #mm / m
            else:
                H_est[:3,3] = t_est

            H_est[:3,:3] = R_est
            print 'translation from camera: ',  H_est[:3,3]

            if self._camPose:
                H_est = np.dot(camPose, H_est)

            all_pose_estimates.append(H_est)
            all_class_idcs.append(clas_idx)

        return (all_pose_estimates, all_class_idcs)