import io
import os
import time
import picamera
import picamera.array
from base_camera import BaseCamera

import time

import cv2
from time import sleep
import datetime

import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile

from distutils.version import StrictVersion
from collections import defaultdict
from io import StringIO

from utils import ops as utils_ops
from utils import label_map_util

# What model to download.
MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17'
MODEL_FILE = MODEL_NAME + '.tar.gz'
DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'

# Path to frozen detection graph. This is the actual model that is used for the object detection.
PATH_TO_FROZEN_GRAPH = MODEL_NAME + '/frozen_inference_graph.pb'

# List of the strings that is used to add correct label for each box.
PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt')

opener = urllib.request.URLopener()
opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)
tar_file = tarfile.open(MODEL_FILE)
for file in tar_file.getmembers():
  file_name = os.path.basename(file.name)
  if 'frozen_inference_graph.pb' in file_name:
    tar_file.extract(file, os.getcwd())

detection_graph = tf.Graph()
with detection_graph.as_default():
  od_graph_def = tf.compat.v1.GraphDef()
  with tf.io.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid:
    serialized_graph = fid.read()
    od_graph_def.ParseFromString(serialized_graph)
    tf.import_graph_def(od_graph_def, name='')

category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)

def load_image_into_numpy_array(image):
  (im_width, im_height) = image.size
  return np.array(image.getdata()).reshape(
      (im_height, im_width, 3)).astype(np.uint8)

def run_inference_for_single_image(tensor_dict, image, graph):
  # Run inference
  output_dict = sess.run(tensor_dict, feed_dict={image_tensor: image})

  # all outputs are float32 numpy arrays, so convert types as appropriate
  output_dict['num_detections'] = int(output_dict['num_detections'][0])
  output_dict['detection_classes'] = output_dict[
      'detection_classes'][0].astype(np.int64)
  output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
  output_dict['detection_scores'] = output_dict['detection_scores'][0]
  if 'detection_masks' in output_dict:
    output_dict['detection_masks'] = output_dict['detection_masks'][0]
  return output_dict

class Camera(BaseCamera):
    @staticmethod
    def frames():
        with picamera.PiCamera() as camera:
            #camera.vflip = True
            #camera.hflip = True
            resX = 768
            resY = 512
            image_np = np.empty((resY, resX, 3), dtype=np.uint8)
            camera.resolution = (resX, resY)

            with   detection_graph.as_default():
              with tf.compat.v1.Session() as sess:
                # Get handles to input and output tensors
                ops = tf.compat.v1.get_default_graph().get_operations()
                all_tensor_names = {output.name for op in ops for output in op.outputs}
                tensor_dict = {}
                for key in [
                    'num_detections', 'detection_boxes', 'detection_scores',
                    'detection_classes', 'detection_masks'
                ]:
                  tensor_name = key + ':0'
                  if tensor_name in all_tensor_names:
                    tensor_dict[key] = tf.compat.v1.get_default_graph().get_tensor_by_name(
                        tensor_name)
                if 'detection_masks' in tensor_dict:
                  # The following processing is only for single image
                  detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
                  detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
                  # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
                  real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
                  detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
                  detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
                  detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
                      detection_masks, detection_boxes, image.shape[1], image.shape[2])
                  detection_masks_reframed = tf.cast(
                      tf.greater(detection_masks_reframed, 0.5), tf.uint8)
                  # Follow the convention by adding back the batch dimension
                  tensor_dict['detection_masks'] = tf.expand_dims(
                      detection_masks_reframed, 0)
                image_tensor = tf.compat.v1.get_default_graph().get_tensor_by_name('image_tensor:0')

                # let camera warm up
                # time.sleep(2) camera startup duirng Tensorflow startup

                ''' Works, latency
                stream = io.BytesIO()
                for _ in camera.capture_continuous(stream, 'jpeg',use_video_port=True):
                    # return current frame
                    stream.seek(0)
                    yield stream.read()
                    # reset stream for next frame
                    stream.seek(0)
                    stream.truncate()
                '''

                ''' works, latency
                while True:
                    camera.capture(image_np, format='bgr')
                    yield cv2.imencode('.jpg', image_np)[1].tobytes()
                '''

                # 309 MS
                stream = io.BytesIO()
                for _ in camera.capture_continuous(stream, 'bgr',use_video_port=True):
                    # return current frame
                    stream.seek(0)

                    image_np = np.frombuffer(stream.read(), dtype=np.uint8, count=resX*resY*3).reshape((resY, resX,3))

                    # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
                    image_np_expanded = np.expand_dims(image_np, axis=0)
                    # Actual detection.
                    #output_dict = run_inference_for_single_image(tensor_dict, image_np_expanded, detection_graph)
                    output_dict = sess.run(tensor_dict, feed_dict={image_tensor: image_np_expanded})

                    # all outputs are float32 numpy arrays, so convert types as appropriate
                    output_dict['num_detections'] = int(output_dict['num_detections'][0])
                    output_dict['detection_classes'] = output_dict[
                        'detection_classes'][0].astype(np.int64)
                    output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
                    output_dict['detection_scores'] = output_dict['detection_scores'][0]
                    if 'detection_masks' in output_dict:
                        output_dict['detection_masks'] = output_dict['detection_masks'][0]

                    # Visualization of the results of a detection.
                    for i in range(output_dict['num_detections']):
                        if output_dict['detection_scores'][i] > 0.5:
                            pt1 = (int(output_dict['detection_boxes'][i][1]*image_np.shape[1]), int(output_dict['detection_boxes'][i][0]*image_np.shape[0]))
                            pt2 = (int(output_dict['detection_boxes'][i][3]*image_np.shape[1]), int(output_dict['detection_boxes'][i][2]*image_np.shape[0]))
                            cv2.rectangle(image_np,pt1,pt2,(0,255,0),3)
                            classTxt = '{} ({:2})'.format(category_index[output_dict['detection_classes'][i]]['name'], output_dict['detection_scores'][i])
                            cv2.putText(image_np,classTxt,pt1, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,0,0))

                    yield cv2.imencode('.jpg', image_np)[1].tobytes()
                    # reset stream for next frame
                    stream.seek(0)
                    stream.truncate()


"""

"""
import numpy as np
import tensorflow as tf

"""
Reads Darknet config and weights and creates Keras model with TF backend.
"""
import os.path, sys
sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
import argparse
import numpy as np
import tensorflow as tf
import cv2
from model.yolov4_tiny import Yolov4_tiny
import os
os.environ['CUDA_VISIBLE_DEVICES']='-1'
from model.nms import yolov4_nms
from model.nms import NonMaxSuppression
import random
import albumentations as A
from utils.preprocess import resize_img

def parse_args(args):
    parser = argparse.ArgumentParser(description='convert Darknet weight To tensorflow format.')
    parser.add_argument('--mode', default='test',help="convert weight or test weight")
    parser.add_argument('--model-type', default='tiny')
    parser.add_argument('--num-classes', default=80)
    parser.add_argument('--nms', default='diou_nms', help="choices=['hard_nms','diou_nms']")
    parser.add_argument('--nms-max-box-num', default=300)
    parser.add_argument('--nms-iou-threshold', default=0.2, type=float)
    parser.add_argument('--score-threshold', default=0.5, type=float)
    parser.add_argument('--scales-x-y', default=[2., 2., 2., 2., 2.])
    parser.add_argument('--darknet-weights-path', help='Path to Darknet weights file.',
                        default='pretrain/yolov4-tiny.weights')
    parser.add_argument('--output-weight-path', help='Path to output tf model file.',
                        default='pretrain/yolov4-tiny')
    parser.add_argument('--test-img-path', help='Path to a image to test a model.',
                        default="images/dog.jpg")

    #test model
    parser.add_argument('--tta', default=True)
    parser.add_argument('--img-size', default=(416, 416))
    parser.add_argument('--class-names', default='dataset/coco.names')
    return parser.parse_args(args)

def convert_weight(args):
    print("model is converting...")
    weights_file = open(args.darknet_weights_path, 'rb')
    major, minor, revision, seen, _ = np.fromfile(weights_file, dtype=np.int32, count=5)
    model = Yolov4_tiny(args, training=False)
    # model.summary()

    backbone_layers_list = []
    max_block_num = 5
    for i in range(1,max_block_num+1):
        j = 1
        while True:
            try:
                layer_name = 'block_{}_{}_batch_normalization'.format(i,j)
                model.get_layer(layer_name)
                backbone_layers_list.append(layer_name)
            except:
                pass
            try:
                layer_name = 'block_{}_{}_conv2d'.format(i,j)
                model.get_layer(layer_name)
                backbone_layers_list.append(layer_name)
            except:
                break
            j+=1

    head_layers_list=[]
    head_num = 2
    for i in range(1,head_num+1):
        j = 1
        while True:
            try:
                layer_name = 'yolov3_head_{}_{}_batch_normalization'.format(i,j)
                model.get_layer(layer_name)
                head_layers_list.append(layer_name)
            except:
                pass
            try:
                layer_name = 'yolov3_head_{}_{}_conv2d'.format(i,j)
                model.get_layer(layer_name)
                head_layers_list.append(layer_name)
            except:
                break
            j+=1

    model_layer_list = backbone_layers_list[0:-2]+head_layers_list[3:]+backbone_layers_list[-2:]+head_layers_list[0:3]
    i=0
    while i<len(model_layer_list):
        layer_name = model_layer_list[i]
        if 'batch_normalization' in layer_name:
            conv_name = layer_name.replace('batch_normalization', 'conv2d')
            filters = model.get_layer(conv_name).filters
            size = model.get_layer(conv_name).kernel_size[0]
            in_dim = model.get_layer(conv_name).get_input_shape_at(0)[-1]
            # darknet [beta, gamma, mean, variance]
            bn_weights = np.fromfile(
                weights_file, dtype=np.float32, count=4 * filters)
            # tf [gamma, beta, mean, variance]
            bn_weights = bn_weights.reshape((4, filters))[[1, 0, 2, 3]]

            conv_shape = (filters, in_dim, size, size)
            conv_weights = np.fromfile(
                weights_file, dtype=np.float32, count=np.product(conv_shape))
            # tf shape (height, width, in_dim, out_dim)
            conv_weights = conv_weights.reshape(
                conv_shape).transpose([2, 3, 1, 0])

            model.get_layer(conv_name).set_weights([conv_weights])
            model.get_layer(layer_name).set_weights(bn_weights)
            i+=2

        elif 'conv2d' in layer_name:
            filters = model.get_layer(layer_name).filters
            size = model.get_layer(layer_name).kernel_size[0]
            in_dim = model.get_layer(layer_name).get_input_shape_at(0)[-1]

            conv_bias = np.fromfile(weights_file, dtype=np.float32, count=filters)
            conv_shape = (filters, in_dim, size, size)
            conv_weights = np.fromfile(
                weights_file, dtype=np.float32, count=np.product(conv_shape))
            # tf shape (height, width, in_dim, out_dim)
            conv_weights = conv_weights.reshape(
                conv_shape).transpose([2, 3, 1, 0])
            model.get_layer(layer_name).set_weights([conv_weights, conv_bias])
            i += 1

    assert len(weights_file.read()) == 0, 'failed to read all data'
    weights_file.close()
    model.save(args.output_weight_path)
    # model.save_weights('/home/wangem1/papers_coding/111111/coco_pretrain')
    print("the converting process is finished!")

tta_transform_1 = A.Compose([
    A.HorizontalFlip(p=1.),
    A.RandomBrightnessContrast(p=0.2),
])
tta_transform_2 = A.Compose([
    A.RandomBrightnessContrast(p=0.2),
])
def detect_batch_img(img,model,args):
    img = img / 255
    pre_nms_decoded_boxes, pre_nms__scores = model(img, training=False)
    pre_nms_decoded_boxes = pre_nms_decoded_boxes.numpy()
    pre_nms__scores = pre_nms__scores.numpy()
    boxes, scores, classes, valid_detections = yolov4_nms(args)(pre_nms_decoded_boxes, pre_nms__scores, args)
    return boxes, scores, classes, valid_detections
    # pre_nms_decoded_boxes,pre_nms__scores = model(img)
    # pre_nms_decoded_boxes = pre_nms_decoded_boxes.numpy()
    # pre_nms__scores = pre_nms__scores.numpy()
    # return pre_nms_decoded_boxes, pre_nms__scores
def tta_nms(boxes,scores,classes,valid_detections,args):
    all_boxes = []
    all_scores = []
    all_classes = []
    batch_index = 0
    valid_boxes = boxes[batch_index][0:valid_detections[batch_index]]
    valid_boxes[:, (0, 2)] = (1.-valid_boxes[:,(2,0)])
    all_boxes.append(valid_boxes)
    all_scores.append(scores[batch_index][0:valid_detections[batch_index]])
    all_classes.append(classes[batch_index][0:valid_detections[batch_index]])
    for batch_index in range(1,boxes.shape[0]):
        all_boxes.append(boxes[batch_index][0:valid_detections[batch_index]])
        all_scores.append(scores[batch_index][0:valid_detections[batch_index]])
        all_classes.append(classes[batch_index][0:valid_detections[batch_index]])
    all_boxes = np.concatenate(all_boxes,axis=0)
    all_scores = np.concatenate(all_scores, axis=0)
    all_classes = np.concatenate(all_classes, axis=0)
    all_boxes,all_scores,all_classes = np.array(all_boxes), np.array(all_scores), np.array(all_classes)
    boxes, scores, classes, valid_detections = NonMaxSuppression.diou_nms_np_tta(np.expand_dims(all_boxes,0),np.expand_dims(all_scores,0),np.expand_dims(all_classes,0),args)
    boxes, scores, classes, valid_detections = np.squeeze(boxes), np.squeeze(scores), np.squeeze(classes), np.squeeze(valid_detections)
    return boxes[:valid_detections], scores[:valid_detections], classes[:valid_detections]

def plot_one_box(img, box, color=None, label=None, line_thickness=None):
    # Plots one bounding box on image img
    tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1  # line/font thickness
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 7, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
        cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 7, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)

def plot_boxes(img,boxes,scores,classes,class_names,args):
    for i in range(len(boxes)):
            if scores[i] < args.score_threshold:
                continue
            x1y1 = (boxes[i][0:2] * img.shape[0:2][::-1]).astype(np.int)
            x2y2 = (boxes[i][2:4] * img.shape[0:2][::-1]).astype(np.int)
            plot_one_box(img,[x1y1[0],x1y1[1],x2y2[0],x2y2[1]],(255,0,255),label=str(class_names[classes[i]]) + "," + str("%0.2f" % scores[i]))
            # cv2.rectangle(img, tuple(x1y1), tuple(x2y2), (0, 255, 0), 2)
            # cv2.putText(img, str(class_names[classes[i]]) + " " + str("%0.2f" % scores[i]),
            #             (x1y1[0], max(x1y1[1] - 10, 0)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
def test_model(args):
    model = tf.keras.models.load_model(args.output_weight_path)
    with open(args.class_names) as f:
        class_names = f.read().splitlines()
    img = cv2.imread(args.test_img_path)
    img_ori, _, _ = resize_img(img, args.img_size)
    img_copy = img_ori.copy()
    aug_imgs = []
    if args.tta:
        aug_imgs.append(tta_transform_1(image=img_copy)['image'])
        aug_imgs.append(tta_transform_2(image=img_copy)['image'])
    aug_imgs.append(img_copy)
    batch_img = np.array(aug_imgs)
    boxes, scores, classes, valid_detections = detect_batch_img(batch_img, model, args)
    boxes, scores, classes = tta_nms(boxes, scores, classes, valid_detections, args)
    plot_boxes(img_copy, boxes, scores, classes, class_names, args)
    cv2.imshow("d", img_copy / 255)
    cv2.waitKey(0)
import sys
if __name__ == '__main__':
    args = parse_args(sys.argv[1:])
    if args.mode=='convert':
        convert_weight(args)
    else:
        test_model(args)


# -*- coding: utf-8 -*-
"""
Created on Wed Jan 11 11:30:11 2017

@author: PRABLANC P.
"""

import numpy as np
import scipy.signal as sp

class ProsodicModificationRealTime(object):
    """Modification of the pitch of the voice with or without preservation of the spectrale envelope

    Notes: resampling step is not handled properly leading to "clic" sounds. The scipy.resample()
        function works with a Fourier-domain method which makes it impossible to deal with boundaries
        issues. It would be necessary to incorporate in the frame to resample a few samples before and after
        in order to deal with boundarie issues. A time-domain-based method would be probably more appropriate
        such as scipy.resample_poly()
    """

    def __init__(self, fs=16000.0, frame_format='Float32'):
        self.fs = np.float(fs)
        self._pitch_rate = 1.0
        self.MIN_INTENSITY_VALUE = -40
        self.fs_norm = self.fs / 16000.0
        self._mid_buffer_size = int(256.0 * self.fs_norm)
        self._buffer = 0.001*np.random.rand(5 * self._mid_buffer_size, 1)
        self.win_length = int(128.0 * self.fs_norm)  # based on fs = 16 kHz
        win = np.hanning(self.win_length*2)
        self.win_G = np.reshape(win[0:self.win_length], [self.win_length, 1])
        self.win_D = np.reshape(win[self.win_length:], [self.win_length, 1])
        self.win_s = int(256*self.fs_norm)

        # initialization related to pitch change
        self.initialize(self._pitch_rate, frame_format)

    def initialize(self, pitch_rate, frame_format='Float32'):
        #==============================================================================
        # Pitch-shifting initialization
        #==============================================================================
        self.frame_format = frame_format
        # if pitch_rate < 1, the weights can't overflow the pointer
        self.length_weight = int(120 * self.fs_norm)  # length of the weight window
        diff = np.ceil(self.win_s * pitch_rate) - self.length_weight
        if diff < 0:
            self.length_weight = int(self.length_weight + diff)
        tmp = np.arange(-self.length_weight, self.length_weight+1)
        # weight used in frame intercorrelation
        self.weight = np.ones(self.length_weight * 2 + 1)
        self.weight[0:self.length_weight] = 1 - np.abs(tmp[0:self.length_weight]) / float(self.length_weight) * 0.5
        self.weight[self.length_weight+1:] = 1 - np.abs(tmp[self.length_weight+1:]) / float(self.length_weight) * 0.5
        self.weight = self.weight.reshape([self.length_weight * 2 + 1, 1])
        self.frame_length = 1 + int(np.ceil(self.win_s * pitch_rate)) + 2*self.length_weight + self.win_s - 1
        self.shift_mod = int(np.ceil(self.win_s*pitch_rate))
        self.zeros_frame_mod = np.zeros([self.shift_mod, 1])
        if pitch_rate < 1:
            self.buffer_mod = np.zeros([5 * self.win_s, 1])
            print(pitch_rate)
        else:
            self.buffer_mod = np.zeros([5 * self.shift_mod, 1])

        #==============================================================================
        # LPC initialization
        #==============================================================================
        self.alpha = 0.97
        self.n_lpc = 16
        self.w_hamming = np.hamming(self.shift_mod)

    def pitchshifting(self, new_frame):
        new_frame = new_frame.reshape([new_frame.size, 1])
        self._buffer = self.buffershift(self._buffer,new_frame)

        frame_mod_prev = \
            self.buffer_mod[(self.shift_mod + 1 - self.length_weight) : (1 + self.shift_mod + self.length_weight + self.win_s)]

        frame_s = np.reshape(self._buffer[0:self.frame_length], [self.frame_length, 1])
        intensity = 10 * np.log10(np.sum(frame_s**2) + 1e-30)

        if intensity > self.MIN_INTENSITY_VALUE:
            ic = sp.correlate(frame_mod_prev, frame_s[0:self.win_s])
            intercorr = ic[self.win_s-2:self.win_s-1+2*self.length_weight] * self.weight
            I = intercorr.argmax()
            k_pos = I-(self.length_weight+1)
        else:
            k_pos = 0
        # Not clean, indices overflow buffer_mod size.
        # It works because overflowed values are taken into account.
        self.buffer_mod[self.shift_mod+k_pos+self.win_length:self.shift_mod+k_pos+self.frame_length] = \
            frame_s[self.win_length:self.frame_length]
        self.buffer_mod[self.shift_mod+k_pos:self.shift_mod+k_pos+self.win_length] = \
            + self.buffer_mod[self.shift_mod+k_pos:self.shift_mod+k_pos+self.win_length]*self.win_D \
            + frame_s[0:self.win_length]*self.win_G

        # shift modified buffer
        self.buffer_mod = self.buffershift(self.buffer_mod,self.zeros_frame_mod)
#        frame_residue = lpcanalysis(self.buffer_mod[:self.shift_mod])
        # resample
        frame_output = sp.resample(self.buffer_mod[:self.shift_mod], new_frame.size)
        return frame_output

    def lpcanalysis(self, frame):
        raise NotImplementedError('Not yet implemented');
#        frame_preemphasis = sp.lfilter([1, -self.alpha], 1, frame.reshape([frame.size]))
#        frame_preemphasis_windowed = frame_preemphasis[self.n_lpc:] \
#            *self.w_hamming
#        R = sp.correlate(frame_preemphasis_windowed, frame_preemphasis_windowed)  # coefficients de corrélation
#        Ri = R[window_length-1:window_length+self.n_lpc]
#
#        lpc_durbin = self.durbin(Ri)  # calcul des ai
#        ai = lpc_durbin['a']
#        ai = ai.reshape([ai.size])
#        self._lpc_coeff[:, n] = ai
#        frame_filt = sp.lfilter(ai, 1.0, frame_preemphasis)  # filtrage du signal
#        frame_residue = frame_filt[self.n_lpc:]

    def buffershift(self, buffer_loc, frame):
        buffer_loc[:-frame.size] = buffer_loc[frame.size:]
        buffer_loc[-frame.size:] = frame.reshape([frame.size, 1])
        return buffer_loc

    def str2numpy(self, frame):
        return np.fromstring(frame,self.frame_format)

    def numpy2str(self, frame):
        frame = frame.astype(self.frame_format)
        return frame.tostring()

    def durbin(self,frame_corr):
        """Perform linear predictive coding with Levinson-Durbin recursion.

        Args:
            frame_corr (numpy array): First samples of auto-correlation frame.

        Returns:
            {
                'a': a, the filter coefficient of an auto-regressive model
                'k': k, the reflexion coefficient
                'En': En, the prediction error
            }
        """
        R0 = frame_corr[0]
        frame_corr = np.reshape(frame_corr/R0, [frame_corr.size, 1])
        p = frame_corr.size - 1
        k = np.zeros([p, 1])
        a = 1
        for n in range(0, p):
            a = np.append(a, 0.)
            a = a.reshape([a.size, 1])
            r = frame_corr[0:n+2]
            En = np.sum(r*a)
            Bn = np.sum(np.flipud(r)*a)
            ki = -Bn/En
            a = a + np.flipud(a)*ki
            k[n] = ki
        En = R0*np.sum(frame_corr*a)
        return {
            'a': a,
            'k': k,
            'En': En,
            }

    def lsf2poly(lsf):
        """Convert line spectral frequencies to prediction filter coefficients
        returns a vector a containing the prediction filter coefficients from a vector lsf of line spectral frequencies.

        """
        #   Reference: A.M. Kondoz, "Digital Speech: Coding for Low Bit Rate Communications
        #   Systems" John Wiley & Sons 1994 ,Chapter 4

        # Line spectral frequencies must be real.

        lsf = np.array(lsf)

        if max(lsf) > np.pi or min(lsf) < 0:
            raise ValueError('Line spectral frequencies must be between 0 and pi.')

        p = len(lsf) # model order

        # Form zeros using the LSFs and unit amplitudes
        z  = np.exp(1.j * lsf)

        # Separate the zeros to those belonging to P and Q
        rQ = z[0::2]
        rP = z[1::2]

        # Include the conjugates as well
        rQ = np.concatenate((rQ, rQ.conjugate()))
        rP = np.concatenate((rP, rP.conjugate()))

        # Form the polynomials P and Q, note that these should be real
        Q  = np.poly(rQ);
        P  = np.poly(rP);

        # Form the sum and difference filters by including known roots at z = 1 and
        # z = -1

        if p%2:
            # Odd order: z = +1 and z = -1 are roots of the difference filter, P1(z)
            P1 = sp.convolve(P, [1, 0, -1])
            Q1 = Q
        else:
            # Even order: z = -1 is a root of the sum filter, Q1(z) and z = 1 is a
            # root of the difference filter, P1(z)
            P1 = sp.convolve(P, [1, -1])
            Q1 = sp.convolve(Q, [1,  1])

        # Prediction polynomial is formed by averaging P1 and Q1

        a = .5 * (P1+Q1)
        return a[0:-1:1] # do not return last element

    def poly2lsf(a):
        """Prediction polynomial to line spectral frequencies.

        converts the prediction polynomial specified by A,
        into the corresponding line spectral frequencies, LSF.
        normalizes the prediction polynomial by A(1).

        """

        #Line spectral frequencies are not defined for complex polynomials.

        # Normalize the polynomial

        a = np.array(a)
        if a[0] != 1:
            a/=a[0]

        if max(np.abs(np.roots(a))) >= 1.0:
            raise ValueError('The polynomial must have all roots inside of the unit circle.');

        # Form the sum and difference filters

        p  = len(a)-1   # The leading one in the polynomial is not used
        a1 = np.concatenate((a, np.array([0])))
        a2 = a1[-1::-1]
        P1 = a1 - a2        # Difference filter
        Q1 = a1 + a2        # Sum Filter

        # If order is even, remove the known root at z = 1 for P1 and z = -1 for Q1
        # If odd, remove both the roots from P1

        if p%2: # Odd order
            P, r = sp.deconvolve(P1,[1, 0 ,-1])
            Q = Q1
        else:          # Even order
            P, r = sp.deconvolve(P1, [1, -1])
            Q, r = sp.deconvolve(Q1, [1,  1])

        rP  = np.roots(P)
        rQ  = np.roots(Q)

        aP  = np.angle(rP[1::2])
        aQ  = np.angle(rQ[1::2])

        lsf = sorted(np.concatenate((-aP,-aQ)))

        return lsf

    @property
    def mid_buffer_size(self):
        return self._mid_buffer_size

    @property
    def buffer(self):
        return self._buffer

    @property
    def pitch_rate(self):
        print("pitch_rate = {}".format(self._pitch_rate))
        return self._pitch_rate

    @pitch_rate.setter
    def pitch_rate(self, new_pitch_rate):
        self._pitch_rate = new_pitch_rate
        self.initialize(self._pitch_rate, self.frame_format)


import rospy
import numpy as np
import copy
import tf
import datetime
import gc
from tools import *
from pprint import pprint
from pyquaternion import Quaternion
from gpd.msg import GraspConfigList
from moveit_python import *
from moveit_msgs.msg import Grasp, PlaceLocation
from geometry_msgs.msg import PoseStamped, Vector3, Pose
from trajectory_msgs.msg import JointTrajectoryPoint
from visualization_msgs.msg import MarkerArray
from std_msgs.msg import Header, ColorRGBA
from moveit_python.geometry import rotate_pose_msg_by_euler_angles, translate_pose_msg
from tf.transformations import *
import geometry_msgs.msg #for pose2 simple
import math
import sys
import moveit_commander
import moveit_msgs.msg
import geometry_msgs.msg
from math import pi
from std_msgs.msg import String
from moveit_commander.conversions import pose_to_list
from tf.msg import tfMessage
import time
from send_gripper_robotiq import gripper_client
from send_gripper_robotiq import gripper_client_2
from tf import TransformListener
import copy
from control_msgs.msg import *
from trajectory_msgs.msg import *
from sensor_msgs.msg import JointState
import sensor_msgs
import actionlib
from filter_pointcloud_client import call_pointcloud_filter_service
from moveit_commander import MoveGroupCommander, RobotCommander
from moveit_msgs.msg import Constraints, OrientationConstraint, PositionConstraint, JointConstraint
from copy import deepcopy
from pointcloud_operations import create_mesh_and_save
from sensor_msgs import point_cloud2
from show_pose_marker import place_marker_at_pose
from std_srvs.srv import Empty
client = None
from moveit_msgs.msg import MoveItErrorCodes

# Build a useful mapping from MoveIt error codes to error names
moveit_error_dict = {}
for name in MoveItErrorCodes.__dict__.keys():
    if not name[:1] == '_':
        code = MoveItErrorCodes.__dict__[name]
        moveit_error_dict[code] = name

class GpdPickPlace(object):
    grasps = []
    grasps_cartesian = []
    mark_pose = False
    grasp_offset = -0.08
    grasp_offset_cartesian = -0.25
    finger_indexes = None
    con_joints_indexes = None
    joint1_con = 0
    joint2_con = 0
    joint3_con = 0
    gripper_closed = False
    global objects_grasped_not_placed
    grasps_received = False

    def __init__(self, mark_pose=False):
        self.grasp_subscriber = rospy.Subscriber("/detect_grasps/clustered_grasps", GraspConfigList, self.grasp_callback)
        if mark_pose:
            self.mark_pose = True
            self.marker_publisher = rospy.Publisher('visualization_marker_array', MarkerArray, queue_size=1)
        self.p = PickPlaceInterface(group="manipulator", ee_group="endeffector", verbose=True, ns="/summit_xl/")
        self.tf = tf.TransformListener()

    def grasp_callback(self, msg):
        self.grasps = msg.grasps
        self.grasps_received = True
        if (len(msg.grasps)==0):
            pevent("No grasps found, aborting!")
        else:
            pevent("Received new grasps")

    def show_grasp_pose(self, publisher, grasp_pose):
        place_marker_at_pose(publisher, grasp_pose)

    def get_gpd_grasps(self):
        pevent("Waiting for grasps to arrive")
        while (self.grasps_received == False):
            rospy.sleep(0.01)
        return self.grasps

    def generate_grasp_msgs(self, selected_grasps):
            self.grasps = []
            formatted_grasps = []
            self.grasps_cartesian = []
            formatted_grasps_cartesian = []
            tot_grasps = len(selected_grasps)
            cont = 0
            filtered_orientation = 0
            for i in range(0, len(selected_grasps)):
                z_axis_unit = (0, 0, -1)
                ap_axis = (selected_grasps[i].approach.x, selected_grasps[i].approach.y, selected_grasps[i].approach.z)
                angle = numpy.dot(z_axis_unit, ap_axis)
                if (angle >= 0):
                    # filter it out, because grasp coming from below the ground
                    filtered_orientation += 1
                    print(repr(filtered_orientation) + " Grasp filtered because coming from underneath the ground")
                    continue
                tf_listener_.waitForTransform('/arm_camera_depth_optical_frame', '/summit_xl_base_footprint',
                                              rospy.Time(), rospy.Duration(2.0))
		g = Grasp()
                g.id = "dupa"
                gp = PoseStamped()
                gp.header.frame_id = "arm_camera_depth_optical_frame"
                org_q = self.trans_matrix_to_quaternion(selected_grasps[i])

                quat = org_q

                gp.pose.position.x = selected_grasps[i].surface.x + self.grasp_offset * selected_grasps[i].approach.x
                gp.pose.position.y = selected_grasps[i].surface.y + self.grasp_offset * selected_grasps[i].approach.y
                gp.pose.position.z = selected_grasps[i].surface.z + self.grasp_offset * selected_grasps[i].approach.z
                gp.pose.orientation.x = float(quat.elements[1])
                gp.pose.orientation.y = float(quat.elements[2])
                gp.pose.orientation.z = float(quat.elements[3])
                gp.pose.orientation.w = float(quat.elements[0])

		translated_pose = tf_listener_.transformPose("summit_xl_base_footprint", gp)

                g.grasp_pose = translated_pose
                g.pre_grasp_approach.direction.header.frame_id = "arm_ee_link"
                g.pre_grasp_approach.direction.vector.x = 1.0
                g.pre_grasp_approach.min_distance = 0.06
                g.pre_grasp_approach.desired_distance = 0.1
               	#g.allowed_touch_objects = ["<octomap>", "obj"]
                g.allowed_touch_objects = ["obj"]
             	#g.max_contact_force = 0.0
                g.grasp_quality = selected_grasps[0].score.data
                formatted_grasps.append(g)

                # Add config for cartesian pick
                gp_cartesian = PoseStamped()
                gp_cartesian.header.frame_id = "arm_camera_depth_optical_frame"
                gp_cartesian.pose.position.x = selected_grasps[i].surface.x + self.grasp_offset_cartesian * selected_grasps[i].approach.x
                gp_cartesian.pose.position.y = selected_grasps[i].surface.y + self.grasp_offset_cartesian * selected_grasps[i].approach.y
                gp_cartesian.pose.position.z = selected_grasps[i].surface.z + self.grasp_offset_cartesian * selected_grasps[i].approach.z
                gp_cartesian.pose.orientation.x = float(quat.elements[1])
                gp_cartesian.pose.orientation.y = float(quat.elements[2])
                gp_cartesian.pose.orientation.z = float(quat.elements[3])
                gp_cartesian.pose.orientation.w = float(quat.elements[0])

		translated_pose = tf_listener_.transformPose("summit_xl_base_footprint", gp_cartesian)

                g_cartesian = Grasp ()
                g_cartesian.id = "cart"
                g_cartesian.grasp_pose = translated_pose
                g_cartesian.allowed_touch_objects = ["obj"]
                formatted_grasps_cartesian.append(g_cartesian)

            # sort grasps using z (get higher grasps first)
            formatted_grasps.sort(key=lambda grasp: grasp.grasp_pose.pose.position.z, reverse=True)
            formatted_grasps_cartesian.sort(key=lambda grasp: grasp.grasp_pose.pose.position.z, reverse=True)
            return formatted_grasps, formatted_grasps_cartesian

    def trans_matrix_to_quaternion(self, grasp):
        r = np.array([[grasp.approach.x, grasp.binormal.x, grasp.axis.x],
                      [grasp.approach.y, grasp.binormal.y, grasp.axis.y],
                      [grasp.approach.z, grasp.binormal.z, grasp.axis.z]])
        return Quaternion(matrix=r)

    def pick(self, grasps_list, verbose=False):
        failed_grasps = 0
        pevent("Pick sequence started")

	# Add object mesh to planning scene
        self.add_object_mesh()

	group.set_goal_tolerance(0.02)
        for single_grasp in grasps_list:
            if self.mark_pose:
                self.show_grasp_pose(self.marker_publisher, single_grasp.grasp_pose)
                rospy.sleep(1)
            pevent("Planning grasp:")
            pprint(single_grasp.grasp_pose)
            group.set_start_state_to_current_state()
            group.detach_object("obj")

            ### start code using pick interface ###
            pick_result = group.pick("obj", single_grasp)
            pevent("Planner returned: " + get_moveit_error_code(pick_result))
            if pick_result == 1:
              pevent("Grasp successful!")
              attach_link = "robotiq_85_base_link"
              touch_links = ["robotiq_base_link","robotiq_85_left_finger_link","robotiq_85_left_finger_tip_link","robotiq_85_right_finger_link","robotiq_85_right_finger_tip_link","robotiq_85_left_knuckle_link", "robotiq_85_left_inner_knuckle_link","robotiq_85_right_knuckle_link", "robotiq_85_right_inner_knuckle_link"]
              group.attach_object("obj", attach_link, touch_links)
              group.stop()
              group.clear_pose_targets()
              return single_grasp
            else:
              failed_grasps += 1
              group.stop()
              group.clear_pose_targets()
            ### end code using pick interface ###

            ### start code NOT using pick interface ###
         #   group.set_pose_target(single_grasp.grasp_pose.pose)
         #   plan = group.plan()
         #   if (len(plan.joint_trajectory.points) != 0):
         #      inp = raw_input("Have a look at the planned motion. Do you want to proceed? y/n: ")[0]
         #      if (inp == 'y'):
         #          pevent("Executing grasp: ")
         #          pick_result = group.execute(plan, wait=True)
         #          if pick_result == True:
         #              pevent("Grasp successful!")
         #              attach_link = "robotiq_85_base_link"
         #              touch_links = ["robotiq_base_link","robotiq_85_left_finger_link","robotiq_85_left_finger_tip_link","robotiq_85_right_finger_link","robotiq_85_right_finger_tip_link","robotiq_85_left_knuckle_link", "robotiq_85_left_inner_knuckle_link","robotiq_85_right_knuckle_link", "robotiq_85_right_inner_knuckle_link"]
         #              group.attach_object("obj", attach_link, touch_links)
         #              return single_grasp
         #          else:
         #              failed_grasps += 1
         #          group.stop()
         #          group.clear_pose_targets()
         #          group.clear_path_constraints()
         #      elif (inp == 'exit'):
         #          group.stop()
         #          group.clear_pose_targets()
         #          group.clear_path_constraints()
         #          exit(1)
        ### end code NOT using pick interface ###

        self.grasps = []

    def pick_cartesian(self, grasps_list, grasps_list_cartesian, verbose=False):
        failed_grasps = 0
        pevent("Cartesian pick sequence started")

	# Add object mesh to planning scene
        self.add_object_mesh()

	group.set_goal_tolerance(0.01)
        cont_c = 0
        for single_grasp in grasps_list_cartesian:
            if self.mark_pose:
                self.show_grasp_pose(self.marker_publisher, single_grasp.grasp_pose)
                rospy.sleep(1)
            pevent("Planning grasp:")
            pprint(single_grasp.grasp_pose)
            group.set_start_state_to_current_state()
            group.detach_object("obj")
            group.set_pose_target(single_grasp.grasp_pose.pose)
            plan = group.plan()
            if (len(plan.joint_trajectory.points) != 0):
                inp = raw_input("Have a look at the planned motion. Do you want to proceed? y/n: ")
                if (inp == 'y'):
                    pevent("Executing grasp: ")
                    pick_result = group.execute(plan, wait=True)
                    if pick_result == True:
                        group.stop()
                        group.clear_pose_targets()
                        group.clear_path_constraints()
                        group.set_start_state_to_current_state()
                        if self.mark_pose:
                            self.show_grasp_pose(self.marker_publisher, grasps_list[cont_c].grasp_pose)
                            rospy.sleep(1)
                        group.set_goal_tolerance(0.01)
                        waypoints = []
                        wpose = grasps_list[cont_c].grasp_pose.pose
                        waypoints.append(copy.deepcopy(wpose))

                        # We want the Cartesian path to be interpolated at a resolution of 1 cm
                        # which is why we will specify 0.01 as the eef_step in Cartesian
                        # translation.  We will disable the jump threshold by setting it to 0.0 disabling:
                        (plan, fraction) = group.compute_cartesian_path(
                            waypoints,	# waypoints to follow
                            0.01,	# eef_step
                            0.0)	# jump_threshold
                        waypoints_result = group.execute(plan, wait=True)
                        if waypoints_result == True:
                            pevent("Grasp successful!")
                            attach_link = "robotiq_85_base_link"
                            touch_links = ["robotiq_base_link","robotiq_85_left_finger_link","robotiq_85_left_finger_tip_link","robotiq_85_right_finger_link","robotiq_85_right_finger_tip_link","robotiq_85_left_knuckle_link", "robotiq_85_left_inner_knuckle_link","robotiq_85_right_knuckle_link", "robotiq_85_right_inner_knuckle_link"]
                            group.attach_object("obj", attach_link, touch_links)
                            return single_grasp
                        else:
                            failed_grasps += 1
                            group.stop()
                            group.clear_pose_targets()
                            group.clear_path_constraints()
                    else:
                        failed_grasps += 1
                        group.stop()
                        group.clear_pose_targets()
                        group.clear_path_constraints()
                elif (inp == 'exit'):
                    group.stop()
                    group.clear_pose_targets()
                    group.clear_path_constraints()
                    exit(1)
            cont_c += 1
        self.grasps = []

    def pick_two_steps(self, grasps_list, grasps_list_cartesian, verbose=False):
        failed_grasps = 0
        pevent("Two step pick sequence started")

	# Add object mesh to planning scene
        self.add_object_mesh()

        group.set_goal_tolerance(0.01)
        group.set_planning_time(5)
        cont_c = 0
        for single_grasp in grasps_list_cartesian:
            if self.mark_pose:
                self.show_grasp_pose(self.marker_publisher, single_grasp.grasp_pose)
                rospy.sleep(1)
            pevent("Planning grasp:")
            pprint(single_grasp.grasp_pose)
            group.set_start_state_to_current_state()
            group.detach_object("obj")
            group.set_pose_target(single_grasp.grasp_pose.pose)
            plan = group.plan()
            if (len(plan.joint_trajectory.points) != 0):
                inp = raw_input("Have a look at the planned motion. Do you want to proceed? y/n: ")
                if (inp == 'y'):
                    pevent("Executing grasp: ")
                    pick_result = group.execute(plan, wait=True)
                    if pick_result == True:
                        group.stop()
                        group.clear_pose_targets()
                        group.clear_path_constraints()
                        group.set_start_state_to_current_state()
                        if self.mark_pose:
                            self.show_grasp_pose(self.marker_publisher, grasps_list[cont_c].grasp_pose)
                            rospy.sleep(1)
                        group.set_pose_target(grasps_list[cont_c].grasp_pose.pose)
                        plan2 = group.go()
                        if (plan2 == True):
                            attach_link = "robotiq_85_base_link"
                            touch_links = ["robotiq_base_link", "robotiq_85_left_finger_link",
                                           "robotiq_85_left_finger_tip_link", "robotiq_85_right_finger_link",
                                           "robotiq_85_right_finger_tip_link", "robotiq_85_left_knuckle_link",
                                           "robotiq_85_left_inner_knuckle_link", "robotiq_85_right_knuckle_link",
                                           "robotiq_85_right_inner_knuckle_link"]
                            group.attach_object("obj", attach_link, touch_links)
                            return single_grasp
                        else:
                            failed_grasps += 1
                            group.stop()
                            group.clear_pose_targets()
                            group.clear_path_constraints()
                    else:
                        failed_grasps += 1
                        group.stop()
                        group.clear_pose_targets()
                        group.clear_path_constraints()
                elif (inp == 'exit'):
                    group.stop()
                    group.clear_pose_targets()
                    group.clear_path_constraints()
                    exit(1)
            cont_c += 1
        self.grasps = []

    def place2(self, place_pose):
        #returns True or False if the place was successfull or not
        pevent("Place sequence started")
        pose_goal = geometry_msgs.msg.Pose()
        pose_goal.position.x = 0.516344249249
        pose_goal.position.y = -0.636391639709
        pose_goal.position.z =   0.603573918343
        pose_goal.orientation.w = 0.185884654522
        pose_goal.orientation.x = -0.681892871857
        pose_goal.orientation.y = 0.682668983936
        pose_goal.orientation.z = 0.185558646917
        group.set_start_state_to_current_state()
        group.set_pose_target(pose_goal)
        pevent("Planning pose:")
        pprint(pose_goal)
        group.set_pose_target(pose_goal)
        plan = group.plan()
        rospy.sleep(1)
        cont_plan_place=0
        place_successful=False
        while ((len(plan.joint_trajectory.points) == 0) and (cont_plan_place <10)):
            plan = group.plan()
            rospy.sleep(1)
            cont_plan_place+=1
        if (len(plan.joint_trajectory.points) != 0):
            inp = raw_input("Have a look at the planned motion. Do you want to proceed? y/n: ")[0]
            if (inp == 'y'):
                pevent("Executing place: ")
                pick_result = group.execute(plan, wait=True)
                if pick_result == True:
                    pevent("Pose successful!")
                    group.detach_object("obj")
                    place_successful=True
                else:
                    pevent("Pose failed!")
                # Calling `stop()` ensures that there is no residual movement
                group.stop()
                group.clear_pose_targets()
                place_successful = False
            elif (inp == 'exit'):
                group.stop()
                group.clear_pose_targets()
                place_successful = False
                exit(1)
        else:
            place_successful=False
            group.stop()
            group.clear_pose_targets()
        return place_successful

    def place(self, place_pose):
        pevent("Place sequence started")
        places = self.generate_place_poses(place_pose)
        place_result = self.p.place_with_retry("obj", places, support_name="<octomap>", planning_time=9001,
                                  goal_is_eef=True)
        #pevent("Planner returned: " + get_moveit_error_code(place_result.error_code.val))

    def generate_place_poses(self, initial_place_pose):
        places = list()
        l = PlaceLocation()
        l.id = "dupadupa"
        l.place_pose.header.frame_id = "summit_xl_base_footprint"
        q = Quaternion(initial_place_pose.grasp_pose.pose.orientation.w,
                        initial_place_pose.grasp_pose.pose.orientation.x,
                        initial_place_pose.grasp_pose.pose.orientation.y,
                        initial_place_pose.grasp_pose.pose.orientation.z)
	# Load successful grasp pose
        l.place_pose.pose.position = initial_place_pose.grasp_pose.pose.position
        l.place_pose.pose.orientation.w = q.elements[0]
        l.place_pose.pose.orientation.x = q.elements[1]
        l.place_pose.pose.orientation.y = q.elements[2]
        l.place_pose.pose.orientation.z = q.elements[3]

	# Move 20cm to the right
        l.place_pose.pose.position.y += 0.2

	# Fill rest of the msg with some data
        l.post_place_posture = initial_place_pose.grasp_posture
        l.post_place_retreat = initial_place_pose.post_grasp_retreat
        l.pre_place_approach = initial_place_pose.pre_grasp_approach
        places.append(copy.deepcopy(l))

	# Rotate place pose to generate more possible configurations for the planner
        m = 16  # Number of possible place poses
        for i in range(0, m - 1):
            l.place_pose.pose = rotate_pose_msg_by_euler_angles(l.place_pose.pose, 0, 0, 2 * math.pi / m)
            places.append(copy.deepcopy(l))
        return places

    def add_object_mesh(self):
        #rospy.sleep(2)
        obj_pose = PoseStamped()
        obj_pose.header.frame_id = "arm_camera_depth_optical_frame"
        obj_pose.pose.position.x = 0
        obj_pose.pose.position.y = 0
        obj_pose.pose.position.z = 0
        obj_pose.pose.orientation.x = 0
        obj_pose.pose.orientation.y = 0
        obj_pose.pose.orientation.z = 0
        obj_pose.pose.orientation.w = 1
        translated_pose = tf_listener_.transformPose("summit_xl_base_footprint", obj_pose)

	# Remove collision object from previous run
        planning.removeCollisionObject("obj")
        rospy.sleep(1)
        planning.addMesh("obj", translated_pose.pose, "object.stl")
        print("Collision object is:")
        rospy.sleep(3)
        pprint(planning.getKnownCollisionObjects())

    def set_pose_constraints(self,tol_joint1,tol_joint2,tol_joint3):
        angle_constraints = Constraints()
        joint_constraint = JointConstraint()
        angle_constraints.name = "angle"
        joint_constraint.position = self.joint1_con
        joint_constraint.tolerance_above = tol_joint1
        joint_constraint.tolerance_below = tol_joint1
        joint_constraint.weight = 0.1
        joint_constraint.joint_name = "arm_shoulder_pan_joint"
        angle_constraints.joint_constraints.append(joint_constraint)

        joint_constraint2 = JointConstraint()
        joint_constraint2.position = self.joint2_con
        joint_constraint2.tolerance_above = tol_joint2
        joint_constraint2.tolerance_below = tol_joint2
        joint_constraint2.weight = 0.2
        joint_constraint2.joint_name = "arm_wrist_2_joint"
        angle_constraints.joint_constraints.append(joint_constraint2)

        joint_constraint3 = JointConstraint()
        joint_constraint3.position = self.joint3_con
        joint_constraint3.tolerance_above = tol_joint3
        joint_constraint3.tolerance_below = tol_joint3
        joint_constraint3.weight = 0.3
        joint_constraint3.joint_name = "arm_wrist_1_joint"
        angle_constraints.joint_constraints.append(joint_constraint3)
        group.set_path_constraints(angle_constraints)

    def set_upright_constraints(self,pose):
        upright_constraints = Constraints()
        orientation_constraint = OrientationConstraint()
        upright_constraints.name = "upright"
        orientation_constraint.header = pose.header
        orientation_constraint.link_name = group.get_end_effector_link()
        orientation_constraint.orientation = pose.pose.orientation
        orientation_constraint.absolute_x_axis_tolerance = 3.14
        orientation_constraint.absolute_y_axis_tolerance = 0.1
        orientation_constraint.absolute_z_axis_tolerance = 0.1
        #orientation_constraint.absolute_z_axis_tolerance = 3.14 #ignore this axis
        orientation_constraint.weight = 1
        upright_constraints.orientation_constraints.append(orientation_constraint)
        group.set_path_constraints(upright_constraints)

    def start_con_setup(self):
        self.joints_subscriber = rospy.Subscriber('/summit_xl/joint_states', JointState, self.joints_callback)

    def stop_con_setup(self):
        self.joints_subscriber.unregister()

    def joints_callback(self, data):
        if (self.con_joints_indexes == None):
            names = data.name
            j1_index = names.index("arm_shoulder_pan_joint")
            j2_index = names.index("arm_wrist_2_joint")
            j3_index = names.index("arm_wrist_1_joint")
            self.con_joints_indexes = (j1_index, j2_index, j3_index)

        self.joint1_con = data.position[self.con_joints_indexes[0]]
        self.joint2_con = data.position[self.con_joints_indexes[1]]
        self.joint3_con = data.position[self.con_joints_indexes[2]]

    def remove_pose_constraints(self):
        group.clear_path_constraints()

    def drop_obj_on_robot(self, successful_grasp):
        pevent("Dropping object on robot")
        pose_goal = geometry_msgs.msg.Pose()
        pose_goal.position.x = -0.2
        pose_goal.position.y = 0
        pose_goal.position.z = 0.4+successful_grasp.grasp_pose.pose.position.z
        pose_goal.orientation.x = 0+successful_grasp.grasp_pose.pose.orientation.x
        pose_goal.orientation.y = 0+successful_grasp.grasp_pose.pose.orientation.y
        pose_goal.orientation.z = 0+successful_grasp.grasp_pose.pose.orientation.z
        pose_goal.orientation.w = successful_grasp.grasp_pose.pose.orientation.w
        group.set_start_state_to_current_state()
        group.set_goal_tolerance(0.05)
        group.set_pose_target(pose_goal)

        # The go command can be called with joint values, poses, or without any
        # parameters if you have already set the pose or joint target for the group
        #success = group.go(wait=True)
        # if (success==True):
        #     result = gripper_client_2(8)
        #     print("Gripper opened")
        #     group.detach_object("obj")
        # group.clear_pose_targets()
        # group.stop()
        # return success

        plan = group.plan()
        rospy.sleep(1)
        cont_plan_drop = 0
        while ((len(plan.joint_trajectory.points) == 0) and (cont_plan_drop < 10)):
            plan = group.plan()
            rospy.sleep(1)
            cont_plan_drop += 1
        if (len(plan.joint_trajectory.points) != 0):
            pevent("Executing dropping: ")
            result = group.execute(plan, wait=True)
            rospy.sleep(1)
            if result == True:
                pevent("Dropping successful!")
                gripper_client_2(0.1)
                print("Gripper opened")
                group.detach_object("obj")
                group.stop()
                group.clear_pose_targets()
                return True
            else:
                pevent("Dropping failed!")
                group.stop()
                group.clear_pose_targets()
                return False
        else:
            pevent("Dropping position planning failed. Aborting")
            return False

    def initial_pose(self):
        pevent("Initial constrained pose sequence started")
        pose_goal = geometry_msgs.msg.Pose()
        pose_goal.position.x = 1.069
        pose_goal.position.y = -0.012
        pose_goal.position.z = 0.831
        pose_goal.orientation.x = -0.534
        pose_goal.orientation.y = 0.445
        pose_goal.orientation.z = 0.478
        pose_goal.orientation.w = 0.536
        group.set_start_state_to_current_state()
        group.set_goal_tolerance(0.05)
        group.set_pose_target(pose_goal)
        group.set_planning_time(5)

        # The go command can be called with joint values, poses, or without any
        # parameters if you have already set the pose or joint target for the group
        # group.go(joint_goal, wait=True)
        plan = group.plan()
        rospy.sleep(1)
        cont_plan = 0
        while ((len(plan.joint_trajectory.points) == 0) and (cont_plan < 10)):
            plan = group.plan()
            rospy.sleep(1)
            cont_plan += 1
        group.clear_path_constraints()
        if (len(plan.joint_trajectory.points) != 0):
            pevent("Executing place: ")
            result = group.execute(plan, wait=True)
            rospy.sleep(1)
            if result == True:
                pevent("Initial pose successful!")
                group.stop()
                group.clear_pose_targets()
                return True
            else:
                pevent("Initial pose failed!")
                group.stop()
                group.clear_pose_targets()
                return False
        else:
            pevent("Initial position planning failed. Aborting")
            return False

    def wait_for_pcl_and_save(self):
        pinfo("Subscribing to pointcloud to generate pointcloud")
        self.obj_pc_subscriber = rospy.Subscriber("/cloud_indexed_pc_only", sensor_msgs.msg.PointCloud2,
                                                  self.obj_pointcloud_callback_pcd)

    def obj_pointcloud_callback_pcd(self, msg):
        pinfo("Pointcloud received")
        cloud = []
        for p in point_cloud2.read_points(msg, skip_nans=True):
            cloud.append([p[0], p[1], p[2]])
        create_pcd_and_save(cloud)
        pinfo("PCD generated")
        self.obj_pc_subscriber.unregister()

    def create_pcd_and_save(cloud):
        np_cloud = np.asarray(cloud)

    def wait_for_mesh_and_save(self):
      pinfo("Subscribing to pointcloud to generate mesh")
      self.obj_pc_subscriber = rospy.Subscriber("/cloud_indexed_pc_only", sensor_msgs.msg.PointCloud2 , self.obj_pointcloud_callback)

    def obj_pointcloud_callback(self, msg):
      pinfo("Pointcloud received")
      cloud = []
      for p in point_cloud2.read_points(msg, skip_nans=True):
                cloud.append([p[0], p[1], p[2]])
      create_mesh_and_save(cloud)
      pinfo("Mesh generated")
      self.obj_pc_subscriber.unregister()

    def gripper_callback(self, data):
        if (self.finger_indexes == None):
            names = data.name
            lf_index = names.index("gripper_left_finger_base_joint")
            rf_index = names.index("gripper_right_finger_base_joint")
            self.finger_indexes = (lf_index, rf_index)

        lf_joint = data.position[self.finger_indexes[0]]
        rf_joint = data.position[self.finger_indexes[1]]

        closed_range = 0.003

        if (lf_joint < closed_range and rf_joint < closed_range):
            if (not self.gripper_closed):
                pprint("Gripper closed, we probably lost the grip")
            self.gripper_closed = True

    def start_grasp_check(self):
        # Subscribe to topic to derive gripper position
        self.fingers_subscriber = rospy.Subscriber('/summit_xl/joint_states', JointState, self.gripper_callback)

    def stop_grasp_check(self):
        self.fingers_subscriber.unregister()
        return self.gripper_closed

if __name__ == "__main__":
    start_time = datetime.datetime.now()
    rospy.init_node("gpd_pick_and_place",anonymous=True)
    tf_listener_ = TransformListener()
    pnp = GpdPickPlace(mark_pose=True)
    group_name = "manipulator"
    group = moveit_commander.MoveGroupCommander(group_name, robot_description="/summit_xl/robot_description", ns="/summit_xl")

    group.set_planner_id("BiTRRT")
    #group.set_max_velocity_scaling_factor(0.05)
    #group.set_goal_orientation_tolerance(0.01)
    group.set_planning_time(5)
    #group.allow_replanning(True)
    planning = PlanningSceneInterface("summit_xl_base_footprint", ns="/summit_xl/")
    planning.clear()
    rospy.sleep(1)
    num_objects = 3
    succesfull_objects_placements = 0
    objects_grasped_lost = 0
    objects_grasped_not_placed = 0
    rospy.wait_for_service('/summit_xl/clear_octomap')
    clear_octomap = rospy.ServiceProxy('/summit_xl/clear_octomap', Empty)
    clear_octomap()

    #num_view = 1
    for r in range (0, num_objects):
        planning.clear()
        # Subscribe for grasps
        print("--- Move Arm to Initial Position---")
        pnp.remove_pose_constraints()
        pnp.start_con_setup()
        pnp.set_pose_constraints(1.57, 3.14, 3.14)
        pnp.stop_con_setup()
        while (pnp.initial_pose() == False):
            print("Initial arm positioning failed!")
        print("Initial arm positioning performed")
        clear_octomap()
        # We have to add a check, so that this is called only if the initial_pose was successful
        call_pointcloud_filter_service()
        pnp.wait_for_mesh_and_save()
    	# Wait for grasps from gpd, wrap them into Grasp msg format and start picking
        pnp.grasps_received = False
        selected_grasps = pnp.get_gpd_grasps()
        [formatted_grasps, formatted_grasps_cartesian] = pnp.generate_grasp_msgs(selected_grasps)
        gripper_client_2(0.1)
        print("Gripper opened")
        pnp.remove_pose_constraints()
        #pnp.start_con_setup()
        #pnp.set_pose_constraints(1.57, 1.57, 1.57)
        #pnp.stop_con_setup()
        #successful_grasp = pnp.pick(formatted_grasps, verbose=True)
        #successful_grasp = pnp.pick_cartesian(formatted_grasps, formatted_grasps_cartesian, verbose=True)
        successful_grasp = pnp.pick_two_steps(formatted_grasps, formatted_grasps_cartesian, verbose=True)
        if successful_grasp is not None:
            gripper_client_2(0.71)
            print("Gripper closed")
            #pnp.start_grasp_check()
            pnp.remove_pose_constraints()
            #pnp.start_con_setup()
            #rospy.sleep(1)
            #pnp.set_pose_constraints(3.14, 1.0, 1.0)
            #pnp.stop_con_setup()
     	    #pnp.set_upright_constraints(successful_grasp.grasp_pose)
            while (pnp.drop_obj_on_robot(successful_grasp) == False):
                print("Object placing failed!")
            #if success == False:
            #    objects_grasped_not_placed += 1

            #check_gripper_closed = pnp.stop_grasp_check()
            #if(check_gripper_closed == False):
            succesfull_objects_placements += 1
            #else:
            #	objects_grasped_lost += 1

        else:
            print("Grasp NOT performed")
            pnp.remove_pose_constraints()

        perc_successful_grasps = (100.0 * succesfull_objects_placements/num_objects)

    pnp.remove_pose_constraints()
    print(str(succesfull_objects_placements) + " out of " + str(num_objects) + " succesfull grasps, that is: %.2f" % perc_successful_grasps + "%. Of the grasped objects, we lost: " + str(objects_grasped_lost) + " and of the grasped once the not placed ones are " + str(objects_grasped_not_placed))
    pinfo("Demo runtime: " + str(datetime.datetime.now() - start_time))


from __future__ import (absolute_import, division,
                        print_function, unicode_literals)
import numpy as np
from scipy import stats
from .profiles import GaussProfile
from .profiles import UserProfile
from .profiles import DataProfile
from ..utils.utils import make_quant, shift_t

class Pulsar(object):
    """class for pulsars

    The minimal data to instantiate a pulsar is the period, Smean, and
    pulse profile. The Profile is supplied via a :class:`PulseProfile`-like
    object.

    Parameters
    ----------

    period : float
        Pulse period (sec)

    Smean : float
        Mean pulse flux density (Jy)

    profile : :class:`PulseProfile`
        Pulse profile or 2-D pulse portrait

    name : str
        Name of pulsar
    """
    #TODO Other data could be supplied via a `.par` file.
    def __init__(self, period, Smean, profiles=None, name=None):
        self._period = make_quant(period, 's')
        self._Smean = make_quant(Smean, 'Jy')

        self._name = name

        # Assign profile class; default to GaussProfile if nothing is specified
        if profiles is None:
            self._Profiles = GaussProfile()
        else:
            self._Profiles = profiles

    def __repr__(self):
        namestr = "" if self.name is None else self.name+", "
        return "Pulsar("+namestr+"{})".format(self.period.to('ms'))

    @property
    def Profiles(self):
        return self._Profiles

    @property
    def name(self):
        return self._name

    @property
    def period(self):
        return self._period

    @property
    def Smean(self):
        return self._Smean

    def make_pulses(self, signal, tobs):
        """generate pulses from Profiles, :class:`PulsePortrait` object

        Required Args:
            signal (:class:`Signal`-like): signal object to store pulses
            tobs (float): observation time (sec)
        """
        signal._tobs = make_quant(tobs, 's')

        # init base profile at correct sample rate
        Nph = int((signal.samprate * self.period).decompose())

        # If there isn't enough frequencies in the profiles
        # And if it is a :class:`Profile` instance, reshape.
        self.Profiles.init_profiles(Nph, signal.Nchan)

        # if (signal.Nchan != self.Profiles.profiles.shape[0]
        #     and hasattr(self.Profiles, 'set_Nchan')):
        #     self.Profiles.set_Nchan(signal.Nchan)

        # select pulse generation method
        if signal.sigtype in ["RFSignal", "BasebandSignal"]:
            self._make_amp_pulses(signal)
        elif signal.sigtype is "FilterBankSignal":
            self._make_pow_pulses(signal)
        else:
            msg = "no pulse method for signal: {}".format(signal.sigtype)
            raise NotImplementedError(msg)

        # compute Smax (needed for radiometer noise level)
        #pr = self.Profiles()
        pr = self.Profiles._max_profile
        nbins = len(pr) # Think this assumes a single profile for now...
        signal._Smax = self.Smean * nbins / np.sum(pr)

    def _make_amp_pulses(self, signal):
        """generate amplitude pulses

        This method should be used for radio frequency and basebanded
        pulses.

        Parameters
        ----------

        signal : :class:`Signal`-like
            Signal object to store pulses.
        """
        # generate several pulses in time
        distr = stats.norm()

        Nsamp = int((signal.tobs * signal.samprate).decompose())
        signal.init_data(Nsamp)

        # TODO break into blocks
        # TODO phase from .par file
        # calc profile at phases
        phs = (np.arange(Nsamp) /
                (signal.samprate * self.period).decompose().value)
        phs %= 1  # clip integer part

        # convert intensity profile to amplitude!
        full_prof = np.sqrt(self.Profiles.calc_profiles(phs))

        signal._data = full_prof * distr.rvs(size=signal.data.shape)

    def _make_pow_pulses(self, signal):
        """generate a power pulse

        This method should be used for filter bank pulses

        Parameters
        ----------

        signal : :class:`Signal`-like
            Signal object to store pulses.
        """
        if signal.fold:
            # Determine how many subints to make
            if signal.sublen is None:
                signal._sublen = signal.tobs
                signal._nsub = 1
            else:
                # This should be an integer, if not, will round
                signal._nsub = int(np.round(signal.tobs / signal.sublen))

            # determine the number of data samples necessary
            signal._nsamp = int((signal.nsub*(self.period*signal.samprate)).decompose())
            # Need to make an initial empty data array
            signal.init_data(signal.nsamp)

            # Tile the profiles to number of desired subints
            sngl_prof = np.tile(self.Profiles(), signal.nsub)

            # changed to number of subints
            signal._Nfold = (signal.sublen / self.period).decompose()
            distr = stats.chi2(df=signal.Nfold)
            signal._set_draw_norm(df=signal.Nfold)

            #Why is there a second call to init_data?
            signal.init_data(sngl_prof.shape[1])
            signal._data = (sngl_prof * distr.rvs(size=signal.data.shape)
                            * signal._draw_norm)
        else:
            # fold is false and will make single pulses
            signal._sublen = self.period
            # This should be an integer, if not, will round; may not be exact
            signal._nsub = int(np.round((signal.tobs / signal.sublen).decompose()))

            # generate several pulses in time
            distr = stats.chi2(df=1)
            signal._set_draw_norm(df=1)

            signal._nsamp = int((signal.tobs * signal.samprate).decompose())
            signal.init_data(signal.nsamp)

            # TODO break into blocks
            # TODO phase from .par file
            # calc profiles at phases
            phs = (np.arange(signal.nsamp) /
                   (signal.samprate * self.period).decompose().value)
            phs %= 1  # clip integer part
            full_prof = self.Profiles.calc_profiles(phs,signal.Nchan)

            signal._data = (full_prof * distr.rvs(size=signal.data.shape)
                            * signal._draw_norm)

    def null(self, signal, null_frac, length=None, frequency=None):
        """
        Function to simulate pulsar pulse nulling. Given some nulling fraction,
        will replace simulated pulses with noise until nulling fraction is met.
        This function should only be run after running any ism or other delays
        have been added, e.g. disperison, FD, profile evolution, etc., but
        should be run before adding the radiometer noise ('telescope.observe()`),
        if nulling is desired.

        Parameters
        ----------

        signal [class] : signal class containing the simulated pulses
        null_frac [float] : desired pulsar nulling fraction, given as a
            decimal; range of 0.0 to 1.0.
        length [float] : desired length of each null in seconds. If not given,
                         will randomly null pulses. Default is None.
        frequency [float] : frequency of pulse nulling, e.g. how often the pulsar
                            nulls per hour. E.g. if frequency is 2, then the
                            pulsar will null twice per hour for some length
                            of time. If not given, will randomly null pulses.
                            Default is None.
        """
        # Determine how many pulses need to be nulled out
        null_pulses = int(np.round(signal.nsub * null_frac)) # needs to be int, may not be exact
        # Get the number of bins per pulse
        Nph = int((signal.samprate * self.period).decompose())
        # determine the off pulse window
        opw = self.Profiles._calcOffpulseWindow(Nphase = Nph)
        # define the random noise distribution
        if signal.fold:
            distr = stats.chi2(df=signal.Nfold)
        else:
            distr = stats.chi2(df=1)
        # have a test ditribution to determine null bins if Nfold is 1
        if not signal.fold or signal.Nfold < 100:
            check_distr = stats.chi2(df=100)
        else:
            check_distr = stats.chi2(df=signal.Nfold)
        # Figure out how off-center the pulses are
        shift_val = Nph//2 - np.where(signal.data[0,:Nph]==np.max(signal.data[0,:Nph]))[0]
        # Now null randomly if no length or frequency is given
        if length==None and frequency==None:
            # randomly draw pulse numbers to null
            rand_pulses = np.random.choice(signal.nsub, null_pulses, replace=False)
            # replace each pulse number with random noise
            if signal.delay == None:
                for p in rand_pulses:
                    # convert to bins
                    null_bins = np.arange(Nph*p, Nph*(p+1)) + shift_val
                    # make sure we don't go beyond the data array length
                    null_bins = null_bins[null_bins<np.shape(signal.data)[1]]
                    # replace pulses with noise
                    if len(null_bins) != Nph:
                        noise = (distr.rvs(size=len(null_bins)) * signal._draw_norm)
                    else:
                        noise = (distr.rvs(size=Nph) * signal._draw_norm)
                    # if no extra delays have been added to the signal
                    signal._data[:,null_bins] = noise * np.mean(self.Profiles._max_profile[opw.astype(int)])
            # if if signal has been delayed (e.g dispersed, etc.)
            else:
                # First initialize new array
                null_array = np.zeros(np.shape(signal.data))
                for p in rand_pulses:
                    #null_bins = np.arange(Nph*p, Nph*(p+1))
                    null_bins = np.arange(Nph*p, Nph*(p+1)) + shift_val
                    # make sure we don't go beyond the data array length
                    null_bins = null_bins[null_bins<np.shape(signal.data)[1]]
                    # replace the appropriate arrays
                    if len(null_bins) != Nph:
                        null_array[:,null_bins] = (check_distr.rvs(size=len(null_bins)) * signal._draw_norm)
                    else:
                        null_array[:,null_bins] = (check_distr.rvs(size=Nph) * signal._draw_norm)
                # now shift the data
                freq_array = signal._dat_freq
                shift_dt = (1/signal._samprate).to('ms')
                for ii, freq in enumerate(freq_array):
                    null_array[ii,:] = shift_t(null_array[ii,:],
                                                 signal.delay[ii].value,
                                                 dt=shift_dt.value)
                # Now replace pulses with noise
                off_pulse_mean = np.mean(self.Profiles._max_profile[opw.astype(int)])
                noise_shape = np.shape(np.where(null_array>1))[1]
                noise = (distr.rvs(size=noise_shape) * signal._draw_norm)
                signal._data[np.where(null_array>1)] = noise*off_pulse_mean

        else:
            raise NotImplementedError("Length and Frequency not been implimented yet")


import cv2
import numpy as np
from keras.models import load_model

face_cascade = cv2.CascadeClassifier('./face_detector/haarcascade_frontalface_default.xml')
video_capture = cv2.VideoCapture(0)

labels_dict={0:'YES',1:'NO'}
color_dict={0:(0,255,0),1:(0,0,255)}
model=load_model('MaskNet.model')

while True:
    ret, frame = video_capture.read()
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    faces=face_cascade.detectMultiScale(gray,1.1,4)
    for (x, y, w, h) in faces:
        #cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
        face_img=frame[y:y+h,x:x+w]
        resize=cv2.resize(face_img,(224,224))
        resize=resize/255.0
        reshape=np.reshape(resize,(1,224,224,3))
        result=model.predict(reshape)
        #print(result)
        label=np.argmax(result[0])
        #print(label)
        cv2.rectangle(frame,(x,y),(x+w,y+h),color_dict[label],2)
        cv2.rectangle(frame,(x,y-40),(x+w,y),color_dict[label],-1)
        cv2.putText(frame,labels_dict[label],(x,y-10),cv2.FONT_HERSHEY_SIMPLEX,0.8,(225,225,225),2)

    cv2.imshow('Video', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

video_capture.release()
cv2.destroyAllWindows()


'''
This file describes the class Drift(), which takes raw MOSFIRE data and
structures it such that it can be used to track drift and measure the seeing
as a function of frame (i.e., time).
'''

__author__ = 'Taylor Hutchison'
__email__ = 'aibhleog@tamu.edu'
__version__ = 'Oct2019'

import numpy as np
import matplotlib.pyplot as plt
import astropy.io.fits as fits
from scipy.interpolate import interp1d
import matplotlib.gridspec as gridspec
import matplotlib.patheffects as PathEffects
from scipy.optimize import curve_fit
from datetime import datetime as dt
from astropy.stats import sigma_clip
import image_registration as ir # github.com/keflavich/image_registration
import collapse_profile as coll # written by TAH
import pandas as pd
import shutil
import os

import sys
if sys.version_info[0] < 3:
	raise Exception('Must be using Python 3+')

print('This class assumes you are based in the directory directly ' +\
		  'above the MOSFIRE data. If this is not true, add the beginning of ' +\
		  'your path to the kwarg "home" for each function.',end='\n\n')

class Drift:
	'''
	This class measures the drift of a given MOSFIRE dataset,
	provided there is a star in the mask. Check the input parameters
	to see what is needed to run this script successfully.
	'''

	# INPUT PARAMETERS
	date = ''       # string date for directories, ex. 2018nov25
	mask = ''       # mask name, ex. UDS_MOSFIRE_J
	dither = 0.0    # the dither in " of the mask (currently assumes ABAB)
	band = ''       # MOSFIRE band; choose from Y, J, H, K
	home = ''       # path for the MOSFIRE data

	# information about the star in the mask
	row_start = 0   # the yvalue for the bottom of the star's slit
	row_end = 0     # the yvalue for the top of the star's slit
	col_start = 0   # the xvalue for the first column (avoid skylines please)
	col_end = 0     # the xvalue for the last column (avoid skylines please)

	# LIST OF RAW FRAMES FOR CHOSEN MASK
	# -- full list of frames
	def mask_frames(self):
		'''
		Returns the list of MOSFIRE raw frames that are targeting the mask
		of choice. Will make distinction between calibrations and on-sky data.
		'''
		path = self.home+'%s/'%self.date
		#print(path)

		# converts the mask name into the format for the file names
		mfile = dt.strptime(self.date,'%Y%b%d').strftime('%y%m%d')
		#print('m'+mfile)

		allfiles = np.asarray(os.listdir(path))
		#allfiles = allfiles[]
		mfiles = [f[:7] for f in allfiles] # fast way to ignore other files
		raw_frames = allfiles[np.asarray(mfiles) == 'm'+mfile]
		raw_frames = np.sort(raw_frames)

		mask_frames = []
		for f in raw_frames:
			#print(f)
			head = fits.getheader(path+f)
			try:
			    if head['OBJECT'] == self.mask:
			        #print(f)
			        mask_frames.append(f)
			except KeyError: pass
		return mask_frames

	# -- splitting into both dithers
	def split_dither(self):
		'''
		Returns the list of MOSFIRE raw frames that are targeting the mask
		of choice, split into the two dithers.

		NOTE: rather than nod_A and nod_B, want a dictionary where the keys
		can be nod_A/B/C/D/E; read YOFFSET in as float and have the key be the
		string of the float with 2 decimal places.  If YOFFSET has a key that
		matches that value, append; otherwise, make new key.
		At the end, can read the number of things in the dictionary, and that
		represents the number of nods.
		'''
		path = self.home+'%s/'%self.date
		mask_frames = self.mask_frames()

		nod_A, nod_B = [],[]
		for filename in mask_frames:
			head = fits.getheader(path+filename)
			if head['GRATMODE'] == 'spectroscopy': # removes alignment frames
			    if head['YOFFSET'] == self.dither: nod_A.append(filename)
			    elif head['YOFFSET'] == -self.dither: nod_B.append(filename)
			    else: raise Exception('ABAB dither pattern not found.')
			else: pass

		print('Number of frames in nod A: %s, in nod B: %s'%(len(nod_A),len(nod_B)),end='\n\n')
		return nod_A, nod_B

	# FITTING MODEL TO STAR'S PROFILE
	def cut_out(self,filename):
		'''
		Creates the cutout for the star's 2D spectrum.

		NOTE: when the new code is implemented here (see collapse_profile.py), the 2D cutout
		returned will be used for the seeing map and the star drift tracker (like usual).
		HOWEVER, the "blank2D" things in the collapse_profile.py module will be used to run
		cross-correlations & return the shift of the mask.
		'''
		path = self.home+'%s/'%self.date

		star_img = fits.getdata(path+filename)
		profile_2D = star_img[self.row_start:self.row_end,self.col_start:self.col_end].copy()

		# CLIPPING OUT COSMIC RAYS
		# --> the threshold is sigma=2 because the code runs on the
		# --> rows of the 2D, where every value should be consistent
		# --> i.e., "spectrally"
		mask = sigma_clip(profile_2D,sigma=2,axis=1)
		profile_2D[mask.mask] = np.median(profile_2D)
		# replacing the cosmic ray pixels with the median

		return profile_2D

	def fit_model(self,filename):
		'''
		Creating the mask star's profile given a handful of columns to sum
		over (to increase the S/N), this function fits this profile.

		NOTE: when the new code is implemented here (see collapse_profile.py), this function
		will have the skylines masked out and the whole thing will be collapsed.  Then, we
		won't need to specify the columns.

		Add. NOTE: Collapse 2D spectrally, find star location, mask where star is; use that
		to collapse spatially, mask where skylines are.

		Need threshold param: if peak count is N times the line, then mask.

		For the fit, when using the entire slit, separate into chunks and collapse,
		getting the fits for those chunks; average the results together
		---> because raw spectrum curves

		INPUTS ---- filename:   str, name of raw MOSFIRE file to be read in
			        row_start:  int, the yvalue for the bottom of the star's slit
			        row_end:    int, the yvalue for the top of the star's slit
			        col_start:  int, the xvalue for the first column (no skylines please)
			        col_end:    int, the xvalue for the last column (no skylines please)

		RETURNS --- mean:       float, center pixel of profile
			        A:          float, value of counts at peak
			        sig:        float, standard deviation -- can get FWHM by ~2.35*sig

		'''
		path = self.home+'%s/'%self.date

		def gauss(xaxis, mean, A, sig, B):
			# simple gaussian fit
			return A * np.exp(-np.power(xaxis-mean, 2.)/(2*np.power(sig, 2.))) + B

		# -- reading in data for the star
		profile_2D = self.cut_out(filename)
		profile = np.sum(profile_2D,axis=1) # summing over a few columns to increase S/N

		# ---- fitting the profile of the star
		x = np.arange(len(profile))
		peak = profile.tolist().index(max(profile)) # index for spatial peak of emission
		#print(peak,end=',')
		popt, wavcov = curve_fit(gauss,x,profile,p0=[x[peak],profile[peak],4.,0.],\
					bounds=(0,[np.inf,np.inf,30,np.inf]))

		mean, A, sig, B = popt
		return mean, A, sig

	def cross_correlations(self,reference,filename):
		'''
		Takes filenames of reference frame & another frame, reads in data, & calculates
		the shift of the second frame from the reference using cross-correlations.

		INPUTS ---- reference:      str, name of raw MOSFIRE file to use as ref.
			        filename:       str, name of raw MOSFIRE file to be read in

		RETURNS --- xshift,yshift:  (float,float), shift of frame relative to reference
		'''
		path = self.home+'%s/'%self.date
		ref_frame = fits.getdata(path+reference)
		raw_frame = fits.getdata(path+filename)

		# masking out rows with signal in both
		# default sigma clipping: upper_sig=2.5, lower_sig=5
		ref_frame = coll.return_blank2D(ref_frame)
		raw_frame = coll.return_blank2D(raw_frame)

		# running the cross-correlation
		xshift,yshift = ir.cross_correlation_shifts(ref_frame,raw_frame)
		return xshift,yshift # shift from ref_frame to raw_frame
			                 # essentially tracks the drift of the slit

	def get_UTC(self,filename):
		'''
		Pulling UTC information for a given filename or multiple filenames.

		INPUTS ---- filename:   str, name of raw MOSFIRE file to be read in
		RETURNS --- utc_value:  string or array, UTC information for frame(s)
		'''
		path = self.home+'%s/'%self.date

		try: len(filename); one = False # more than one file
		except TypeError: one = True

		if one == True:
			head = fits.getheader(path+filename)
			utc = head['utc']
		else:
			utc = []
			for i in range(len(filename)):
			    head = fits.getheader(path+filename[i])
			    utc.append(head['utc'])

		return utc

	def get_airmass(self,filename):
		'''
		Pulling airmass information for a given filename or multiple filenames.

		INPUTS ---- filename:   str, name of raw MOSFIRE file to be read in
		RETURNS --- airmass:  	float or array, airmass information for frame(s)
		'''
		path = self.home+'%s/'%self.date

		try: len(filename); one = False # more than one file
		except TypeError: one = True

		if one == True:
			head = fits.getheader(path+filename)
			air = head['utc']
		else:
			air = []
			for i in range(len(filename)):
			    head = fits.getheader(path+filename[i])
			    air.append(head['airmass'])

		return air

	def get_pa_el(self,filename):
		'''
		Pulling PA & Elevation information for a given filename or multiple filenames.

		INPUTS ---- filename:   str, name of raw MOSFIRE file to be read in
		RETURNS --- pa:         float or array, rotpposn information for frame(s)
			        el:         float or array, elevation information for frame(s)
		'''
		path = self.home+'%s/'%self.date

		try: len(filename); one = False # more than one file
		except TypeError: one = True

		if one == True:
			head = fits.getheader(path+filename)
			pa,el = head['rotpposn'],head['el']
		else:
			pa,el = [],[]
			for i in range(len(filename)):
			    head = fits.getheader(path+filename[i])
			    pa.append(head['rotpposn'])
			    el.append(head['el'])

		return pa,el

	# CONVENIENCE FUNCTIONS
	# retrieving the fit for an entire dataset
	def fit_all(self,frames):
		'''
		Runs the fitting function on a large number of frames.
		Returns the fit parameters and the frame numbers.
		'''
		all_centers, all_As, all_sigs, frame_number = [],[],[],[]
		for filename in frames:
			#print(filename)
			mean, A, sig = self.fit_model(filename)
			all_centers.append(mean)
			all_As.append(A)
			all_sigs.append(sig)
			frame_number.append(int(filename[-8:-5]))

		return all_centers, all_As, all_sigs, frame_number

	# plotting all of the profiles for inspection
	def show_me_all_profiles(self,frames):
		'''
		Given frames, makes the cutouts, collapses spectrally,
		and plots all of the profiles for inspection.
		'''
		path = self.home+'%s/'%self.date

		# plotting the profiles
		plt.figure(figsize=(9,6))
		for filename in frames:
			# -- reading in data for the star
			profile_2D = self.cut_out(filename)
			profile = np.sum(profile_2D,axis=1)
			plt.plot(profile)

		plt.tight_layout()
		plt.show()
		plt.close('all')

import numpy,from scipy,import astropy,from astropy


import torch
from torch import nn

from .data import cfg_re50
from .models.retinaface import RetinaFace
from .layers.functions.prior_box import PriorBox
from .utils.nms.py_cpu_nms import py_cpu_nms
from .utils.box_utils import decode, decode_landm

import cv2
import numpy as np

class Retina(nn.Module):
    """Retina face detection module.

    Keyword Arguments:
        threshold -- confidence_threshold
        device {torch.device} -- The device on which to run neural net passes. Image tensors and
            models are copied to this device before running forward passes. (default: {None})
    """
    def __init__(
        self, threshold=0.9, device=None
    ):
        super().__init__()

        self.threshold = threshold

        self.cfg = cfg_re50
        self.net = RetinaFace(cfg=self.cfg, phase = 'test')
        net = self.load_model(self.net, './weights/Resnet50_Final.pth', False)

        self.device = torch.device('cpu')
        if device is not None:
            self.device = device
            self.to(device)

    def forward(self, img):
        """Run Retina face detection on a CV image.

        Arguments:
            img -- A CV image.

        Example:
        >>> import Retina
        >>> retina = Retina()
        >>> boxes = retina(img)
        """

        # Detect faces
        boxes = self.detect_faces(img)
        return boxes

    def check_keys(self, model, pretrained_state_dict):
        ckpt_keys = set(pretrained_state_dict.keys())
        model_keys = set(model.state_dict().keys())
        used_pretrained_keys = model_keys & ckpt_keys
        unused_pretrained_keys = ckpt_keys - model_keys
        missing_keys = model_keys - ckpt_keys
        print('Missing keys:{}'.format(len(missing_keys)))
        print('Unused checkpoint keys:{}'.format(len(unused_pretrained_keys)))
        print('Used keys:{}'.format(len(used_pretrained_keys)))
        assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
        return True

    def remove_prefix(self, state_dict, prefix):
        ''' Old style model is stored with all names of parameters sharing common prefix 'module.' '''
        print('remove prefix \'{}\''.format(prefix))
        f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
        return {f(key): value for key, value in state_dict.items()}

    def load_model(self, model, pretrained_path, load_to_cpu):
        print('Loading pretrained model from {}'.format(pretrained_path))
        if load_to_cpu:
            pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage)
        else:
            device = torch.cuda.current_device()
            pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))
        if "state_dict" in pretrained_dict.keys():
            pretrained_dict = self.remove_prefix(pretrained_dict['state_dict'], 'module.')
        else:
            pretrained_dict = self.remove_prefix(pretrained_dict, 'module.')
        self.check_keys(model, pretrained_dict)
        model.load_state_dict(pretrained_dict, strict=False)
        return model

    def detect_faces(self, img, origin_size = False, nms_threshold=0.4):
        img = np.float32(img)

        # testing scale
        target_size = 1600
        max_size = 1000#2150
        im_shape = img.shape
        im_size_min = np.min(im_shape[0:2])
        im_size_max = np.max(im_shape[0:2])
        resize = float(target_size) / float(im_size_min)
        # prevent bigger axis from being more than max_size:
        if np.round(resize * im_size_max) > max_size:
            resize = float(max_size) / float(im_size_max)
        if origin_size:
            resize = 1

        if resize != 1:
            img = cv2.resize(img, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)
        im_height, im_width, _ = img.shape
        scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
        img -= (104, 117, 123)
        img = img.transpose(2, 0, 1)
        img = torch.from_numpy(img).unsqueeze(0)
        img = img.to(self.device)
        scale = scale.to(self.device)

        loc, conf, landms = self.net(img)  # forward pass

        priorbox = PriorBox(self.cfg, image_size=(im_height, im_width))
        priors = priorbox.forward()
        priors = priors.to(self.device)
        prior_data = priors.data
        boxes = decode(loc.data.squeeze(0), prior_data, self.cfg['variance'])
        boxes = boxes * scale / resize
        boxes = boxes.cpu().numpy()
        scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
        landms = decode_landm(landms.data.squeeze(0), prior_data, self.cfg['variance'])
        scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2],
                                img.shape[3], img.shape[2], img.shape[3], img.shape[2],
                                img.shape[3], img.shape[2]])
        scale1 = scale1.to(self.device)
        landms = landms * scale1 / resize
        landms = landms.cpu().numpy()

        # ignore low scores
        inds = np.where(scores > self.threshold)[0]
        boxes = boxes[inds]
        landms = landms[inds]
        scores = scores[inds]

        # keep top-K before NMS
        order = scores.argsort()[::-1]
        # order = scores.argsort()[::-1][:args.top_k]
        boxes = boxes[order]
        landms = landms[order]
        scores = scores[order]

        # do NMS
        dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
        keep = py_cpu_nms(dets, nms_threshold)
        # keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
        dets = dets[keep, :]
        landms = landms[keep]

        # keep top-K faster NMS
        # dets = dets[:args.keep_top_k, :]
        # landms = landms[:args.keep_top_k, :]

        #dets = np.concatenate((dets, landms), axis=1)
        probs = dets[:,4]
        points = landms.reshape(-1,5,2)
        boxes = dets[:,:4]
        return boxes, probs, points