#
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# Utility functions for performing image inference

import tensorflow as tf
from PIL import Image
import numpy as np

import utils.model as model_utils # UFF conversion uttils

# This class is similar as TRTInference inference, but it manages Tensorflow
class TensorflowInference(object):
    def __init__(self, pb_model_path):
        self.detection_graph = tf.Graph()
        with self.detection_graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(pb_model_path, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')
        self.sess = tf.Session(graph=self.detection_graph)

    def infer(self, image_path):
        img_np = self._load_img(image_path)
        return self._run_tensorflow_graph(np.expand_dims(img_np, axis=0))

    def infer_batch(self, image_paths):
        img_np = self._load_imgs(image_paths)
        return self._run_tensorflow_graph(img_np)

    def _run_tensorflow_graph(self, image_input):
        ops = self.detection_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'
        ]:
            tensor_name = key + ':0'
            if tensor_name in all_tensor_names:
                tensor_dict[key] = self.detection_graph.get_tensor_by_name(
                    tensor_name)

        image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
        output_dict = self.sess.run(tensor_dict,
            feed_dict={image_tensor: image_input})

        # All outputs are float32 numpy arrays, so convert types as appropriate
        output_dict['num_detections'] = output_dict['num_detections'].astype(np.int32)
        output_dict['detection_classes'] = output_dict[
            'detection_classes'].astype(np.uint8)

        return output_dict

    def _load_image_into_numpy_array(self, image):
        (im_width, im_height) = image.size
        return np.array(image).reshape(
            (im_height, im_width, model_utils.ModelData.get_input_channels())
        ).astype(np.uint8)

    def _load_imgs(self, image_paths):
        numpy_array = np.zeros((len(image_paths),) + (300, 300, 3))
        for idx, image_path in enumerate(image_paths):
            img_np = self._load_img(image_path)
            numpy_array[idx] = img_np
        return numpy_array

    def _load_img(self, image_path):
        img = Image.open(image_path)
        img_np = self._load_image_into_numpy_array(img)
        return img_np


import torch
from tts_modules.synthesizer.configs.hparams import hparams
# from tts_modules.synthesizer.utils import audio
from tts_modules.synthesizer.utils.symbols import symbols
from tts_modules.synthesizer.utils.text import text_to_sequence
from tts_modules.synthesizer.models.tacotron import Tacotron
from pathlib import Path
from typing import Union, List
import numpy as np
import librosa

class SynthesizerManager:

    def __init__(self, configs, test_dataloader=None, train_dataloader=None,
                 model=None, checkpoint_path=None):
        self.configs = configs
        self.test_dataloader = test_dataloader
        self.train_dataloader = train_dataloader
        self.model = model
        self.checkpoint_path = checkpoint_path
        if torch.cuda.is_available():
            self.device = torch.device("cuda")
        else:
            self.device = torch.device("cpu")

        if self.model is None:
            self.__init_tacotron()

    def __call__(self, *args, **kwargs):
        return self.synthesize_spectrograms(*args, **kwargs)

    def synthesize_spectrograms(self, texts: List[str],
                                embeddings: Union[np.ndarray, List[np.ndarray]],
                                return_alignments=False,
                                do_save_spectrograms=True):
        """
        Synthesizes mel spectrograms from texts and speaker embeddings.
        :param texts: a list of N text prompts to be synthesized
        :param embeddings: a numpy array or list of speaker embeddings of shape (N, 256)
        :param return_alignments: if True, a matrix representing the alignments between the
        characters
        and each decoder output step will be returned for each spectrogram
        :return: a list of N melspectrograms as numpy arrays of shape (80, Mi), where Mi is the
        sequence length of spectrogram i, and possibly the alignments.
        """
        # Preprocess text inputs
        inputs = [text_to_sequence(text.strip(), hparams.tts_cleaner_names) for text in texts]
        if not isinstance(embeddings, list):
            embeddings = [embeddings]

        # Batch inputs
        batched_inputs = [inputs[i:i+hparams.synthesis_batch_size]
                             for i in range(0, len(inputs), hparams.synthesis_batch_size)]
        batched_embeds = [embeddings[i:i+hparams.synthesis_batch_size]
                             for i in range(0, len(embeddings), hparams.synthesis_batch_size)]

        specs = []
        for i, batch in enumerate(batched_inputs, 1):

            # Pad texts so they are all the same length
            text_lens = [len(text) for text in batch]
            max_text_len = max(text_lens)
            chars = [SynthesizerManager.pad1d(text, max_text_len) for text in batch]
            chars = np.stack(chars)

            # Stack speaker embeddings into 2D array for batch processing
            speaker_embeds = np.stack(batched_embeds[i-1])

            # Convert to tensor
            chars = torch.tensor(chars).long().to(self.device)
            speaker_embeddings = torch.tensor(speaker_embeds).float().to(self.device)

            # Inference
            _, mels, alignments = self.model.generate(chars, speaker_embeddings)
            mels = mels.detach().cpu().numpy()
            for m in mels:
                # Trim silence from end of each spectrogram
                while np.max(m[:, -1]) < hparams.tts_stop_threshold:
                    m = m[:, :-1]
                specs.append(m)

        if do_save_spectrograms:
            if return_alignments:
                self.save_spectrograms(specs, alignments)
            else:
                self.save_spectrograms(specs)

        return (specs, alignments) if return_alignments else specs

    def save_spectrograms(self, specs, alignments=None):
        pass

    @staticmethod
    def pad1d(x, max_len, pad_value=0):
        return np.pad(x, (0, max_len - len(x)), mode="constant", constant_values=pad_value)

    def __init_tacotron(self):
        self.model = Tacotron(embed_dims=hparams.tts_embed_dims,
                              num_chars=len(symbols),
                              encoder_dims=hparams.tts_encoder_dims,
                              decoder_dims=hparams.tts_decoder_dims,
                              n_mels=hparams.num_mels,
                              fft_bins=hparams.num_mels,
                              postnet_dims=hparams.tts_postnet_dims,
                              encoder_K=hparams.tts_encoder_K,
                              lstm_dims=hparams.tts_lstm_dims,
                              postnet_K=hparams.tts_postnet_K,
                              num_highways=hparams.tts_num_highways,
                              dropout=hparams.tts_dropout,
                              stop_threshold=hparams.tts_stop_threshold,
                              speaker_embedding_size=hparams.speaker_embedding_size
                              ).to(self.device)
        if self.checkpoint_path is not None:
            self.__load_model()
        self.model.eval()

    def __load_model(self):
        if self.checkpoint_path is not None:
            self.model.load(self.checkpoint_path)
        else:
            print('Synthesizer was not loaded!!!')

    def __save_checkpoint(self, path):
        pass


import cv2
import redis
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import sys
import struct
from model import Model
import time

parser = argparse.ArgumentParser()
parser.add_argument('-in','--in_path',help='Video path',type=str)
parser.add_argument('-out','--out_path',help='output video path',type=str)
parser.add_argument('-db','--database',help='redis database',type=str)

if __name__=='__main__':
#     start = time.time()
    args = parser.parse_args()
    in_path = args.in_path
    out_path = args.out_path
    redis_db = args.database

    cap = cv2.VideoCapture(in_path)
    ret,fr = cap.read()
    r = redis.Redis(db=redis_db)
    threshold = 0.23

    try:

        stored_vectors = r.get('vectors')
        mapping = r.hgetall('mapping')

    except(redis.ConnectionError):

        raise("Redis Server is not UP!")

    if stored_vectors is None:
        raise('Please Register First!')
#     print(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    h,w = struct.unpack('>II',stored_vectors[:8])           #getting the shape of the vectors as it is stored as bytes
    stored_vectors = np.frombuffer(stored_vectors[8:])
    stored_vectors = stored_vectors.reshape(h,w)

    frame_rate = cap.get(cv2.CAP_PROP_FPS)
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    height,width = fr.shape[1],fr.shape[0]
    writer = cv2.VideoWriter(f'{out_path}/output.avi', fourcc, frame_rate, (height,width))

    while ret:

        alligned_faces,faces,kps,boxes = Model.get_input(fr[:,:,::-1])

        if boxes is None:
            writer.write(cv2.resize(fr,(height,width)))
            ret,fr = cap.read()
            continue

        generated_vectors = Model.get_batch_feature(alligned_faces)

        sim  = cosine_similarity(stored_vectors,generated_vectors)

        max_sim_indices = sim.argmax(axis=0)

        matched_faces = np.where(sim[max_sim_indices,range(len(generated_vectors))]>threshold)[0]

        left_faces = [i for i in range(len(generated_vectors)) if i not in matched_faces]

        if len(matched_faces) != 0:

            matched_stored_faces = max_sim_indices[matched_faces]

            for stored_idx,gen_idx in zip(matched_stored_faces, matched_faces):
                box = boxes[gen_idx]

                cv2.rectangle(fr, (box[0], box[1]), (box[2], box[3]), (0,255,0), 2)
                name = mapping[f'{stored_idx}'.encode()].decode()

                cv2.putText(
                        fr,str(name),
                        (box[0],box[1]),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        1,
                        (0,255,0),
                        2
                           )
        for idx in left_faces:
            box = boxes[idx]
            cv2.rectangle(fr, (box[0], box[1]), (box[2], box[3]), (0,0,255), 2)

        fr = cv2.resize(fr,(height,width))

        writer.write(fr)
        ret,fr = cap.read()

    print('Done')
#         frame_id+=1
#     print(time.time()-start)


# May 2020
# Synthetic variant(s) of Epic Kitchens dataset a.k.a. Epic Splitchens or SynthEpic

import os, time
import pandas as pd
import numpy as np
from dataset_epic import *
from epic_data import *

class synthepic_action_pair(data.Dataset):
    '''
    Epic Kitchens variant that contains videos concatenating a sequence of two actions
    (action_A, gap, cut, gap, action_B), based on labeled segments.
    Uses frames generated by the synthepic_curate.ipynb Jupyter notebook.
    '''

    def __init__(self, mode='train', transform=None, seq_len=5, num_seq=8,
                 pred_step=3, downsample=6, drive='ssd',
                 sample_method='within', sample_offset=0,
                 exact_cuts=False):
        '''
        sample_method: within / match_cut
            - within = given a clip, sample uniformly randomly within.
            - match_cut = transition between subclips is precisely when we start predicting (t -> t+1).
        sample_offset: Number of video blocks to shift sequence sampling by (in case of match_cut only).
        '''
        self.mode = mode
        self.transform = transform
        self.seq_len = seq_len
        self.num_seq = num_seq
        self.pred_step = pred_step
        self.downsample = downsample
        self.seq_frames = seq_len * num_seq * downsample # number of frames in one complete sequence
        self.drive = drive
        self.sample_method = sample_method
        self.sample_offset = sample_offset
        self.exact_cuts = exact_cuts

        # Gather video info for train / val / test. Every row of video_info is:
        # (video_root, n_frames, cut_idx, verb_A, noun_A, verb_B, noun_B)
        if exact_cuts:
            path_prefix = '/proj/vondrick/shared/uncertainty/synthepic_exact_'
        else:
            path_prefix = '/proj/vondrick/shared/uncertainty/synthepic_'
        if mode == 'train':
            split = path_prefix + 'train_split_%s.csv' % self.drive
            video_info = pd.read_csv(split)
        elif mode == 'val':
            split = path_prefix + 'val_split_%s.csv' % self.drive
            video_info = pd.read_csv(split)
        else:
            raise ValueError('Unknown dataset mode: ' + mode)

        # Drop videos that are too short
        drop_idx = []
        for idx, row in video_info.iterrows():
            vlen = row['n_frames']
            if vlen - self.seq_frames <= 0:
                drop_idx.append(idx)
        self.video_info = video_info.drop(drop_idx, axis=0)
        print('SynthEPIC size:', len(self.video_info), 'sample method:', sample_method, sample_offset,
              'exact:', exact_cuts)

        self.last_random_state = 1754

    def index_sampler(self, n_frames, cut_idx):
        ''' Chooses and returns all frame indices for a single sequence. '''
        if n_frames - self.seq_frames <= 0:
            return None

        if self.sample_method == 'within':
            # Get uniformly random start frame number
            start_idx = np.random.randint(0, n_frames - self.seq_frames)
        elif self.sample_method == 'match_cut':
            start_idx = cut_idx - self.seq_len * (self.num_seq - self.pred_step) * self.downsample
            start_idx += self.sample_offset * self.seq_len * self.downsample

        # Verify clip range
        if start_idx < 0 or start_idx + self.seq_frames >= n_frames:
            return None

        # Get 1D list of frame numbers representing START of every block
        seq_idx = np.expand_dims(
            np.arange(self.num_seq), -1) * self.downsample * self.seq_len + start_idx

        # Get 2D list of frame numbers of ALL frames within every block
        seq_idx_block = seq_idx + \
            np.expand_dims(np.arange(self.seq_len), 0) * self.downsample

        # Return list of used frame indices
        return seq_idx_block

    def __getitem__(self, index):
        if self.mode == 'train':
            # For train, sample fully random sequences each time
            # Force randomness (this might help especially with num_workers > 0)
            random_state = index * 1830 + \
                int((time.time() * 123456) % 321654) + self.last_random_state
            np.random.seed(random_state)
            self.last_random_state = random_state % 654321
        elif self.mode in ['val', 'test']:
            # For val and test, set deterministic random_state
            random_state = index

        # Get video path and associated total number of frames
        row = self.video_info.sample(
            weights=self.video_info['n_frames'], random_state=random_state).iloc[0]
        vpath, n_frames, cut_idx = row['vpath'], row['n_frames'], row['n_frames_A']

        # Sample uniformly random sequence of frames within this video
        idx_block = self.index_sampler(n_frames, cut_idx)
        if idx_block is None:
            # When match_cut didn't work out, simply call this method again to preserve statistical properties
            return self.__getitem__(np.random.randint(self.__len__()))

        # Decode list of frame numbers
        # dimensions = (block number, frame within block)
        assert(idx_block.shape == (self.num_seq, self.seq_len))
        # most significant = block number
        idx_block = idx_block.reshape(self.num_seq * self.seq_len)

        # Read all frames as images
        seq = [pil_loader(os.path.join(vpath, 'frame_%010d.jpg' % (i + 1)))
               for i in idx_block]

        # Transform sequence of images (NOT necessarily the same for all frames!)
        t_seq = self.transform(seq)

        # Convert list into tensor
        try:
            (C, H, W) = t_seq[0].size()
            t_seq = torch.stack(t_seq, 0)
        except:
            (C, H, W) = t_seq[0][0].size()
            tmp = [torch.stack(i, 0) for i in t_seq]
            assert len(tmp) == 5
            t_seq = torch.stack(tmp, 1)

        # View as e.g. (8, 5, 3, 128, 128) and then transpose to (8, 3, 5, 128, 128)
        t_seq = t_seq.view(self.num_seq, self.seq_len, C, H, W).transpose(1, 2)

        # Return all useful information in a dictionary
        result = {'t_seq': t_seq, 'idx_block': idx_block, 'vpath': vpath,
                  'n_frames': n_frames, 'cut_idx': cut_idx,
                  'n_frames_A': row['n_frames_A'], 'n_frames_B': row['n_frames_B'],
                  'video_id_A': row['video_id_A'], 'video_id_B': row['video_id_B'],
                  'src_start_A': row['src_start_A'], 'src_end_A': row['src_end_A'],
                  'src_start_B': row['src_start_B'], 'src_end_B': row['src_end_B'],
                  'verb_A': row['verb_A'], 'noun_A': row['noun_A'],
                  'verb_A_class': row['verb_A_class'], 'noun_A_class': row['noun_A_class'],
                  'verb_B': row['verb_B'], 'noun_B': row['noun_B'],
                  'verb_B_class': row['verb_B_class'], 'noun_B_class': row['noun_B_class']}
        return result

    def __len__(self):
        return len(self.video_info)


import numpy as np
from .slit_selection_ui import SlitSelectionUI
from .rectangle_slit import RectanglePixelSlit, RectangleSkySlit

class SlitController:
    """
    Controller that stores MOSViz slits.
    """
    def __init__(self, mosviz_viewer=None):
        self.mosviz_viewer = mosviz_viewer

        self._slits = []

    @property
    def slits(self):
        return self._slits

    @property
    def has_slits(self):
        return len(self.slits) > 0

    @property
    def x_bounds(self):
        """
        Given all slits, return the lowermost
        and uppermost x pixel values.
        None if n/a.
        """
        if not self.has_slits:
            return
        x_max_list = []
        x_min_list = []
        for slit in self.slits:
            bounds = slit.x_bounds
            if bounds is None:
                continue
            xmin, xmax = bounds
            x_min_list.append(xmin)
            x_max_list.append(xmax)

        if len(x_max_list) > 0:
            return [min(x_min_list), max(x_max_list)]

    @property
    def y_bounds(self):
        """
        Given all slits, return the lowermost
        and uppermost y pixel values.
        None if n/a.
        """
        if not self.has_slits:
            return
        y_max_list = []
        y_min_list = []
        for slit in self.slits:
            bounds = slit.y_bounds
            if bounds is None:
                continue
            ymin, ymax = bounds
            y_min_list.append(ymin)
            y_max_list.append(ymax)

        if len(y_max_list) > 0:
            return [min(y_min_list), max(y_max_list)]

    def get_cutout_limit(self):
        """
        The limits of the cutout image viewer
        needs to match the y axis length of
        the slit to match the 2d spectrum viewer.
        This function calculates the x and y limits
        such that the viewer is displaying a y by y image.

        Returns
        -------
        plot_limits : list or None
            [x_min, x_max, y_min, y_max]
        """
        if not self.has_slits:
            return None

        # We want the image dimensions to be dy x dy
        y_bounds = self.y_bounds
        if y_bounds is None:
            return None
        y_min, y_max = y_bounds

        x_bounds = self.x_bounds
        if x_bounds is None:
            return None
        x_min, x_max = x_bounds

        # Calculate new x bounds (must be length of dy)
        dy = y_max - y_min
        x = (x_max + x_min) / 2.
        new_x_min, new_x_max = (x - dy / 2., x + dy / 2.)

        return [new_x_min, new_x_max, y_min, y_max]

    def add_slit(self, slit):
        self._slits.append(slit)

    def add_rectangle_sky_slit(self, wcs, ra, dec, width, length):
        """
        Slit constructed using WCS and coords information.
        Utilizes `regions` package.
        Parameters
        ----------
        wcs : `astropy.wcs.WCS`
        ra, dec : float
            Center (ra, dec) of slit in deg.
        width, length : float
            Angular width and length of slit in arcsec.
        """
        self.add_slit(RectangleSkySlit(wcs, ra, dec, width, length))

    def add_rectangle_pixel_slit(self, x, y, width, length):
        """
        Slit constructed using WCS and coords information.
        Utilizes `regions` package.
        Parameters
        ----------
        x, y : float
            Center (x, y) of slit in pix.
        width, length : float
            width and length of slit in pix.
        """
        self.add_slit(RectanglePixelSlit(x, y, width, length))

    def draw_slits(self, ax):
        for slit in self.slits:
            if slit.is_active:
                slit.remove()
            slit.draw(ax=ax)

    def clear_slits(self):
        """
        Reset slit controller and its variables.
        Remove the patch from the axes its drawn in.
        """
        for slit in self.slits:
            if slit.is_active:
                slit.remove()
            del slit
        self._slits = []

    def launch_slit_ui(self):
        """
        Launches UI for slit selection.
        """
        return SlitSelectionUI(self.mosviz_viewer, self.mosviz_viewer)


import cv2
import numpy as np
from keras.models import load_model
import glob, os

def rgb2gray(rgb):
    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
    return gray

def predict_images(model):
    num_of_imges = 0
    num_of_predicted_imges = 0
    num_of_angry = 0
    num_of_fear = 0
    num_of_happy = 0
    num_of_sad = 0
    num_of_surprised = 0
    num_of_neutral = 0
    newf = open("emotion_data_test.txt", "w+")

    # Dictionary for emotion recognition model output and emotions
    emotions = {0: 'Angry', 1: 'Fear', 2: 'Happy', 3: 'Sad', 4: 'Surprised', 5: 'Neutral'}

    olddir = os.getcwd()

    newdir = olddir + "\\face_age_add_age_test"
    os.chdir(newdir)

    for file in glob.glob("*.png"):
        image = rgb2gray(cv2.imread(file));
        test_image = cv2.resize(image, (48, 48))
        test_image = test_image.reshape([-1, 48, 48, 1])

        test_image = np.multiply(test_image, 1.0 / 255.0)
        probab = model.predict(test_image)[0] * 100
        # Finding label from probabilities
        # Class having highest probability considered output label
        label = np.argmax(probab)
        max_probab = probab[label]
        predicted_emotion = emotions[label]

        if max_probab >= 55:
            newf.write(str(file) + "-" + str(predicted_emotion))
            newf.write("\n")
            num_of_predicted_imges += 1
            if predicted_emotion == "Angry":
                num_of_angry += 1
            elif predicted_emotion == "Fear":
                num_of_fear += 1
            elif predicted_emotion == "Happy":
                num_of_happy += 1
            elif predicted_emotion == "Sad":
                num_of_sad += 1
            elif predicted_emotion == "Surprised":
                num_of_surprised += 1
            elif predicted_emotion == "Neutral":
                num_of_neutral += 1

        num_of_imges += 1

    newf.close()
    print("Total number of pictures: " + str(num_of_imges))
    print("Number of labeled pictures: " + str(num_of_predicted_imges))
    print("Distribution of emotions: ")
    print("Angry: " + str(num_of_angry))
    print("Fear: " + str(num_of_fear))
    print("Happy: " + str(num_of_happy))
    print("Sad: " + str(num_of_sad))
    print("Surprised: " + str(num_of_surprised))
    print("Neutral: " + str(num_of_neutral))

def main():
    # Creating objects for emotion detection
    emotion_model = load_model('./model/emotion_recognition.h5')
    predict_images(emotion_model)

if __name__ == '__main__':
    main()


import itertools
import glob
import json
import numpy as np
import os
import pickle
import random

from nltk import word_tokenize
import torch
from torch import nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
from tqdm import tqdm

from keras.preprocessing import sequence

"""
LOSS_FUNC = nn.BCEWithLogitsLoss()
IMDB_DIR = 'data/aclImdb'
LM_FILE = 'data/lm_scores/imdb_all.txt'
COUNTER_FITTED_FILE = 'data/counter-fitted-vectors.txt'
"""
class AdversarialModel_Snli(object):
    def __init__(self, model, tokenizer, maxlen):
        super(AdversarialModel_Snli, self).__init__()
        model.eval()
        self.model = model
        self.tokenizer = tokenizer
        self.maxlen = maxlen
    def query(self, x_p, x_h, device):
        x_p = self.tokenizer.texts_to_sequences([x_p])
        x_p_mask = [1 for i in range(len(x_p[0]))]+[0 for i in range(self.maxlen-len(x_p[0]))]
        x_p_mask = [x_p_mask]
        x_p_mask = np.array(x_p_mask)
        x_p_mask=torch.from_numpy(x_p_mask).to(device).to(torch.float)

        x_p = sequence.pad_sequences(x_p, maxlen=self.maxlen, padding='post', truncating='post')
        x_p=torch.from_numpy(x_p).to(device).to(torch.long)

        x_h = self.tokenizer.texts_to_sequences([x_h])
        x_h_mask = [1 for i in range(len(x_h[0]))]+[0 for i in range(self.maxlen-len(x_h[0]))]
        x_h_mask = [x_h_mask]
        x_h_mask = np.array(x_h_mask)
        x_h_mask=torch.from_numpy(x_h_mask).to(device).to(torch.float)

        x_h = sequence.pad_sequences(x_h, maxlen=self.maxlen, padding='post', truncating='post')
        x_h=torch.from_numpy(x_h).to(device).to(torch.long)
        #input_vector = np.zeros((1, self.maxlen))
        #input_vector = x

        logit = self.model(mode="text_to_logit",x_p=x_p, x_h=x_h, x_p_mask=x_p_mask, x_h_mask=x_h_mask).squeeze(0)

        return logit.detach().cpu().numpy()
        #p = F.softmax(logit).detach().cpu().numpy()
        #return p.squeeze()

class AdversarialModel(object):
    def __init__(self, model, tokenizer, maxlen, test_bert=False):
        super(AdversarialModel, self).__init__()
        model.eval()
        self.model = model
        self.tokenizer = tokenizer
        self.maxlen = maxlen
        self.test_bert = test_bert

    def tokenize(self, x):
        if self.test_bert:
            #print(len(x.split()))
            #print(x)
            token = self.tokenizer.encode_plus(x, None, add_special_tokens=True, max_length=self.maxlen, pad_to_max_length=True)

            return token
        else:
            token = self.tokenizer.texts_to_sequences([x])
            token = sequence.pad_sequences(token, maxlen=self.maxlen, padding='post', truncating='post')
            return token

    def query(self, x, device):

        #input_vector = self.tokenizer.texts_to_sequences([x])
        #input_vector = sequence.pad_sequences(input_vector, maxlen=self.maxlen, padding='post', truncating='post')

        if self.test_bert:
            token = self.tokenize(x)

            text = np.array([token['input_ids']])
            text = torch.tensor(text,dtype=torch.long).to(device)
            mask = np.array([token['attention_mask']])
            #print(mask.sum())
            mask = torch.tensor(mask,dtype=torch.long).to(device)
            token_type_ids = np.array([token["token_type_ids"]])
            token_type_ids = torch.tensor(token_type_ids,dtype=torch.long).to(device)

            logit = self.model(mode="text_to_logit", input=text, bert_mask=mask, bert_token_id=token_type_ids).squeeze(0)

        else:
            token = self.tokenize(x)
            input_vector=torch.from_numpy(token).to(device).to(torch.long)
            logit = self.model(mode="text_to_logit",input=input_vector).squeeze(0)

        p = F.softmax(logit).detach().cpu().numpy()
        return p

class Adversary(object):
    def __init__(self, attack_surface):
        self.attack_surface = attack_surface

    def run(self, model, dataset, device, opts=None):
        """Run adversary on a dataset.

        Args:
        model: a TextClassificationModel.
        dataset: a TextClassificationDataset.
        device: torch device.
        Returns: pair of
        - list of 0-1 adversarial loss of same length as |dataset|
        - list of list of adversarial examples (each is just a text string)
        """
        raise NotImplementedError

class GeneticAdversary_Snli(Adversary):
    def __init__(self, attack_surface, num_iters=20, pop_size=60, margin_goal=0):
        super(GeneticAdversary_Snli, self).__init__(attack_surface)
        self.num_iters = num_iters
        self.pop_size = pop_size
        self.margin_goal = margin_goal

    def get_margins(self, model_output, gold_labels):
        logits = model_output

        true_class_pred = logits[gold_labels]

        temp = logits.copy()
        temp[gold_labels]=-1e20
        highest_false_pred = temp.max()
        value_margin = true_class_pred - highest_false_pred
        return value_margin

    def perturb(self, x_p, hypo, choices, model, y, device):
        if all(len(c) == 1 for c in choices):
            value_margin = self.get_margins( model.query(x_p, ' '.join(hypo), device), y)
            return hypo, value_margin.item()
        good_idxs = [i for i, c in enumerate(choices) if len(c) > 1]
        idx = random.sample(good_idxs, 1)[0]
        x_h_list = [' '.join(hypo[:idx] + [w_new] + hypo[idx+1:])
                for w_new in choices[idx]]
        querry_list = [model.query(x_p, x_h, device) for x_h in x_h_list]
        best_replacement_idx = None
        worst_margin = float('inf')
        for idx_in_choices, logits in enumerate(querry_list):
            value_margin = self.get_margins(logits, y)
            if best_replacement_idx is None or value_margin.item() < worst_margin:
                best_replacement_idx = idx_in_choices
                worst_margin = value_margin.item()

        cur_words = list(hypo)
        cur_words[idx] = choices[idx][best_replacement_idx]
        return cur_words, worst_margin

    def run(self, model, texts_p, texts_h, labels, device, genetic_test_num, opts=None):
        is_correct = []
        adv_exs = []
        total = 0
        acc = 0

        #texts_p = texts_p[opts.h_test_start:opts.h_test_start+opts.genetic_test_num]
        #texts_h = texts_h[opts.h_test_start:opts.h_test_start+opts.genetic_test_num]
        #labels = labels[opts.h_test_start:opts.h_test_start+opts.genetic_test_num]

        for x_p, x_h, y in zip(texts_p, texts_h, labels):

            words = x_h.split()
            if not self.attack_surface.check_in(words):
                continue

            print(acc, "/", total)
            if total >= genetic_test_num:
                break
            total += 1
            y = np.argmax(y)
            # First query the example itself
            orig_pred = model.query(x_p, x_h, device)
            if np.argmax(orig_pred) != y :
                print('ORIGINAL PREDICTION WAS WRONG')
                is_correct.append(0)
                adv_exs.append((x_p, x_h))
                continue
            # Now run adversarial search
            x_h_words = x_h.split()
            swaps = self.attack_surface.get_swaps(x_h_words)
            choices = [[w] + cur_swaps for w, cur_swaps in zip(x_h_words, swaps)]
            found = False
            population = [self.perturb(x_p, x_h_words, choices, model, y, device)
                            for i in range(self.pop_size)]
            for g in range(self.num_iters):
                best_idx = min(enumerate(population), key=lambda x: x[1][1])[0]
                #print('Iteration %d: %.6f' % (g, population[best_idx][1]))
                if population[best_idx][1] < self.margin_goal:
                    found = True
                    is_correct.append(0)
                    adv_exs.append(' '.join(population[best_idx][0]))
                    #print('ADVERSARY SUCCESS on ("%s", %d): Found "%s" with margin %.2f' % (x, y, adv_exs[-1], population[best_idx][1]))
                    print('ADVERSARY SUCCESS')
                    break
                new_population = [population[best_idx]]

                margins = np.array([m for c, m in population])
                adv_probs = 1 / (1 + np.exp(margins)) + 1e-4
                # Sigmoid of negative margin, for probabilty of wrong class
                # Add 1e-6 for numerical stability
                sample_probs = adv_probs / np.sum(adv_probs)

                for i in range(1, self.pop_size):
                    parent1 = population[np.random.choice(range(len(population)), p=sample_probs)][0]
                    parent2 = population[np.random.choice(range(len(population)), p=sample_probs)][0]
                    child = [random.sample([w1, w2], 1)[0] for (w1, w2) in zip(parent1, parent2)]
                    child_mut, new_margin = self.perturb(x_p, child, choices, model, y, device)
                    new_population.append((child_mut, new_margin))
                population = new_population
            else:
                is_correct.append(1)
                adv_exs.append([])
                acc+=1
                #print('ADVERSARY FAILURE on ("%s", %d)' % (x, y))
                print('ADVERSARY FAILURE', 'Iteration %d: %.6f' % (g, population[best_idx][1]))
        if total != 0:
            return acc*1.0/total
        else:
            return 0

class GeneticAdversary(Adversary):
    def __init__(self, attack_surface, num_iters=20, pop_size=60, margin_goal=0.5):
        super(GeneticAdversary, self).__init__(attack_surface)
        self.num_iters = num_iters
        self.pop_size = pop_size
        self.margin_goal = margin_goal

    def perturb(self, words, choices, model, y, device):
        if all(len(c) == 1 for c in choices): return words
        good_idxs = [i for i, c in enumerate(choices) if len(c) > 1]
        idx = random.sample(good_idxs, 1)[0]
        x_list = [' '.join(words[:idx] + [w_new] + words[idx+1:])
                for w_new in choices[idx]]
        preds = [model.query(x, device) for x in x_list]
        preds_of_y = [p[y] for p in preds]
        best_idx = min(enumerate(preds_of_y), key=lambda x: x[1])[0]
        cur_words = list(words)
        cur_words[idx] = choices[idx][best_idx]
        return cur_words, preds_of_y[best_idx]

    def run(self, model, texts, labels, device, genetic_test_num, opts=None):
        is_correct = []
        adv_exs = []
        total = 0
        acc = 0

        #texts = texts[opts.h_test_start:opts.h_test_start+opts.genetic_test_num]
        #labels = labels[opts.h_test_start:opts.h_test_start+opts.genetic_test_num]

        for x, y in zip(texts, labels):
            words = x.split()
            if not self.attack_surface.check_in(words):
                continue

            print(acc, "/", total)
            if total >= genetic_test_num:
                break
            total += 1
            y = np.argmax(y)
            # First query the example itself
            orig_pred = model.query(x, device)
            if np.argmax(orig_pred) != y :
                print('ORIGINAL PREDICTION WAS WRONG')
                is_correct.append(0)
                adv_exs.append(x)
                continue
            # Now run adversarial search
            words = x.split()
            swaps = self.attack_surface.get_swaps(words)
            choices = [[w] + cur_swaps for w, cur_swaps in zip(words, swaps)]
            found = False
            population = [self.perturb(words, choices, model, y, device)
                            for i in range(self.pop_size)]
            for g in range(self.num_iters):
                best_idx = min(enumerate(population), key=lambda x: x[1][1])[0]
                #print('Iteration %d: %.6f' % (g, population[best_idx][1]))
                if population[best_idx][1] < self.margin_goal:
                    found = True
                    is_correct.append(0)
                    adv_exs.append(' '.join(population[best_idx][0]))
                    #print('ADVERSARY SUCCESS on ("%s", %d): Found "%s" with margin %.2f' % (x, y, adv_exs[-1], population[best_idx][1]))
                    print('ADVERSARY SUCCESS')
                    break
                new_population = [population[best_idx]]
                p_y = np.array([m for c, m in population])
                temp = 1-p_y + 1e-8
                sample_probs = (temp) / np.sum(temp)
                #sample_probs = sample_probs + 1e-8
                for i in range(1, self.pop_size):
                    parent1 = population[np.random.choice(range(len(population)), p=sample_probs)][0]
                    parent2 = population[np.random.choice(range(len(population)), p=sample_probs)][0]
                    child = [random.sample([w1, w2], 1)[0] for (w1, w2) in zip(parent1, parent2)]
                    child_mut, new_p_y = self.perturb(child, choices, model, y, device)
                    new_population.append((child_mut, new_p_y))
                population = new_population
            else:
                is_correct.append(1)
                adv_exs.append([])
                acc+=1
                #print('ADVERSARY FAILURE on ("%s", %d)' % (x, y))
                print('ADVERSARY FAILURE', 'Iteration %d: %.6f' % (g, population[best_idx][1]))
        if total != 0:
            return acc*1.0/total
        else:
            return 0


from __future__ import print_function
import numpy as np
from discretize.utils import mkvc

########################################################################################
#                  DIRECT CURRENT RESISTIVITY AND INDUCED POLARIZATION
########################################################################################

def read_dcip3d_ubc(file_name, data_type):
    """
    Read UBC DCIP3D formatted survey, predicted or observed data files.

    This function loads 3D DC or IP data formatted for the UBC DCIP3D
    coding package and outputs a SimPEG data object. The function will
    recognize if the file is a survey, predicted data or observed data file.

    Parameters
    ----------
    file_name : str
        The file path to the data file
    data_type: str {'volt', 'apparent_chargeability', 'secondary_potential'}
        Defining the input data type:

        - 'volt': DC resistivity data as voltages
        - 'apparent_chargeability': IP data as apparent chargeabilities
        - 'secondary_potential': IP data as secondary potentials

    Returns
    -------
    data
        A SimPEG.data.Data object containing:

        - The survey
        - Observed/predicted data (if present in the data file)
        - Uncertainties (if present in the data file). Note that predicted DC data
        files contain the apparent resistivities, which are loaded into SimPEG and
        defined as uncertainties.

    """
    assert data_type.lower() in [
        "volt",
        "apparent_chargeability",
        "secondary_potential",
    ], "Parameter 'data_type' must be one of {'volt', 'apparent_chargeability', 'secondary_potential'}"

    return _read_dcip_3d_or_octree_ubc(file_name, data_type, "dcip3d")

def read_dcipoctree_ubc(file_name, data_type):
    """
    Read UBC DCIPoctree formatted survey, predicted or observed data files.

    This function loads 3D DC or IP data formatted for the UBC DCIPoctree
    coding package and outputs a SimPEG data object. The function requires
    the user to define whether the data are DC resistivity or IP.

    Parameters
    ----------
    file_name : str
        The file path to the data file
    data_type: str {'volt', 'apparent_chargeability'}
        Defining the input data type:

        - 'volt': DC resistivity data as voltages
        - 'apparent_chargeability': IP data as apparent chargeabilities

    Returns
    -------
    data_object
        A SimPEG.data.Data object containing:

        - The survey
        - Observed/predicted data (if present in the data file)
        - Uncertainties (if present in the data file). Note that predicted DC data
        files contain the apparent resistivities, which are loaded into SimPEG and
        defined as uncertainties.

    """

    # Unused for now but it will be when we manage IP types better.
    assert data_type.lower() in [
        "volt",
        "apparent_chargeability",
    ], "Parameter 'data_type' must be one of {'volt', 'apparent_chargeability'}"

    return _read_dcip_3d_or_octree_ubc(file_name, data_type, "dcipoctree")

def _read_dcip_3d_or_octree_ubc(file_name, data_type, code_type):
    """
    Read 3D DC/IP survey, predicted and observation files in UBC-GIF format.

    Parameters
    ----------

    file_name : str
        The file path to the data file
    code_type : str {'dcip3d', 'dcipoctree'}
        Code type. Choose from {'dcip3d', 'dcipoctree'}

    Returns
    -------
    data_object
        A SimPEG.data.Data object containing:

        - The survey
        - Observed/predicted data (if present in the data file)
        - Uncertainties (if present in the data file). Note that predicted DC data
        files contain the apparent resistivities, which are loaded into SimPEG and
        defined as uncertainties.

    """
    assert data_type.lower() in [
        "volt",
        "apparent_chargeability",
        "secondary_potential",
    ], "Parameter 'data_type' must be one of {'volt', 'apparent_chargeability', 'secondary_potential'}"

    assert code_type.lower() in [
        "dcip3d",
        "dcipoctree",
    ], "Parameter 'code_type' must be one of {'dcip3d', 'dcipoctree'}"

    # Prevent circular import
    from ...electromagnetics.static import resistivity as dc
    from ...data import Data

    # Load file
    obsfile = np.genfromtxt(file_name, delimiter=" \n", dtype=np.str, comments="!")

    # Pre-allocate
    source_list = []
    receiver_list = []
    d = []
    wd = []

    # Flag for z value provided
    is_surface = False
    is_pole_tx = False
    is_pole_rx = False

    # IP data for dcip3d has a line with a flag we can remove.
    if (code_type == "dcip3d") & (data_type != "volt"):
        obsfile = obsfile[1:]

    # Since SimPEG defines secondary potential from IP as voltage,
    # we must use this type when defining the receivers.
    if data_type.lower() == "secondary_potential":
        data_type = "volt"

    # Countdown for number of obs/tx
    count = 0
    for ii in range(obsfile.shape[0]):

        if not obsfile[ii]:
            continue

        # Extract transmitter ii and the number of receivers
        if count == 0:
            rx = []
            temp = np.fromstring(obsfile[ii], dtype=float, sep=" ").T
            count = int(temp[-1])

            # Check if z value is provided, if False -> nan
            if len(temp) == 5:
                # check if pole|dipole
                if np.allclose(temp[0:2], temp[2:4]):
                    tx = np.r_[temp[0:2], np.nan]
                    is_pole_tx = True
                else:
                    tx = np.r_[temp[0:2], np.nan, temp[2:4], np.nan]
                is_surface = True

            else:
                # check if pole|dipole
                if np.allclose(temp[0:3], temp[3:6]):
                    tx = np.r_[temp[0:3]]
                    is_pole_tx = True
                else:
                    tx = temp[:-1]

            continue

        # Extract receivers
        temp = np.fromstring(obsfile[ii], dtype=float, sep=" ")

        if is_surface:
            data_column_index = 4  # Since dpred for dc has app_res

            # Check if Pole Receiver
            if np.allclose(temp[0:2], temp[2:4]):
                is_pole_rx = True
                rx.append(temp[:2])
            else:
                rx.append(np.r_[temp[0:2], np.nan, temp[2:4], np.nan])

        else:
            data_column_index = 6  # Since dpred for dc has app_res

            # Check if Pole Receiver
            if np.allclose(temp[0:3], temp[3:6]):
                is_pole_rx = True
                rx.append(temp[:3])
            else:
                rx.append(temp[:6])

        # Predicted/observed data
        if len(temp) == data_column_index + 1:
            d.append(temp[data_column_index])

        # Observed data or predicted DC data (since app res column)
        elif len(temp) == data_column_index + 2:
            d.append(temp[data_column_index])
            wd.append(temp[data_column_index + 1])

        count = count - 1

        # Reach the end of transmitter block
        if count == 0:
            rx = np.asarray(rx)
            if is_pole_rx:
                Rx = dc.receivers.Pole(rx[:, :3], data_type=data_type)
            else:
                Rx = dc.receivers.Dipole(rx[:, :3], rx[:, 3:], data_type=data_type)
            if is_pole_tx:
                source_list.append(dc.sources.Pole([Rx], tx[:3]))
            else:
                source_list.append(dc.sources.Dipole([Rx], tx[:3], tx[3:]))

    # Define survey type
    if is_pole_tx:
        str1 = "pole-"
    else:
        str1 = "dipole-"

    if is_pole_rx:
        str2 = "pole"
    else:
        str2 = "dipole"

    electrode_configuration = str1 + str2
    survey = dc.survey.Survey(source_list, survey_type=electrode_configuration)
    data_out = Data(survey=survey)

    if len(d) > 0:
        data_out.dobs = d

    if len(wd) > 0:
        data_out.standard_deviation = wd

    return data_out

def write_dcip3d_ubc(
    file_name,
    data_object,
    data_type,
    file_type,
    format_type="general",
    electrode_configuration=None,
    comment_lines="",
):
    """
    Write UBC DCIP3D formatted survey, predicted or observation files.

    Parameters
    ----------
    file_name:
    data_object:
    file_type: 'survey', 'dpred', 'dobs'
    format_type: 'general', 'surface', 'simple'
    data_type: 'volt', 'apparent_chargeability', 'secondary_potential'
    electrode_configuration: 'pole-pole', 'pole-dipole', 'dipole-pole', 'dipole-dipole'
    comment_lines:)
    """

    assert data_type.lower() in [
        "volt",
        "apparent_chargeability",
        "secondary_potential",
    ], "Parameter 'data_type' must be one of {'volt', 'apparent_chargeability', 'secondary_potential'}"

    assert file_type.lower() in [
        "survey",
        "dpred",
        "dobs",
    ], "Parameter 'file_type' must be one of {'survey', 'dpred', 'dobs'}"

    _write_dcip_3d_or_octree_ubc(
        file_name,
        data_object,
        data_type,
        file_type,
        format_type=format_type,
        electrode_configuration=electrode_configuration,
        code_type="dcip3d",
        comment_lines=comment_lines,
    )

def write_dcipoctree_ubc(
    file_name,
    data_object,
    data_type,
    file_type,
    format_type="general",
    electrode_configuration=None,
    comment_lines="",
):
    """
    Write UBC DCIPoctree formatted survey, predicted or observation files.

    Parameters
    ----------
    file_name:
    data_object:
    file_type: 'survey', 'dpred', 'dobs'
    format_type: 'general', 'surface', 'simple'
    data_type: 0 (DC), 1 (IP), 2 (another IP)
    electrode_configuration: 'pole-pole', 'pole-dipole', 'dipole-pole', 'dipole-dipole'
    comment_lines:)
    """

    assert data_type.lower() in [
        "volt",
        "apparent_chargeability",
    ], "Parameter 'data_type' must be one of {'volt', 'apparent_chargeability'}"

    assert file_type.lower() in [
        "survey",
        "dpred",
        "dobs",
    ], "Parameter 'file_type' must be one of {'survey', 'dpred', 'dobs'}"

    _write_dcip_3d_or_octree_ubc(
        file_name,
        data_object,
        data_type,
        file_type,
        format_type="general",
        electrode_configuration=electrode_configuration,
        code_type="dcipoctree",
        comment_lines=comment_lines,
    )

def _write_dcip_3d_or_octree_ubc(
    file_name,
    data_object,
    data_type,
    file_type,
    format_type="general",
    electrode_configuration=None,
    code_type="dcip3d",
    comment_lines=None,
):
    """
    Write UBC DCIP3D formatted survey, predicted or observation files.

    Parameters
    ----------
    file_name:
    data_object:
    file_type: 'survey', 'dpred', 'dobs'
    format_type: 'general', 'surface', 'simple'
    data_type: {'volt', 'apparent_chargeability', 'secondary_potential'}
    electrode_configuration: 'pole-pole', 'pole-dipole', 'dipole-pole', 'dipole-dipole'
    code_type: 'dcip3d', 'dcipoctree'
    comment_lines:
    """

    # Prevent circular import
    from ...electromagnetics.static import resistivity as dc
    from ...electromagnetics.static.utils.static_utils import apparent_resistivity
    from ...data import Data

    # Validate inputs
    if not isinstance(data_object, Data):
        raise Exception(
            "A Data instance ({datacls}: <{datapref}.{datacls}>) must be "
            "provided as the second input. The provided input is a "
            "{providedcls} <{providedpref}.{providedcls}>".format(
                datacls=Data.__name__,
                datapref=Data.__module__,
                providedcls=data_object.__class__.__name__,
                providedpref=data_object.__module__,
            )
        )

    assert data_type.lower() in [
        "volt",
        "apparent_chargeability",
        "secondary_potential",
    ], "Parameter 'data_type' must be one of {'volt', 'apparent_chargeability', 'secondary_potential'}"

    assert file_type.lower() in [
        "survey",
        "dpred",
        "dobs",
    ], "Parameter 'file_type' must be one of {'survey', 'dpred', 'dobs'}"

    assert code_type.lower() in [
        "dcip3d",
        "dcipoctree",
    ], "Parameter 'code_type' must be one of {'dcip3d', 'dcipoctree'}"

    format_type = format_type.lower()
    if format_type not in ["surface", "general", "simple"]:
        raise Exception(
            "format_type must be 'surface' | 'general' | 'simple' "
            " not {}".format(format_type)
        )

    if electrode_configuration is None:
        electrode_configuration = data_object.survey.survey_type

    # Predicted DC data will automatically contain apparent resistivity column.
    # Here we compute the apparent resistivities and treat it like an uncertainties column.
    if (file_type.lower() == "dpred") & (data_type == "volt"):
        data_object.standard_deviation = apparent_resistivity(data_object)
        file_type = "dobs"

    # Write comments and IP type (if applicable)
    with open(file_name, "w") as fid:
        fid.write(f"! {format_type} FORMAT\n")

        if comment_lines is not None and len(comment_lines) > 0:
            # ensure comment_lines ends with a new line character
            if comment_lines[-1] != "\n":
                comment_lines += "\n"
            fid.write(comment_lines)

        # DCIP3D will allow user to choose definition of IP data. DC data has no flag.
        # DCIPoctree IP data is always apparent chargeability.
        if (code_type.lower() == "dcip3d") & (data_type == "apparent_chargeability"):
            fid.write("IPTYPE=1\n")
        elif (code_type.lower() == "dcip3d") & (data_type == "secondary_potential"):
            fid.write("IPTYPE=2\n")

        # Index deciding if z locations are written
        if format_type.lower() == "surface":
            end_index = 2
        elif format_type.lower() == "general":
            end_index = 3

        # Loop over all sources
        count = 0
        for src in data_object.survey.source_list:

            # Write Source
            nD = src.nD

            if electrode_configuration.lower() in ["pole-dipole", "pole-pole"]:
                tx = np.r_[src.location]
                tx = np.repeat(np.r_[[tx]], 2, axis=0)
            elif electrode_configuration.lower() in ["dipole-dipole", "dipole-pole"]:
                tx = np.c_[src.location]

            fid.writelines("%e " % ii for ii in mkvc(tx[:, 0:end_index].T))
            fid.write(f"{nD}\n")

            # Write receivers
            for rx in src.receiver_list:

                if electrode_configuration.lower() in ["pole-dipole", "dipole-dipole"]:
                    M = rx.locations[0][0:end_index]
                    N = rx.locations[1][0:end_index]
                elif electrode_configuration.lower() in ["pole-pole", "dipole-pole"]:
                    M = rx.locations[0:end_index]
                    N = rx.locations[0:end_index]

                if file_type.lower() != "survey":
                    N = np.c_[N, data_object.dobs[count : count + rx.nD]]

                if file_type.lower() == "dobs":
                    N = np.c_[N, data_object.standard_deviation[count : count + rx.nD]]

                # Write receivers and locations
                if isinstance(N, np.ndarray):
                    np.savetxt(
                        fid, np.c_[M, N], fmt="%e",
                    )
                else:
                    raise Exception(
                        """Uncertainities SurveyObject.std should be set.
                        Either float or nunmpy.ndarray is expected, """
                        "not {}".format(type(data_object.relative_error))
                    )

                fid.write("\n")

                count += rx.nD


import argparse
import sys
import torch
import torch.nn as nn
import numpy as np
from torch.utils.data import DataLoader
from torchvision import transforms
import multiprocessing
from torch.autograd import Variable
from modules.models.Seq2Seq_atn import EncoderCRNN, AttnDecoderRNN
from modules.core.dataloader_test import SequenceDataset
from modules.handlers.TextColor import TextColor
from collections import defaultdict
from modules.handlers.VcfWriter import VCFWriter
from modules.handlers.FileManager import FileManager
import operator
import pickle
from tqdm import tqdm
import os
import time
"""
This script uses a trained model to call variants on a given set of images generated from the genome.
The process is:
- Create a prediction table/dictionary using a trained neural network
- Convert those predictions to a VCF file

INPUT:
- A trained model
- Set of images for prediction

Output:
- A VCF file containing all the variants.
"""
FLANK_SIZE = 10

SNP = 1
IN = 2
DEL = 3
HOM = 0
HET = 1
HOM_ALT = 2

prediction_dict = defaultdict(list)
reference_dict = defaultdict(tuple)

def predict(test_file, batch_size, model_path, gpu_mode, num_workers):
    """
    Create a prediction table/dictionary of an images set using a trained model.
    :param test_file: File to predict on
    :param batch_size: Batch size used for prediction
    :param model_path: Path to a trained model
    :param gpu_mode: If true, predictions will be done over GPU
    :param num_workers: Number of workers to be used by the dataloader
    :return: Prediction dictionary
    """
    # the prediction table/dictionary
    chromosome_name = ''
    transformations = transforms.Compose([transforms.ToTensor()])

    sys.stderr.write(TextColor.PURPLE + 'Loading data\n' + TextColor.END)

    test_dset = SequenceDataset(test_file, transformations)
    testloader = DataLoader(test_dset,
                            batch_size=batch_size,
                            shuffle=False,
                            num_workers=num_workers
                            )

    sys.stderr.write(TextColor.PURPLE + 'Data loading finished\n' + TextColor.END)

    # load the model
    checkpoint = torch.load(model_path, map_location='cpu')
    encoder_state_dict = checkpoint['encoder_state_dict']
    decoder_state_dict = checkpoint['decoder_state_dict']

    from collections import OrderedDict
    new_encoder_state_dict = OrderedDict()
    new_decoder_state_dict = OrderedDict()

    for k, v in encoder_state_dict.items():
        name = k
        if k[0:7] == 'module.':
            name = k[7:]  # remove `module.`
        new_encoder_state_dict[name] = v

    for k, v in decoder_state_dict.items():
        name = k
        if k[0:7] == 'module.':
            name = k[7:]  # remove `module.`
        new_decoder_state_dict[name] = v

    hidden_size = 256
    encoder_model = EncoderCRNN(image_channels=10, hidden_size=hidden_size)
    decoder_model = AttnDecoderRNN(hidden_size=hidden_size, num_classes=6, max_length=1)
    encoder_model.load_state_dict(new_encoder_state_dict)
    decoder_model.load_state_dict(new_decoder_state_dict)
    encoder_model.cpu()
    decoder_model.cpu()

    if gpu_mode:
        encoder_model = encoder_model.cuda()
        encoder_model = torch.nn.DataParallel(encoder_model).cuda()
        decoder_model = decoder_model.cuda()
        decoder_model = torch.nn.DataParallel(decoder_model).cuda()

    # Change model to 'eval' mode (BN uses moving mean/var).
    encoder_model.eval()
    decoder_model.eval()

    sys.stderr.write(TextColor.PURPLE + 'MODEL LOADED\n' + TextColor.END)
    # TO HERE
    with torch.no_grad():
        for images, labels, positional_info in tqdm(testloader, file=sys.stdout, dynamic_ncols=True):
            if gpu_mode:
                # encoder_hidden = encoder_hidden.cuda()
                images = images.cuda()
                labels = labels.cuda()

            decoder_input = torch.LongTensor(labels.size(0), 1).zero_()
            encoder_hidden = torch.FloatTensor(labels.size(0), 2, hidden_size).zero_()

            # if gpu_mode:
            #     decoder_input = decoder_input.cuda()
            #     encoder_hidden = encoder_hidden.cuda()

            chr_name, start_positions, reference_seqs, allele_dict_paths = positional_info

            window_size = images.size(2) - 2 * FLANK_SIZE
            index_start = FLANK_SIZE
            end_index = index_start + window_size
            unrolling_genomic_position = np.zeros((images.size(0)), dtype=np.int64)

            for seq_index in range(index_start, end_index):
                x = images[:, :, seq_index - FLANK_SIZE:seq_index + FLANK_SIZE + 1, :]

                output_enc, hidden_dec = encoder_model(x, encoder_hidden)
                output_dec, decoder_hidden, attn = decoder_model(decoder_input, output_enc, hidden_dec)

                encoder_hidden = decoder_hidden.detach()

                topv, topi = output_dec.topk(1)
                decoder_input = topi.squeeze().detach()  # detach from history as input

                # One dimensional softmax is used to convert the logits to probability distribution
                m = nn.Softmax(dim=1)
                soft_probs = m(output_dec)
                output_preds = soft_probs.cpu()
                # record each of the predictions from a batch prediction
                batches = images.size(0)

                for batch in range(batches):
                    allele_dict_path = allele_dict_paths[batch]
                    chromosome_name = chr_name[batch]
                    reference_seq = reference_seqs[batch]
                    # current_genomic_position = int(start_positions[batch])
                    current_genomic_position = int(start_positions[batch]) + unrolling_genomic_position[batch]

                    ref_base = reference_seq[seq_index]

                    if ref_base == '*':
                        continue

                    # true_label = labels[batch, seq_index - index_start]
                    # fake_probs = [0.0] * 6
                    # fake_probs[true_label] = 1.0
                    # top_n, top_i = torch.FloatTensor(fake_probs).topk(1)
                    # predicted_label = top_i[0].item()
                    # reference_dict[current_genomic_position] = (ref_base, allele_dict_path)
                    # prediction_dict[current_genomic_position].append((predicted_label, fake_probs))

                    preds = output_preds[batch, :].data
                    top_n, top_i = preds.topk(1)
                    predicted_label = top_i[0].item()
                    reference_dict[current_genomic_position] = (ref_base, allele_dict_path)
                    prediction_dict[current_genomic_position].append((predicted_label, preds))

                    if ref_base != '*':
                        unrolling_genomic_position[batch] += 1
    return chromosome_name

def get_record_from_prediction(pos, alleles):
    predictions = prediction_dict[pos]
    genotype, qual, gq = VCFWriter.process_prediction(pos, predictions)
    alts = list(allele[0] for allele in alleles)
    ref_base = reference_dict[pos][0][0]
    return ref_base, alts, genotype, qual, gq

def produce_vcf_records(chromosome_name, output_dir, thread_no, pos_list):
    """
    Convert prediction dictionary to a VCF file
    :param: arg_tuple: Tuple of arguments containing these values:
    - chromosome_name: Chromosome name
    - pos_list: List of positions where we will search for variants
    - prediction_dict: prediction dictionary containing predictions of each image records
    - reference_dict: Dictionary containing reference information
    - bam_file_path: Path to the BAM file
    - sample_name: Name of the sample in the BAM file
    - output_dir: Output directory
    - thread_id: Unique id assigned to each thread
    :return:
    """
    # object that can write and handle VCF
    # vcf_writer = VCFWriter(bam_file_path, sample_name, output_dir, thread_id)
    # collate multi-allelic records to a single record
    current_allele_dict = ''
    allele_dict = {}
    record_file = open(output_dir + chromosome_name + "_" + str(thread_no) + ".tsv", 'w')
    for pos in pos_list:
        allele_dict_path = reference_dict[pos][1]
        if allele_dict_path != current_allele_dict:
            allele_dict = pickle.load(open(allele_dict_path, 'rb'))
            current_allele_dict = allele_dict_path

        if pos not in allele_dict:
            continue
        alleles = allele_dict[pos]

        record = get_record_from_prediction(pos, alleles)

        if record is None:
            continue

        ref_base, alts, genotype, qual, gq = record

        if genotype == '0/0':
            continue
        # print('BEFORE', record)
        record = VCFWriter.get_proper_alleles(record)
        ref, alts, qual, gq, genotype = record
        # print('AFTER', record)
        if len(alts) == 1:
            alts.append('.')
        rec_end = int(pos + len(ref) - 1)
        record_string = chromosome_name + "\t" + str(pos) + "\t" + str(rec_end) + "\t" + ref + "\t" + '\t'.join(alts) \
                        + "\t" + genotype + "\t" + str(qual) + "\t" + str(gq) + "\t" + "\n"
        record_file.write(record_string)

def merge_call_files(vcf_file_directory):
    filemanager_object = FileManager()
    # get all bed file paths from the directory
    file_paths = filemanager_object.get_file_paths_from_directory(vcf_file_directory)
    all_records = []
    for file_path in file_paths:
        with open(file_path, 'r') as tsv:
            for line in tsv:
                chr_name, pos_st, pos_end, ref, alt1, alt2, genotype, qual, gq = line.strip().split('\t')
                alts = []
                pos_st, pos_end, qual, gq = int(pos_st), int(pos_end), float(qual), float(gq)
                if alt1 != '.':
                    alts.append(alt1)
                if alt2 != '.':
                    alts.append(alt2)
                all_records.append((chr_name, pos_st, pos_end, ref, alts, genotype, qual, gq))

    filemanager_object.delete_files(file_paths)
    os.rmdir(vcf_file_directory)

    return all_records

def call_variant(csv_file, batch_size, model_path, gpu_mode, num_workers, bam_file_path, sample_name, output_dir,
                 vcf_dir, max_threads):
    program_start_time = time.time()
    sys.stderr.write(TextColor.GREEN + "INFO: " + TextColor.END + "SAMPLE NAME: " + sample_name + "\n")
    sys.stderr.write(TextColor.GREEN + "INFO: " + TextColor.END + "PLEASE USE --sample_name TO CHANGE SAMPLE NAME.\n")
    sys.stderr.write(TextColor.GREEN + "INFO: " + TextColor.END + "OUTPUT DIRECTORY: " + output_dir + "\n")
    chr_name = predict(csv_file, batch_size, model_path, gpu_mode, num_workers)
    sys.stderr.write(TextColor.GREEN + "INFO: " + TextColor.END + "PREDICTION GENERATED SUCCESSFULLY.\n")
    sys.stderr.write(TextColor.GREEN + "INFO: " + TextColor.END + "COMPILING PREDICTIONS TO CALL VARIANTS.\n")

    pos_list = list(prediction_dict.keys())
    each_chunk_size = int(len(pos_list) / max_threads)
    thread_no = 1
    # produce_vcf_records(chr_name, vcf_dir, thread_no, pos_list)
    # exit()

    for i in tqdm(range(0, len(pos_list), each_chunk_size), file=sys.stdout, dynamic_ncols=True):
        start_position = i
        end_position = min(i + each_chunk_size, len(pos_list))

        sub_pos = pos_list[start_position:end_position]
        # gather all parameters
        args = (chr_name, vcf_dir, thread_no, sub_pos)
        p = multiprocessing.Process(target=produce_vcf_records, args=args)
        p.start()
        thread_no += 1

        # wait until we have room for new processes to start
        while True:
            if len(multiprocessing.active_children()) < max_threads:
                break

    # wait until we have room for new processes to start
    while True:
        if len(multiprocessing.active_children()) == 0:
            break
    sys.stderr.write(TextColor.GREEN + "INFO: " + TextColor.END + "VARIANT CALLING COMPLETE.\n")
    sys.stderr.write(TextColor.GREEN + "INFO: " + TextColor.END + "MERGING FILES.\n")
    all_calls = merge_call_files(vcf_dir)

    # sort based on position
    all_calls.sort(key=operator.itemgetter(1))
    # print(all_calls)
    last_end = 0
    sys.stderr.write(TextColor.GREEN + "INFO: " + TextColor.END + "WRITING VCF.\n")
    vcf_writer = VCFWriter(bam_file_path, sample_name, output_dir)
    for record in all_calls:
        # get the record filter ('PASS' or not)
        rec_filter = VCFWriter.get_filter(record, last_end)
        # get proper alleles. INDEL alleles are handled here.
        # record = VCFWriter.get_proper_alleles(record)
        chrm, st_pos, end_pos, ref, alt_field, genotype, phred_qual, phred_gq = record
        # if genotype is not HOM keep track of where the previous record ended
        if genotype != '0/0':
            # HOM
            last_end = end_pos
        # add the record to VCF
        vcf_writer.write_vcf_record(chrm, st_pos, end_pos, ref, alt_field, genotype, phred_qual, phred_gq, rec_filter)

    sys.stderr.write(TextColor.GREEN + "INFO: " + TextColor.END + "VARIANT CALLING COMPLETE.\n")
    program_end_time = time.time()

    sys.stderr.write(TextColor.PURPLE + "TIME ELAPSED: " + str(program_end_time - program_start_time) + "\n")

def handle_output_directory(output_dir):
    """
    Process the output directory and return a valid directory where we save the output
    :param output_dir: Output directory path
    :return:
    """
    # process the output directory
    if output_dir[-1] != "/":
        output_dir += "/"
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    vcf_path = output_dir + "vcfs" + "/"
    if not os.path.exists(vcf_path):
        os.mkdir(vcf_path)

    return output_dir, vcf_path

if __name__ == '__main__':
    '''
    Processes arguments and performs tasks.
    '''
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--csv_file",
        type=str,
        required=True,
        help="CSV file containing all image segments for prediction."
    )
    parser.add_argument(
        "--bam_file",
        type=str,
        required=True,
        help="Path to the BAM file."
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        required=False,
        default=100,
        help="Batch size for testing, default is 100."
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        required=False,
        default=4,
        help="Batch size for testing, default is 100."
    )
    parser.add_argument(
        "--model_path",
        type=str,
        default='./CNN.pkl',
        help="Saved model path."
    )
    parser.add_argument(
        "--gpu_mode",
        type=bool,
        default=False,
        help="If true then cuda is on."
    )
    parser.add_argument(
        "--sample_name",
        type=str,
        required=False,
        default='NA12878',
        help="Sample name of the sequence."
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        required=False,
        default='vcf_output',
        help="Output directory."
    )
    parser.add_argument(
        "--max_threads",
        type=int,
        default=8,
        help="Number of maximum threads for this region."
    )
    FLAGS, unparsed = parser.parse_known_args()
    FLAGS.output_dir, vcf_dir = handle_output_directory(FLAGS.output_dir)
    call_variant(FLAGS.csv_file,
                 FLAGS.batch_size,
                 FLAGS.model_path,
                 FLAGS.gpu_mode,
                 FLAGS.num_workers,
                 FLAGS.bam_file,
                 FLAGS.sample_name,
                 FLAGS.output_dir,
                 vcf_dir,
                 FLAGS.max_threads)


import myo
import numpy as np
import time
import os
import keyboard
import threading
import collections
import sys
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.models import load_model

Pulgar = [0,1,2,3]
Indice = [0,1,2,3,4,5,6]
Mayor = [0,1,2,3,4,5,6]
Anular = [0,1,2,3,4,5,6]
Menique = [0,1,2,3,4,5,6]

class Listener(myo.DeviceListener):

  def __init__(self, queue_size=8):
    self.lock = threading.Lock()
    self.emg_data_queue = collections.deque(maxlen=queue_size)

  def on_connected(self, event):
    event.device.stream_emg(True)

  def on_emg(self, event):
    with self.lock:
      self.emg_data_queue.append((event.timestamp, event.emg))

  def get_emg_data(self):
    with self.lock:
      return list(self.emg_data_queue)

def step(listener):
    emgs = np.array([x[1] for x in listener.get_emg_data()])
    return emgs

def main(model):
    numb = 0
    queue_size=10
    myo.init()
    hub = myo.Hub()
    listener = Listener(queue_size)
    try:
        threading.Thread(target=lambda: hub.run_forever(listener.on_event)).start()
        t = 0
        flag = 0
        data = []
        while True:
            emgs = step(listener)
            if len(emgs) != 10:
                continue
            if len(data) < 50:
                emgs = tf.keras.utils.normalize(emgs, axis = 0)
                data.append(emgs.reshape(1,80))
                continue
            else:
                predictions = []
                for x in data:
                    predictions.append(model.predict(x).argmax())
                os.system('cls')
                #print(predictions)
                data = []
                print(np.bincount(predictions).argmax())
    finally:
        hub.stop()

if __name__ == '__main__':
    '''Y = np.loadtxt("datasets/PruebaY.csv", delimiter=",")
    X = np.loadtxt("datasets/PruebaX.csv", delimiter=",")
    X = tf.keras.utils.normalize(X, axis=1)
    Y = Y.astype(int)
    config = tf.ConfigProto( device_count = {'GPU': 1})
    sess = tf.Session(config=config)
    tf.keras.backend.set_session(sess)
    model = Sequential()
    model.add(Dense(128, input_dim=X.shape[1]))
    model.add(Activation('relu'))
    model.add(Dropout(0.1))
    model.add(Dense(128))
    model.add(Activation('relu'))
    model.add(Dropout(0.1))
    model.add(Dense(7))
    model.add(Activation('softmax'))
    model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
    model.fit(X, Y[0:,3], epochs=300, batch_size=32)'''
    model = load_model("modelo.h5")
    main(model)


#!/usr/bin/env python
""" LBR iiwa poke an object tracked by RealSense D435
    Depend on:
        roslaunch iiwa_gazebo iiwa_gazebo_with_sunrise.launch (tf defined in 'iiwa_gazebo' package)
"""
from __future__ import absolute_import, division, print_function, unicode_literals

import sys
import os
import time

import cv2
import pyrealsense2 as rs
import numpy as np
from numpy import pi
from glob import glob

import torch

from pysot.core.config import cfg
from pysot.models.model_builder import ModelBuilder
from pysot.tracker.tracker_builder import build_tracker

import rospy
import tf
from robots.lbr import iiwaRobot
from geometry_msgs.msg import PoseStamped, Quaternion
from iiwa_msgs.msg import JointPosition

# iiwa's initial perching pose
JOINT_PERCH = JointPosition()
JOINT_PERCH.position.a1 = -pi/6
JOINT_PERCH.position.a2 = pi/3
JOINT_PERCH.position.a3 = pi/6
JOINT_PERCH.position.a4 = -pi/2
JOINT_PERCH.position.a5 = -pi/6
JOINT_PERCH.position.a6 = -pi/4
JOINT_PERCH.position.a7 = -pi/6

def quat_from_vecs(vec_0, vec_1):
    """
    Compute quaternion from two known vectors
    """
    quat = Quaternion()
    norms = np.sqrt(np.linalg.norm(vec_0)*np.linalg.norm(vec_1)) # |u||v|
    cp = np.cross(vec_0, vec_1) # cross product
    w = norms + np.dot(vec_0,vec_1)
    x, y, z = cp[0], cp[1], cp[2]
    q = np.array([x,y,z,w])
    norm_q = q/np.linalg.norm(q)
    quat.x = norm_q[0]
    quat.y = norm_q[1]
    quat.z = norm_q[2]
    quat.w = norm_q[3]

    return quat

def main():
    # instantiate iiwa
    iiwa = iiwaRobot()
    time.sleep(4) # allow iiwa taking some time to wake up
    # zero joints
    iiwa.move_joint(commit=True)
    # iiwa get ready
    iiwa.move_joint(JOINT_PERCH, commit=True)
    rospy.loginfo("iiwa is ready")
    # Configure realsense D435 depth and color streams
    pipeline = rs.pipeline()
    config = rs.config()
    config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
    config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
    profile = pipeline.start(config)
    # Create an align object
    align_to = rs.stream.color
    align = rs.align(align_to)
    # load siammask config
    cfg.merge_from_file(sys.path[0]+"/siammask_r50_l3/config.yaml")
    cfg.CUDA = torch.cuda.is_available()
    device = torch.device('cuda' if cfg.CUDA else 'cpu')
    # create model
    model = ModelBuilder()
    # load model
    model.load_state_dict(torch.load(sys.path[0]+"/siammask_r50_l3/model.pth",
        map_location=lambda storage, loc: storage.cpu()))
    model.eval().to(device)
    # build tracker
    tracker = build_tracker(model)
    # label object
    video_name = 'D435_color'
    cv2.namedWindow(video_name, cv2.WND_PROP_FULLSCREEN)
    first_frame = True
    FIN_FLAG = False
    GOAL_SET_FLAG = False
    while not FIN_FLAG:
        # wait image stream and select object of interest
        frames = pipeline.wait_for_frames()
        # Align the depth frame to color frame
        aligned_frames = align.process(frames)
        color_frame = aligned_frames.get_color_frame()
        depth_frame = aligned_frames.get_depth_frame()
        depth_intrinsics = rs.video_stream_profile(depth_frame.profile).get_intrinsics()
        # convert image to numpy arrays
        if color_frame:
            color_image = np.asanyarray(color_frame.get_data())
            depth_image = np.asanyarray(depth_frame.get_data())
        if first_frame:
            try:
                init_rect = cv2.selectROI(video_name, color_image, False, False)
            except:
                exit()
            tracker.init(color_image, init_rect)
            first_frame = False
        else:
            # start tracking
            outputs = tracker.track(color_image)
            polygon = np.array(outputs['polygon']).astype(np.int32)
            cv2.polylines(color_image, [polygon.reshape((-1, 1, 2))],
                          True, (0, 255, 0), 3)
            mask = ((outputs['mask'] > cfg.TRACK.MASK_THERSHOLD) * 255)
            mask = mask.astype(np.uint8)
            mask = np.stack([mask, mask*255, mask]).transpose(1, 2, 0)
            color_image = cv2.addWeighted(color_image, 0.77, mask, 0.23, -1)
            bbox = list(map(int, outputs['bbox']))
            poi_pixel = [int(bbox[0]+0.5*bbox[2]), int(bbox[1]+0.5*bbox[3])]
            poi_depth = depth_frame.get_distance(poi_pixel[0], poi_pixel[1])
            poi_rs = rs.rs2_deproject_pixel_to_point(depth_intrinsics, poi_pixel, poi_depth)
            print("Object 3D position w.r.t. camera frame: {}".format(poi_rs))
            if not np.allclose(poi_rs, np.zeros(3)):
                # compute transformed position of poi w.r.t. iiwa_link_0
                transfrom = iiwa.tf_listener.getLatestCommonTime('/iiwa_link_0', '/rs_d435')
                pos_rs = PoseStamped()
                pos_rs.header.frame_id = 'rs_d435'
                pos_rs.pose.orientation.w = 1.
                pos_rs.pose.position.x = poi_rs[0]
                pos_rs.pose.position.y = poi_rs[1]
                pos_rs.pose.position.z = poi_rs[2]
                pos_iiwa = iiwa.tf_listener.transformPose('/iiwa_link_0', pos_rs)
                rospy.loginfo("Object 3D position w.r.t. iiwa base from: {}\n ee w.r.t. iiwa base: {}".format(pos_iiwa.pose.position, iiwa.cartesian_pose.position))
                # vec_ee_poi = np.array([pos_iiwa.pose.position.x,       pos_iiwa.pose.position.y,pos_iiwa.pose.position.z]) - np.array([iiwa.cartesian_pose.position.x,iiwa.cartesian_pose.position.y,iiwa.cartesian_pose.position.z])
                # goal_pos = np.array([pos_iiwa.pose.position.x,       pos_iiwa.pose.position.y,pos_iiwa.pose.position.z]) - vec_ee_poi/np.linalg.norm(vec_ee_poi)*0.167
                # # compute orientation w.r.t. iiwa_link_0
                # goal_quat = quat_from_vecs(vec_ee_poi, np.array([0,0,1]))
                # set cartesian goal
                iiwa.goal_carte_pose.pose.position.x = pos_iiwa.pose.position.x
                iiwa.goal_carte_pose.pose.position.y = pos_iiwa.pose.position.y
                iiwa.goal_carte_pose.pose.position.z = pos_iiwa.pose.position.z
                # iiwa.goal_carte_pose.pose.position.x = goal_pos[0]
                # iiwa.goal_carte_pose.pose.position.y = goal_pos[1]
                # iiwa.goal_carte_pose.pose.position.z = goal_pos[2]
                # iiwa.goal_carte_pose.pose.orientation = goal_quat
                # iiwa.goal_carte_pose.pose.orientation.x = goal_quat[0]
                # iiwa.goal_carte_pose.pose.orientation.y = goal_quat[1]
                # iiwa.goal_carte_pose.pose.orientation.z = goal_quat[2]
                # iiwa.goal_carte_pose.pose.orientation.w = goal_quat[3]
                iiwa.goal_carte_pose.pose.orientation.x = iiwa.cartesian_pose.orientation.x
                iiwa.goal_carte_pose.pose.orientation.y = iiwa.cartesian_pose.orientation.y
                iiwa.goal_carte_pose.pose.orientation.z = iiwa.cartesian_pose.orientation.z
                iiwa.goal_carte_pose.pose.orientation.w = iiwa.cartesian_pose.orientation.w
                iiwa.goal_carte_pose.header.frame_id = 'iiwa_link_0'
                FIN_FLAG = True
                GOAL_SET_FLAG = True

        # display image stream, press 'ESC' or 'q' to terminate
        cv2.imshow(video_name, color_image)
        key = cv2.waitKey(40)
        if key in (27, ord("q")):
            break

    iiwa.move_cartesian(cartesian_pose=iiwa.goal_carte_pose, commit=True)
    time.sleep(1)
    iiwa.move_joint(joint_position=JOINT_PERCH)
    pipeline.stop()
    rospy.loginfo("Finished")

if __name__ == '__main__':
    try:
        main()
    except rospy.ROSInterruptException:
        pass