from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
from __future__ import print_function

import torch as th
from torch.utils.data import Dataset
import torch.nn.functional as F
import pandas as pd
import os
import numpy as np
import re
import random
import librosa

class Youtube_DataLoader(Dataset):
    """Youtube dataset loader."""

    def __init__(
            self,
            csv,
            features_path,
            features_path_audio,
            caption,
            we,
            min_time=10.0,
            feature_framerate=1.0,
            feature_framerate_3D=24.0 / 16.0,
            we_dim=300,
            max_words=30,
            min_words=0,
            n_pair=1,
            num_audio_frames=1024,
            random_audio_windows=False,
    ):
        """
        Args:
        """
        self.csv = pd.read_csv(csv)
        self.features_path = features_path
        self.features_path_audio = features_path_audio if features_path_audio != "" \
                                   else features_path
        self.caption = caption
        self.min_time = min_time
        self.feature_framerate = feature_framerate
        self.feature_framerate_3D = feature_framerate_3D
        self.we_dim = we_dim
        self.max_words = max_words
        self.min_words = min_words
        self.num_audio_frames = num_audio_frames
        self.we = we
        self.n_pair = n_pair
        self.fps = {'2d': feature_framerate, '3d': feature_framerate_3D}
        self.feature_path = {'2d': features_path}
        if features_path != '':
            self.feature_path['3d'] = features_path
        self.random_audio_windows = random_audio_windows

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

    def _zero_pad_tensor(self, tensor, size):
        if len(tensor) >= size:
            return tensor[:size]
        else:
            zero = np.zeros((size - len(tensor), self.we_dim), dtype=np.float32)
            return np.concatenate((tensor, zero), axis=0)

    def _zero_pad_audio(self, audio, max_frames):
        n_frames = audio.shape[1]
        if n_frames >= max_frames:
            return audio[:, 0:max_frames], int(max_frames)
        else:
            p = max_frames - n_frames
            audio_padded = np.pad(audio, ((0, 0),(0, p)), 'constant', constant_values=(0, 0))
            return audio_padded, n_frames

    def _tokenize_text(self, sentence):
        w = re.findall(r"[\w']+", str(sentence))
        return w

    def _words_to_we(self, words):
        words = [word for word in words if word in self.we.vocab]
        if words:
            we = self._zero_pad_tensor(self.we[words], self.max_words)
            return th.from_numpy(we)
        else:
            return th.zeros(self.max_words, self.we_dim)

    def _get_audio_and_text(self, caption, n_pair_max, mel_spec):
        n_caption = len(caption['start'])
        k = n_pair_max
        starts = np.zeros(k)
        ends = np.zeros(k)
        text = th.zeros(k, self.max_words, self.we_dim)
        audio = [0 for i in range(k)]
        nframes = np.zeros(k)
        r_ind = np.random.choice(range(n_caption), k, replace=True)

        for i in range(k):
            ind = r_ind[i]
            audio[i], nframes[i], starts[i], ends[i], text[i] = self._get_single_audio_text(caption, ind, mel_spec)

        audio = th.cat([i.unsqueeze(0) for i in audio], dim=0)
        return audio, nframes, starts, ends, text

    def _get_single_audio_text(self, caption, ind, mel_spec):
        start, end = ind, ind
        words = self._tokenize_text(caption['text'][ind])
        diff = caption['end'][end] - caption['start'][start]
        # Extend the video clip if shorter than the minimum desired clip duration
        while diff < self.min_time:
            if start > 0 and end < len(caption['end']) - 1:
                next_words = self._tokenize_text(caption['text'][end + 1])
                prev_words = self._tokenize_text(caption['text'][start - 1])
                d1 = caption['end'][end + 1] - caption['start'][start]
                d2 = caption['end'][end] - caption['start'][start - 1]
                # Use the closest neighboring video clip
                if d2 <= d1:
                    start -= 1
                    words.extend(prev_words)
                else:
                    end += 1
                    words.extend(next_words)
            # If no video clips after it, use the clip before it
            elif start > 0:
                words.extend(self._tokenize_text(caption['text'][start - 1]))
                start -= 1
             # If no video clips before it, use the clip after it.
            elif end < len(caption['end']) - 1:
                words.extend(self._tokenize_text(caption['text'][end + 1]))
                end += 1
            # If there's no clips before or after
            else:
                break
            diff = caption['end'][end] - caption['start'][start]

        frames = librosa.core.time_to_frames([caption['start'][start], caption['end'][end]], sr=16000, hop_length=160, n_fft=400)
        padded_mel_spec, nframes = self._zero_pad_audio(mel_spec[:, frames[0]: frames[1]], self.num_audio_frames)
        return th.from_numpy(padded_mel_spec), nframes, caption['start'][start], caption['end'][end], self._words_to_we(words)

    def _get_audio_random(self, n_pair_max, mel_spec):
        k = n_pair_max
        starts = np.zeros(k)
        ends = np.zeros(k)
        audio = [0 for i in range(k)]
        nframes = np.zeros(k)
        video_duration_seconds = int(librosa.core.frames_to_time(mel_spec.shape[1], sr=16000, hop_length=160, n_fft=400))
        num_audio_seconds = int(librosa.core.frames_to_time(self.num_audio_frames, sr=16000, hop_length=160, n_fft=400))
        # Sample clips that end before the end of the video
        # If the video is shorter than the desired window, use the entire video
        start_seconds = np.random.choice(range(max(1, video_duration_seconds - (num_audio_seconds + 1))), k, replace=True)

        for i in range(k):
            start_frame = max(0, librosa.core.time_to_frames(start_seconds[i], sr=16000, hop_length=160, n_fft=400))
            audio_window = mel_spec[:, start_frame : start_frame + self.num_audio_frames]
            # Pad in the case that the audio wasn't long enough
            padded_mel_spec, nframes_spec = self._zero_pad_audio(audio_window, self.num_audio_frames)
            end_second = start_seconds[i] + num_audio_seconds
            audio[i], nframes[i], starts[i], ends[i] = th.from_numpy(padded_mel_spec), nframes_spec, start_seconds[i], end_second

        audio = th.cat([i.unsqueeze(0) for i in audio], dim=0)
        return audio, nframes, starts, ends

    def _get_video(self, vid_path, s, e, video_id):
        feature_path = {}
        video = {}
        output = {}
        for k in self.feature_path:
            feature_path[k] = os.path.join(self.feature_path[k], vid_path, video_id + "_{}.npz".format(k))
            np_arr = np.load(feature_path[k])['features']
            video[k] = th.from_numpy(np_arr).float()
            output[k] = th.zeros(len(s), video[k].shape[-1])
            for i in range(len(s)):
                start = int(s[i] * self.fps[k])
                end = int(e[i] * self.fps[k]) + 1
                slice = video[k][start:end]
                if len(slice) < 1:
                    print("missing visual feats; video_id: {}, start: {}, end: {}".format(
                        feature_path[k], start, end))
                else:
                    output[k][i] = F.normalize(th.max(slice, dim=0)[0], dim=0)

        return th.cat([output[k] for k in output], dim=1)

    def __getitem__(self, idx):
        vid_path = self.csv['path'].values[idx].replace("None/", "")
        video_id = vid_path.split("/")[-1]
        audio_path = os.path.join(self.features_path_audio, vid_path, video_id + "_spec.npz")
        mel_spec = np.load(audio_path)['arr_0']
        if self.random_audio_windows:
            audio, nframes, starts, ends = self._get_audio_random(self.n_pair, mel_spec)
        else:
            audio, nframes, starts, ends, text = self._get_audio_and_text(self.caption[video_id], self.n_pair, mel_spec)
        video = self._get_video(vid_path, starts, ends, video_id)
        if self.random_audio_windows:
            return {'video': video, 'audio':th.HalfTensor(audio), 'nframes': th.IntTensor(nframes), 'video_id': video_id}
        else:
            return {'video': video, 'audio':th.HalfTensor(audio), 'nframes': th.IntTensor(nframes), 'video_id': video_id,
                    'text': text}


import os
import coco
import sys
import random
import math
import re
import time
import numpy as np
import cv2
import matplotlib
import matplotlib.pyplot as plt
import skimage.io
import json
from PIL import Image
from config import Config
import utils
import model as modellib
import visualize
from model import log

# Root directory of the project
ROOT_DIR = os.getcwd()

# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")

# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")

# local path to mapillary dataset
DATASET_DIR = os.path.join(ROOT_DIR, "mapillary_dataset")

# Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)

class MapillaryConfig(coco.CocoConfig):
    """Configuration for training on the mapillary dataset.
    Derives from the base Config class and overrides values specific
    to the mapillary dataset.
    """
    # Give the configuration a recognizable name
    NAME = "debug"

    # Train on 1 GPU and 8 images per GPU. We can put multiple images on each
    # GPU because the images are small. Batch size is 8 (GPUs * images/GPU).
    GPU_COUNT = 1
    IMAGES_PER_GPU = 8

    # Number of classes (including background)
    # this MUST be explicitly defined, or will run into index out of bound error
    NUM_CLASSES = 1 + 11  # background + 11 objects

    # Use small images for faster training. Set the limits of the small side
    # the large side, and that determines the image shape.
#     IMAGE_MIN_DIM = 1024
#     IMAGE_MAX_DIM = 1024

    # Use smaller anchors because our image and objects are small
    # RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)  # anchor side in pixels

    # Reduce training ROIs per image because the images are small and have
    # few objects. Aim to allow ROI sampling to pick 33% positive ROIs.
    TRAIN_ROIS_PER_IMAGE = 32

    # Use a small epoch since the data is simple
    STEPS_PER_EPOCH = 1

    # use small validation steps since the epoch is small
    VALIDATION_STEPS = 50

    LEARNING_RATE = 0.0001
    WEIGHT_DECAY = 0.0005

config = MapillaryConfig()
config.display()

class MapillaryDataset(utils.Dataset):
    """Generates the shapes synthetic dataset. The dataset consists of simple
    shapes (triangles, squares, circles) placed randomly on a blank surface.
    The images are generated on the fly. No file access required.
    """

    DEBUG = False
    CLASS_MAP = {}

    CLASSES = ["Bird", "Person", "Bicyclist", "Motorcyclist", "Bench", \
         "Fire Hydrant", "Traffic Light", "Bus", "Car", "Motorcycle", "Truck"]

    # local path to image folder, choose 'dev', 'training', or 'testing'
    SUBSET_DIR = ""

    # local path to images inside development folder
    IMG_DIR = ""

    # local path to instance annotations inside development folder
    INS_DIR = ""

    def load_mapillary(self, dataset_dir, subset, class_ids=None,
                  class_map=None):

        self.SUBSET_DIR = os.path.join(dataset_dir, subset)
        self.IMG_DIR = os.path.join(self.SUBSET_DIR, 'images')
        self.INS_DIR = os.path.join(self.SUBSET_DIR, 'instances')

        # load classes, start with id = 1 to account for background "BG"
        class_id = 1
        for label_id, label in enumerate(class_ids):
            if label["instances"] == True and label["readable"] in self.CLASSES:
                self.CLASS_MAP[label_id] = class_id

                if (self.DEBUG):
                    print("{}: Class {} {} added".format(label_id, class_id, label["readable"]))

                self.add_class("mapillary", class_id, label["readable"])
                class_id = class_id + 1

        # add images
        file_names = next(os.walk(self.IMG_DIR))[2]
        for i in range(len(file_names)):
            if file_names[i] != '.DS_Store':

                image_path = os.path.join(self.IMG_DIR, file_names[i])
                base_image = Image.open(image_path)
                w, h = base_image.size

                if (self.DEBUG):
                    print("Image {} {} x {} added".format(file_names[i], w, h))

                self.add_image("mapillary", image_id = i,
                              path = file_names[i],
                               width = w,
                               height = h
                              )

    def load_image(self, image_id):
        """Generate an image from the specs of the given image ID.
        This function loads the image from a file.
        """
        info = self.image_info[image_id]
        img_path = os.path.join(self.IMG_DIR, info["path"])
        image = Image.open(img_path)
        image_array = np.array(image)
        return image_array

    def image_reference(self, image_id):
        """Return the local directory path of the image."""
        info = self.image_info[image_id]
        img_path = os.path.join(self.IMG_DIR, info["path"])
        return img_path

    def load_mask(self, image_id):
        """Generate instance masks for shapes of the given image ID.
        """
        info = self.image_info[image_id]
        instance_path = os.path.join(self.INS_DIR, info["path"])
        instance_image = Image.open(instance_path.rsplit(".", 1)[0] + ".png")

        # convert labeled data to numpy arrays for better handling
        instance_array = np.array(instance_image, dtype=np.uint16)

        instances = np.unique(instance_array)
        instaces_count = instances.shape[0]

        label_ids = instances // 256
        label_id_count = np.unique(label_ids).shape[0]

        if (self.DEBUG):
            print("There are {} instances, {} classes labelled instances in the image {}."\
                  .format(instaces_count, label_id_count, info["path"]))

        mask = np.zeros([instance_array.shape[0], instance_array.shape[1], instaces_count], dtype=np.uint8)
        mask_count = 0
        loaded_class_ids = []
        for instance in instances:
            label_id = instance // 256
            if (label_id in self.CLASS_MAP):
                m = np.zeros((instance_array.shape[0], instance_array.shape[1]), dtype=np.uint8)
                m[instance_array == instance] = 1
                m_size = np.count_nonzero(m == 1)

                # only load mask greater than threshold size,
                # otherwise bounding box with area zero causes program to crash
                if m_size > 5000:
                    mask[:, :, mask_count] = m
                    loaded_class_ids.append(self.CLASS_MAP[label_id])
                    mask_count = mask_count + 1
                    if (self.DEBUG):
                        print("Non-zero: {}".format(m_size))
                        print("Mask {} created for instance {} of class {} {}"\
                              .format(mask_count, instance, self.CLASS_MAP[label_id], \
                                      self.class_names[self.CLASS_MAP[label_id]]))
        mask = mask[:, :, 0:mask_count]
        return mask, np.array(loaded_class_ids)

# read in config file
with open(os.path.join(DATASET_DIR, 'config.json')) as config_file:
    class_config = json.load(config_file)
# in this example we are only interested in the labels
labels = class_config['labels']

# Training dataset
dataset_train = MapillaryDataset()
dataset_train.load_mapillary(DATASET_DIR, "debug_train", class_ids = labels)
dataset_train.prepare()

# Validation dataset
dataset_val = MapillaryDataset()
dataset_val.load_mapillary(DATASET_DIR, "debug_val", class_ids = labels)
dataset_val.prepare()

print("mapping: ", class_config["mapping"])
print("version: ", class_config["version"])
print("folder_structure:", class_config["folder_structure"])
print("There are {} classes in the config file".format(len(labels)))
print("There are {} classes in the model".format(len(dataset_train.class_names)))
for i in range(len(dataset_train.class_names)):
    print("    Class {}: {}".format(i, dataset_train.class_names[i]))

# Create model in training mode
model = modellib.MaskRCNN(mode="training", config=config,
                          model_dir=MODEL_DIR)

# Which weights to start with?
init_with = "coco"  # imagenet, coco, or last

if init_with == "imagenet":
    model.load_weights(model.get_imagenet_weights(), by_name=True)

elif init_with == "coco":
    # Load weights trained on MS COCO, but skip layers that
    # are different due to the different number of classes
    # See README for instructions to download the COCO weights
    model.load_weights(COCO_MODEL_PATH, by_name=True,
                       exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",
                                "mrcnn_bbox", "mrcnn_mask"])
elif init_with == "last":
    # Load the last model you trained and continue training
    model.load_weights(model.find_last()[1], by_name=True)

# Train the head branches
# Passing layers="heads" freezes all layers except the head
# layers. You can also pass a regular expression to select
# which layers to train by name pattern.
model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE,
            epochs=2,
            layers='heads')


import numpy as np
import os
import shutil
import random
import torch
import sklearn.utils

from vae.model import VAE
# from gen_model.model import AutoDropoutNN
from loader.dataloader import MIDI_Loader,MIDI_Render
from loader.chordloader import Chord_Loader
from torch.nn import functional as F

new_outputs= np.load("model_v7_midi.npy")

test_data_real = np.load("test_data_no_split_8_crop.npy")
render = MIDI_Render(datasetName = "Nottingham", minStep= 0.125)

# vae_model = VAE(130, 2048, 3, 12, 128, 128, 32)
# vae_model.eval()
# dic = torch.load("vae/tr_chord.pt")
# for name in list(dic.keys()):
#     dic[name.replace('module.', '')] = dic.pop(name)
# vae_model.load_state_dict(dic)
# if torch.cuda.is_available():
#     vae_model = vae_model.cuda()
for kkk in range(len(new_outputs)):
    # print(kkk)
    # z1 = new_outputs[kkk][0][:,:128]
    # z2 = new_outputs[kkk][0][:,128:]
    # chord_cond = new_outputs[kkk][1]
    # # print(z1.shape)
    # # print(z2.shape)
    # # print(chord_cond.shape)
    # z1 = torch.from_numpy(z1).float()
    # z2 = torch.from_numpy(z2).float()
    # chord_cond = torch.from_numpy(chord_cond).float()
    # # z1 = z1.cuda()
    # # z2 = z2.cuda()
    # # chord_cond = chord_cond.cuda()
    # res = vae_model.decoder(z1,z2,chord_cond)
    # # print(res.detach().cpu().numpy().shape)
    # q = np.asarray((res.detach().cpu().numpy()))
    # temp = []
    # for i in q:
    #     for j in i:
    #         temp.append(np.argmax(j))
    # q = np.array(temp)

    # song = {"notes":q, "chords": test_data_real[kkk]["chords"]}
    # render.data2midi(data = new_outputs[kkk],output = "results/o/o_" + str(kkk) + ".mid")
    render.data2midi(data = test_data_real[kkk],output = "results/o/o_" + str(kkk) + ".mid")


from numpy.core.defchararray import join
import pybullet as p
import numpy as np
import os
from gym_cricket.assests.cricket_abs import Cricket_abs

class Cricket(Cricket_abs):
    def __init__(self, client, strating_position=[], base_position = [0,0,0.5], normal_forces = 4) -> None:
        """
        Input:
        - clinet : costant
        PyBullet phisic client
        - starting_position : np.array
        starting position for the joints
        - base_position :
         [x,y,z] position of the robot at the beginning of the simulation

        variables:
        self.cricket:
            cricket_id from URDF file
        self.track_joints and self.wheel_ids:
            joints related to the tracks --> list of lists [[wheel1,wheel2],[..],...]
        self.limb_joints,self.limb_ids:
            all the other non-fixed joints (knees, shoulders, and so on)
        """
        self.client = client
        f_path='urdfs/cricket_robot.urdf'
        f_name = os.path.join(os.path.dirname(__file__), f_path)
        self.cricket = p.loadURDF(fileName = f_name,
                                  basePosition = base_position,
                                  physicsClientId=self.client
                                  )
                                  #flags=p.URDF_USE_SELF_COLLISION | p.URDF_USE_SELF_COLLISION_INCLUDE_PARENT)
        self.track_joints,self.wheel_ids,\
            self.limb_joints,self.limb_ids,\
                self.fixed_joints, self.fixed_ids = self.__find_joints()   # completes the above joint lists

        self.num_wheel = 8
        num_tracks = len(self.track_joints)
        num_joints = len(self.limb_ids)

        #Starting positions
        if strating_position == []:
            self.wheel_position = np.zeros((self.num_wheel,))
            self.limb_positions = np.zeros((num_joints))
        else :
            self.wheel_position = np.array(strating_position[:self.num_wheel])
            self.limb_positions = np.array(strating_position[self.num_wheel:])

        for joint_id, joint_pos in zip(self.limb_ids, self.limb_positions):
            p.resetJointState(
                self.cricket,
                joint_id,
                joint_pos,
                0,
                self.client
            )

        self.mass = 100
        # Starting velocities 0
        self.track_velocities = np.zeros(self.num_wheel)
        self.limb_velocities = np.zeros(len(self.limb_positions))

        # Max and Min linear velocity - vx,vy,vz
        velocity = 100
        self.max_lvel = np.array([velocity,velocity,velocity])
        self.min_lvel = np.array([0,0,0])
        # Max and Min angular velocity
        self.max_avel = np.array([velocity,velocity,velocity])
        self.min_avel = np.array([0,0,0])

        # Normal forces
        self.n_normal_f = normal_forces

        # Store all the collision between links and their fixed accessories (which are considered collision by PyBullet)
        self.collision_safe = {}
        for i in self.limb_ids:
            aa, bb = p.getAABB(self.cricket,i,self.client) # return the bounding box of the body (-1) starting from the center of mass
            obs = p.getOverlappingObjects(aa, bb, self.client)
            self.collision_safe.update({i : sorted(obs)})

    # def set_joint_position(self, positions):
    #     for i in range(p.getNumJoints(self.cricket)):
    #         p.resetJointState(self.cricket, i, positions[i], physicsClientId = self.client)

    def perform_action(self, action):
        '''
        The action passed to the robot are the angles of the joints and the rotation of
        the wheel tracks

        [angle1, angle2, ..., anglen, speed1, speed2, ..., anglen]

        action:
            0-3 the value of the tracks movement
            4... the tortion of all the other joints
        '''
        half = int(len(action)/2)
        action = np.array([[a,b] for a,b in zip(action[:half],action[half:])])
        # action -> [(angle1, speed1), (angle2, speed2), ..., (anglen,speedn)]
        # Tracks movement
        wheel_angles = [val for val in action[:4,0] for _ in (0,1)] # take all the first values of the first 4 pairs two time --> e.g. [[1,2],[3,4],[5,6],[7,8]] >> became >> [1,1,3,3,5,5,7,7]
        wheel_velocities = [val for val in action[:4,1] for _ in (0,1)]
        self.wheel_position = np.add(self.wheel_position, wheel_angles)
        self.track_velocities = np.add(self.track_velocities, wheel_velocities)
        p.setJointMotorControlArray(
            self.cricket,
            jointIndices = self.wheel_ids,
            controlMode = p.POSITION_CONTROL,
            targetPositions = self.wheel_position,
            targetVelocities = self.track_velocities,
            physicsClientId = self.client
            )

        # Limbs movement
        limb_angles = [val for val in action[4:,0]]
        limb_velocites = [val for val in action[4:,1]]
        self.limb_positions = np.add(self.limb_positions, limb_angles)
        self.limb_velocities = np.add(self.limb_velocities, limb_velocites)
        p.setJointMotorControlArray(
            self.cricket,
            jointIndices = self.limb_ids,
            controlMode = p.POSITION_CONTROL,
            targetPositions = self.limb_positions,
            targetVelocities = self.limb_velocities,
            physicsClientId = self.client
            )

    # def get_observations(self):
    #     '''Return:
    #         - position (x,y,z)
    #         - orientation (roll around X, pitch around Y, yaw around Z)
    #         - velocity (two 3 values vetors: linear velocity [x,y,z], angular velocity [ωx,ωy,ωz])
    #        of the robot in the simulation'''
    #     pos, ang = p.getBasePositionAndOrientation(self.cricket, self.client)

    #     # convert the quaternion (δ,Θ,ψ,ω) to Euler (δ,Θ,ψ)
    #     angs = p.getEulerFromQuaternion(ang, physicsClientId=self.client) # roll, pitch, yaw

    #     # Get the velocity of the robot
    #     l_vel, a_vel = p.getBaseVelocity(self.cricket, self.client) # linear & angular velocity

    #     return pos,angs,l_vel,a_vel

    # def get_ids(self):
    #     '''Return the basic PyBullet information of the robot

    #     Return:
    #      - ID
    #      - Client
    #     '''
    #     return self.cricket, self.client

    # def get_joint_ids(self):
    #     '''Return the joints indexes:
    #      - track
    #      - limb
    #      - fixed
    #     '''
    #     return self.wheel_ids,\
    #         self.limb_ids, self.fixed_ids

    def __find_joints(self):
        '''Completes the joints in the cricket robot'''
        number_of_joints = p.getNumJoints(self.cricket, physicsClientId=self.client)
        track = []
        track_joints, limb_joints, fixed_joints = [], [], []
        wheel_ids, limb_ids, fixed_ids = [], [], []
        for joint_number in range(number_of_joints):
            joint_info = p.getJointInfo(self.cricket, joint_number, physicsClientId=self.client)
            # [jointIndex, jointName (bytes), jointType, ...]
            if joint_info[2] == 0 and "track" in joint_info[1].decode("utf-8"): # if the joint is revolute/continous and contains the word "track"
                # the joint is one of the wheels
                track.append(joint_info)
                wheel_ids.append(joint_info[0])
                if len(track) == 2:
                    track_joints.append(track)
                    track = []
            elif joint_info[2] == 0:
                # the joint is revolute/continuos
                limb_ids.append(joint_info[0])
                limb_joints.append(joint_info)
            else :
                fixed_ids.append(joint_info[0])
                fixed_joints.append(joint_info)
        return np.array(track_joints, dtype=object), np.array(wheel_ids, dtype=object),\
                np.array(limb_joints, dtype=object), np.array(limb_ids, dtype=object),\
                 np.array(fixed_joints, dtype=object), np.array(fixed_ids, dtype=object)

    def get_joint_positions(self):
        """
         - getJointState(robot_id, jointIndex)
        [*jointPosition*, jointVelocity, jointReactionForces, appliedJointMotorTorque]
        """
        track_pos = [p.getJointState(self.cricket, wheel[0], physicsClientId=self.client)[0] # wheel[0] gets the jointIndex
            for track in self.track_joints for wheel in track]
        track_pos = self.__normalize(track_pos)

        limb_pos = [p.getJointState(self.cricket, joint[0], physicsClientId=self.client)[0]
            for joint in self.limb_joints]
        limb_pos = self.__normalize(limb_pos)
        # should I return two dictionaries? --> dict(zip(list1,list2))
        return track_pos, limb_pos

    def __normalize(self, position_list):
        '''Helps get_joint_position to normalize the returned values'''
        res = []
        for position in position_list:
            if position % (2*np.pi) > np.pi or position == (2*np.pi) :
                res.append(position%np.pi-np.pi )
            else :
                res.append(position%np.pi)
        return res

    # def get_joint_velocities(self):
    #     """
    #      - getJointState(robot_id, jointIndex)
    #     [jointPosition, *jointVelocity*, jointReactionForces, appliedJointMotorTorque]
    #     """
    #     track_vel = [p.getJointState(self.cricket, wheel[0])[1]
    #         for track in self.track_joints for wheel in track]
    #     limb_vel = [p.getJointState(self.cricket, joint[0])[1]
    #         for limb in self.limb_joints for joint in limb]
    #     return track_vel, limb_vel

    # def get_joint_collisions(self):
    #     collision = {}
    #     for id in self.limb_ids:
    #         aa, bb = p.getAABB(self.cricket,id,self.client) # return the bounding box of the body (-1) starting from the center of mass
    #         obs = p.getOverlappingObjects(aa, bb, self.client)
    #         safe_obs = self.collision_safe.get(id)
    #         if sorted(obs) not in safe_obs :
    #             collision.update({id : list(set(obs) - set(safe_obs))})
    #     return collision

    def get_normal_forces(self, planeId : str):
        """Get all the normal forces between the robot and palneId

        If it's not possible to obtain 4 normal forces an array full of zeros will complete it
        if there are more than 4 normal forces per track, it will choose randomly

        Return:
         - a list of 4 lists (one per track) of normal forces, for each list:
            4 normal forces along the wheels and the track between the wheels
        """
        contact_points = []
        for count, wheel_id in enumerate(self.wheel_ids):
            contact_points += [c_point[9] for c_point in p.getContactPoints(self.cricket, planeId, linkIndexA=wheel_id)]

            if count + 1 % 2 != 0 : # update trackId once each two iterations
                track_id = wheel_id + 1
                contact_points += [c_point[9] for c_point in p.getContactPoints(self.cricket, planeId, linkIndexA=track_id)]

        if len(contact_points) < self.n_normal_f :
            contact_points += [0.0] * (self.n_normal_f - len(contact_points))
        elif len(contact_points) > self.n_normal_f :
            contact_points = sorted(contact_points, reverse=True)[:self.n_normal_f]

        return contact_points

    def get_joint_limits(self):
        high_lim, low_lim = [] ,[]
        for joint in self.limb_joints:
            if joint[8] == 0 : # continous joint
                high_lim.append(np.pi)
                low_lim.append(-np.pi)
            else:
                high_lim.append(np.pi/2)
                low_lim.append(-np.pi/2)
        return np.array(high_lim), np.array(low_lim)

    def get_track_limits(self):
        high_lim, low_lim = np.full((self.num_wheel,), np.pi), np.full((self.num_wheel,),-np.pi)
        return high_lim,low_lim

    # def get_normal_forces_limits(self,gravity):
    #     mass = 10000 # grams --> there's a way to get the mass
    #     max_nf, min_nf = np.full((self.n_normal_f,),mass*gravity), np.zeros((self.n_normal_f))
    #     return max_nf,min_nf

    # def get_action_limits(self):
    #     dim = len(self.limb_ids) + len(self.track_joints)
    #     high_lim = np.full((dim,), np.pi)
    #     low_lim = -high_lim
    #     return high_lim,low_lim

    def get_action_velocities_limits(self):
        dim = len(self.limb_ids) + len(self.track_joints)
        high_lim = np.full((dim,), 100)
        low_lim = np.zeros((dim,))
        return high_lim,low_lim

    # def print_info(self):
    #     '''
    #     Prints the main robot characteristics
    #     '''
    #     times = 40
    #     print('='*times)
    #     print(f'ID: {self.cricket}')
    #     print(f'clinet: {self.client}')
    #     print(f'Body mass: {p.getDynamicsInfo(self.cricket,-1,self.client)}')
    #     print('tracks: ')
    #     for count, track in enumerate(self.track_joints):
    #         print(f'{count} - {track}')
    #         print('-'*times)
    #     print('_'*times)
    #     print('limb joints: ')
    #     for count, joint in enumerate(self.limb_joints):
    #         print(f'{count} - {joint}')
    #         print('-'*times)
    #     print('_'*times)
    #     print('fixed joints: ')
    #     for count, joint in enumerate(self.fixed_joints):
    #         print(f'{count} - {joint}')
    #         print('-'*times)
    #     print('_'*times)
    #     print('Link states')
    #     print('='*times)


import time
import cv2
import numpy as np
import os, sys
import pickle
import nep
import threading
import sharo

show_image = 1
try:
    print (sys.argv[1])
    show_image = int(sys.argv[1])
    print ("Show image: " + show_image)
except:
    pass

node = nep.node('face_detection')
sub_image = node.new_sub('robot_image', 'image')
pub_position = node.new_pub('face_positions', 'json')

perception_face = sharo.BooleanPerception(node, "human_detected", "value", 1, 3)  # --------------- Sharo ------------------
frame = ""
bounding_boxes = ""

def thread_function(name):
    global sub_image, frame
    while True:
        s, img = sub_image.listen()
        if s:
            frame = cv2.resize(img, (640,480), interpolation = cv2.INTER_AREA)

face_positions = threading.Thread(target=thread_function, args=(1,))
face_positions.start()

font = cv2.FONT_HERSHEY_COMPLEX_SMALL
model_filename =  'models/haarcascade_frontalface_default.xml'
faceCascade = cv2.CascadeClassifier(model_filename)
time.sleep(1)
people_detected_msg_sended = False
people_non_detected_msg_sended = False
n_detected = 0
non_detected = 0
start_detection =  time.time()
start_alone =  time.time()
print('Start!')

def draw_bounding_box(face_coordinates, image_array, color):
    x, y, w, h = face_coordinates
    cv2.rectangle(image_array, (x, y), (x + w, y + h), color, 2)

main_face = [0,0]
nb_faces = 0
timer_faces_yes = 0
timer_faces_no = 0

while True:
        newImage = frame.copy()
        close = 0
        gray = cv2.cvtColor(newImage, cv2.COLOR_BGR2GRAY)
        matrix = faceCascade.detectMultiScale( gray, scaleFactor=1.1, minNeighbors=5,minSize=(30, 30),)

        # n saves the number of face detected
        try:
            n,m =  matrix.shape
        except:
            n = 0

        if n > 0:
            #boxes = [[int(matrix[0,0]), int(matrix[0,1]), int(matrix[0,0]) + int(matrix[0,2]), int(matrix[0,1]) + int(matrix[0,3])]]
            # If only 1 face that is the main face
            if n == 1:
                x_m = int(matrix[close,0]) + int(matrix[close,2]/2)
                y_m = int(matrix[close,1]) + int(matrix[close,3]/2)
                main_face = [x_m, y_m] # Center of face
            # If more than 1 face track the closest one to the main face
            else:

                dist = 10000000
                close = 0 # Index of closest face
                for i in range(n):

                    x_m = int(matrix[i,0]) + int(matrix[i,2]/2)
                    y_m = int(matrix[i,1]) + int(matrix[i,3]/2)
                    dist_temp = ((main_face[0] - x_m)**2 + (main_face[1] - y_m)**2)**.5
                    if dist_temp < dist:
                        dist = dist_temp
                        close = i
                x_m = int(matrix[close,0]) + int(matrix[close,2]/2)
                y_m = int(matrix[close,1]) + int(matrix[close,3]/2)
                main_face = [x_m, y_m] # Center of face

            boxes = []
            for i in range(n):
                boxes.append([int(matrix[i,0]), int(matrix[i,1]), int(matrix[i,2]), int(matrix[i,3])])
            #Send only main face
            bounding_boxes = np.array(boxes)
            end = time.time()
            # Only send one face
            center = {"x":boxes[close][0], "y":boxes[close][1]}
            size = {"w":boxes[close][2] , "h": boxes[close][3]}
            box = {"x1" : int(matrix[close,0]), "x2" : int(matrix[close,1]),"y1" : int(matrix[close,0] + matrix[i,2]), "y2" : int(matrix[close,1] + matrix[i,3])}
            pub_position.publish({"face": {"center":center, "size":size,"box":box}})
            nb_faces = bounding_boxes.shape[0]

            if(nb_faces > 0):
                timer_faces_yes = timer_faces_yes + 1
                if (timer_faces_yes > 10):
                    timer_faces_yes = 0
                    perception_face.primitive_detected()      # --------------- Sharo ------------------
                    timer_faces_no = 0

            else:
                pub_position.publish({"positions": []})
                if (timer_faces_no > 10):
                    timer_faces_no = timer_faces_no + 1
                    perception_face.primitive_non_detected()  # --------------- Sharo ------------------
                    timer_faces_yes = 0
                    timer_faces_no = 0

        else:

            timer_faces_no = timer_faces_no + 1
            timer_faces_yes = 0
            if (timer_faces_no > 10):
                timer_faces_no = 0
                perception_face.primitive_non_detected()  # --------------- Sharo ------------------
            nb_faces = 0

        if show_image == 1:
            if nb_faces > 0:
                for i in range(nb_faces):
                    if i == close:
                        draw_bounding_box(bounding_boxes[i], newImage, (0, 255, 0))
                    else:
                        draw_bounding_box(bounding_boxes[i], newImage, (0, 0, 200))
            cv2.imshow("Face Detection", newImage)
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break


import torch
import numpy as np
from .. import load_tensor_and_image_from_file, load_image_buffer_to_tensor
from ..params import COCO_INSTANCE_CATEGORY_NAMES as COCO

recognition_nets = ['deeplabv3_resnet101', 'fcn_resnet101']
detection_nets = [
    'alexnet', 'densenet121', 'densenet161', 'densenet169',
    'densenet201', 'googlenet', 'inception_v3', 'mobilenet_v2', 'resnet101',
    'resnet152', 'resnet18', 'resnet34', 'resnet50', 'resnext101_32x8d',
    'resnext50_32x4d', 'shufflenet_v2_x0_5', 'shufflenet_v2_x1_0',
    'squeezenet1_0', 'squeezenet1_1', 'vgg11', 'vgg11_bn', 'vgg13',
    'vgg13_bn', 'vgg16', 'vgg16_bn', 'vgg19', 'vgg19_bn', 'wide_resnet101_2',
    'wide_resnet50_2'
]

def get_available_models_names(force_reload=False):
    """Gives list of model names available on torch hub

    Return:
        models (list[string]): list of pytorch/vision models available
                                on pytorch hub
    """
    return torch.hub.list('pytorch/vision', force_reload=force_reload)

def check_if_new_models_in_hub():
    """Checks if there are new models available on torch hub

    Return:
        boolean: True if new models are available, False otherwise
    """
    hub_list = get_available_models_names(False)
    local_list = recognition_nets + detection_nets
    if len(set(hub_list) & set(local_list)) == len(local_list):
        for i in local_list:
            if i not in hub_list:
                return True
        return False
    return True

class YoloVision:
    """Trained deep neuron network vision interface

    Attributes:
        _models: Pretrained net models torch hub instance name
    """

    _models = 'pytorch/vision',

    def __init__(self, device='cpu'):
        """YoloVisionRecognition constructor

        Args:
            device (string): Name of device used for calculation
        """
        if device not in ['cuda', 'cpu']:
            raise Exception('Wrong device type passed to constructor')
        self._device = device

class YoloVisionRecognition(YoloVision):
    """Trained deep neuron network vision recognition interface
    """

    def __init__(self, nn_model=None, device='cpu'):
        """YoloVisionRecognition constructor

        Args:
            nn_module (string): Deep net vision model
            device (string): Name of device used for calculation
        """
        super().__init__(device)
        if nn_model not in recognition_nets:
            raise Exception('Wrong net model for recognition')
        self._nn_module = torch.hub.load('pytorch/vision', nn_model,
                                         pretrained=True).eval()

    @staticmethod
    def find_boxes(output_predictions):
        """Sets boxed boundaries for predicted parameters

        Args:
            output_predictions (Tensor): AI model output predictions tensor

        Returns:
            prediction_boxed (object) Object with keys of predicted object
                                        names paired with boundary
                                        boxes of each prediction type
        """
        prediction_boxed = {}
        for (x, y), t in np.ndenumerate(output_predictions
                                        .byte()
                                        .cpu()
                                        .numpy()):
            for _, v in np.ndenumerate(t):
                value = v.item()
                if value in prediction_boxed.keys():
                    if prediction_boxed[value]['x_min'] > x:
                        prediction_boxed[value]['x_min'] = x
                    if prediction_boxed[value]['y_min'] > y:
                        prediction_boxed[value]['y_min'] = y
                    if prediction_boxed[value]['x_max'] < x:
                        prediction_boxed[value]['x_max'] = x
                    if prediction_boxed[value]['y_max'] < y:
                        prediction_boxed[value]['y_max'] = y
                else:
                    prediction_boxed[value] = {
                        'x_min': x or 0,
                        'y_min': y or 0,
                        'x_max': x,
                        'y_max': y
                    }
        named = {}
        for k in prediction_boxed.keys():
            named[str(COCO[k])] = prediction_boxed[k]
        return named

    def recognize_local_file(self, image_path):
        """Performs image recognition of local file

        Args:
            image_path (string): path to image

        Returns:
            tuple (input_image, output_predictions) Vectorized image
                                                    to numpy array and AI model
                                                    output predictions tensor
        """
        input_batch, input_image = load_tensor_and_image_from_file(
            image_path, self._device
        )
        with torch.no_grad():
            output = self._nn_module(input_batch)
        output = output['out'][0]
        output_predictions = output.argmax(0)
        return input_image, output_predictions

    def recognize_buffer(self, image_buf):
        """Performs image recognition of given image buffer

        Args:
            image_buf (bytes buffer): The image to recognize

        Returns:
            output_predictions: AI model output predictions tensor
        """
        input_batch = load_image_buffer_to_tensor(image_buf, self._device)
        with torch.no_grad():
            output = self._nn_module(input_batch)
        output = output['out'][0]
        output_predictions = output.argmax(0)
        return output_predictions

class YoloVisionDetection(YoloVision):
    """Trained deep neuron network vision recognition interface
    """

    def __init__(self, nn_model=None, device='cpu'):
        """YoloVisionRecognition constructor

        Args:
            nn_module (string): Deep net vision model
            device (string): Name of device used for calculation
        """
        super().__init__(device)
        if nn_model not in detection_nets:
            raise Exception('Wrong net model for detection')
        self._nn_module = torch.hub.load('pytorch/vision', nn_model,
                                         pretrained=True).eval()

    @staticmethod
    def find_most_probable(output_predictions):
        """Looks in to output predictions for most probable detection

        Args:
            output_predictions: Tensor with probabilities of Yolo detection

        Returns:
            tuple (float, int): tuple of recognition probability and number
                                representing Yolo class of the probability
        """
        max_probability = torch.max(output_predictions)
        yolo_class_index = (output_predictions ==
                            max_probability).nonzero().item()
        return max_probability.item(), yolo_class_index

    def detect_local_file(self, image_path):
        """Performs image recognition of local file

        Args:
            image_path (string): path to image

        Returns:
            output_predictions: AI model output probabilities tensor of
                                of Yolo classes
        """
        input_batch, input_image = load_tensor_and_image_from_file(
            image_path, self._device
        )
        with torch.no_grad():
            output = self._nn_module(input_batch)
        output_predictions = torch.nn.functional.softmax(output[0], dim=0)
        return output_predictions

    def detect_buffer(self, image_buf):
        """Performs object detection on given image buffer

        Args:
            image_buf (bytes buffer): The image to recognize

        Returns:
            output_predictions: AI model output probabilities tensor
                                of Yolo classes
        """
        input_batch = load_image_buffer_to_tensor(image_buf, self._device)
        with torch.no_grad():
            output = self._nn_module(input_batch)
        output_predictions = torch.nn.functional.softmax(output[0], dim=0)
        return output_predictions


import os
import random
from collections import namedtuple
import lmdb
import msgpack_numpy
import numpy as np
import torch
import torch.utils.data as data
from PIL import Image
import os.path as osp
import tqdm
from lilanet.datasets.transforms import Compose, RandomHorizontalFlip, Normalize

class KITTI(data.Dataset):
    """`KITTI LiDAR`_ Dataset.

    Args:
        root (string): Root directory of the ``lidar_2d`` and ``ImageSet`` folder.
        split (string, optional): Select the split to use, ``train``, ``val`` or ``all``
        transform (callable, optional): A function/transform that  takes in distance, reflectivity
            and target tensors and returns a transformed version.
    """

    Class = namedtuple('Class', ['name', 'id', 'color'])

    classes = [
        Class('unknown', 0, (0, 0, 0)),
        Class('car', 1, (0, 0, 142)),
        Class('pedestrian', 2, (220, 20, 60)),
        Class('cyclist', 3, (119, 11, 32)),
    ]

    def __init__(self, root, split='train', transform=None):
        self.root = os.path.expanduser(root)
        self.lidar_path = os.path.join(self.root, 'lidar_2d')
        self.split = os.path.join(self.root, 'ImageSet', '{}.txt'.format(split))
        self.transform = transform
        self._cache = os.path.join(self.root,"lidar_2d_cache")
        if split not in ['train', 'val', 'all']:
            raise ValueError('Invalid split! Use split="train", split="val" or split="all"')
        if not os.path.exists(self._cache):
            os.makedirs(self._cache)
        if not os.path.exists(osp.join(self._cache, split)):
                self.lidar = []

                with open(os.path.join(self.root, 'ImageSet', '{}.txt'.format(split))) as file:
                    images = ['{}.npy'.format(x.strip()) for x in file.readlines()]
                    for img in images:
                        lidar_2d = os.path.join(self.lidar_path, img)
                        self.lidar.append(lidar_2d)

                with lmdb.open(
                        osp.join(self._cache, split), map_size=1 << 33
                ) as lmdb_env, lmdb_env.begin(write=True) as txn:
                    for i in tqdm.trange(len(self.lidar)):
                        fn = self.lidar[i]

                        point_set = np.load(fn).astype(np.float32, copy=False)
                        txn.put(
                            str(i).encode(),
                            msgpack_numpy.packb(
                                dict(pc=point_set, ), use_bin_type=True
                            ),
                        )
        self._lmdb_file = osp.join(self._cache, split)
        with lmdb.open(self._lmdb_file, map_size=1 << 33) as lmdb_env:
            self._len = lmdb_env.stat()["entries"]

        self._lmdb_env = None

    def __getitem__(self, index):
        if self._lmdb_env is None:
            self._lmdb_env = lmdb.open(
                self._lmdb_file, map_size=1 << 33, readonly=True, lock=False
            )

        with self._lmdb_env.begin(buffers=True) as txn:
            ele = msgpack_numpy.unpackb(txn.get(str(index).encode()), raw=False)

        record = ele["pc"]
        record = torch.from_numpy(record.copy()).permute(2, 0, 1).contiguous()
        record.type=torch.float16

        distance = record[3, :, :]
        reflectivity = record[4, :, :]
        label = record[5, :, :]

        if self.transform:
            distance, reflectivity, label = self.transform(distance, reflectivity, label)

        return distance, reflectivity, label

    def __len__(self):
        return self._len

    @staticmethod
    def num_classes():
        return len(KITTI.classes)

    @staticmethod
    def mean():
        return [0.21, 12.12]

    @staticmethod
    def std():
        return [0.16, 12.32]

    @staticmethod
    def class_weights():
        return torch.tensor([1 / 15.0, 1.0, 10.0, 10.0])

    @staticmethod
    def get_colormap():
        cmap = torch.zeros([256, 3], dtype=torch.uint8)

        for cls in KITTI.classes:
            cmap[cls.id, :] = torch.tensor(cls.color, dtype=torch.uint8)

        return cmap

if __name__ == '__main__':
    import matplotlib.pyplot as plt

    joint_transforms = Compose([
        RandomHorizontalFlip(),
        Normalize(mean=KITTI.mean(), std=KITTI.std())
    ])

    def _normalize(x):
        return (x - x.min()) / (x.max() - x.min())

    def visualize_seg(label_map, one_hot=False):
        if one_hot:
            label_map = np.argmax(label_map, axis=-1)

        out = np.zeros((label_map.shape[0], label_map.shape[1], 3))

        for l in range(1, KITTI.num_classes()):
            mask = label_map == l
            out[mask, 0] = np.array(KITTI.classes[l].color[1])
            out[mask, 1] = np.array(KITTI.classes[l].color[0])
            out[mask, 2] = np.array(KITTI.classes[l].color[2])

        return out

    dataset = KITTI('../../data/kitti', transform=joint_transforms)
    distance, reflectivity, label = random.choice(dataset)

    print('Distance size: ', distance.size())
    print('Reflectivity size: ', reflectivity.size())
    print('Label size: ', label.size())

    distance_map = Image.fromarray((255 * _normalize(distance.numpy())).astype(np.uint8))
    reflectivity_map = Image.fromarray((255 * _normalize(reflectivity.numpy())).astype(np.uint8))
    label_map = Image.fromarray((255 * visualize_seg(label.numpy())).astype(np.uint8))

    blend_map = Image.blend(distance_map.convert('RGBA'), label_map.convert('RGBA'), alpha=0.4)

    plt.figure(figsize=(10, 5))
    plt.subplot(221)
    plt.title("Distance")
    plt.imshow(distance_map)
    plt.subplot(222)
    plt.title("Reflectivity")
    plt.imshow(reflectivity_map)
    plt.subplot(223)
    plt.title("Label")
    plt.imshow(label_map)
    plt.subplot(224)
    plt.title("Result")
    plt.imshow(blend_map)

    plt.show()


#!/usr/bin/env python3
import time
import math
from datetime import datetime
from time import sleep
import numpy as np
import random
import cv2
import os
import argparse
import torch

import sys
sys.path.append('./Eval')
sys.path.append('./')

from scipy.spatial.transform import Rotation as R

from env import Engine
from utils_env import get_view,safe_path,cut_frame,point2traj,get_gripper_pos,backup_code

class Engine13(Engine):
    def __init__(self, worker_id, opti, p_id, taskId=5, maxSteps=15, n_dmps=3, cReward=True):
        super(Engine13,self).__init__(opti, wid=worker_id, p_id=p_id, maxSteps=maxSteps, taskId=taskId, n_dmps=n_dmps, cReward=cReward,robot_model=None)
        self.opti = opti

    def init_obj(self):
        self.obj_file = os.path.join(self.urdf_dir, "objmodels/nut.urdf")
        self.obj_position = [0.3637 + 0.055, -0.065, 0.35]
        self.obj_scaling = 2.2
        self.obj_orientation = self.p.getQuaternionFromEuler([math.pi/2+0.2, -math.pi/2, -0.4])
        self.obj_id = self.p.loadURDF(fileName=self.obj_file, basePosition=self.obj_position,baseOrientation=self.obj_orientation,
                                     globalScaling=self.obj_scaling)#,physicsClientId=self.physical_id

        self.box_file = os.path.join (self.urdf_dir, "openbox/openbox.urdf")
        self.box_position = [0.35, 0.03, -0.33]
        self.box_scaling = 0.000225
        self.box_orientation = self.p.getQuaternionFromEuler ([0, 0, math.pi/2])
        self.box_id = self.p.loadURDF (fileName=self.box_file, basePosition=self.box_position,
                                      baseOrientation=self.box_orientation,
                                      globalScaling=self.box_scaling,useFixedBase=True)
        self.p.changeVisualShape (self.obj_id, -1, rgbaColor=[38/255.,0.,128/255.0,1])
        self.p.changeDynamics(self.obj_id,-1,mass=2.0)

        obj_friction_ceof = 2000.0
        self.p.changeDynamics(self.obj_id, -1, lateralFriction=obj_friction_ceof)
        self.p.changeDynamics(self.obj_id, -1, rollingFriction=obj_friction_ceof)
        self.p.changeDynamics(self.obj_id, -1, spinningFriction=obj_friction_ceof)
        self.p.changeDynamics (self.obj_id, -1, mass=0.01)
        self.p.changeDynamics (self.obj_id, -1, linearDamping=20.0)
        self.p.changeDynamics (self.obj_id, -1, angularDamping=20.0)
        self.p.changeDynamics (self.obj_id, -1, contactStiffness=0.1, contactDamping=0.9)

    def reset_obj(self):
        self.p.resetBasePositionAndOrientation(self.obj_id, self.obj_position, self.obj_orientation)

    def init_motion(self):
        self.data_q = np.load (os.path.join(self.robot_recordings_dir,"47-4/q.npy"))
        self.data_gripper = np.load (self.configs_dir + '/init/gripper.npy')
        self.robot.setJointValue(self.data_q[0],gripper=self.data_gripper[0])

    def init_grasp(self):
        self.p.resetBasePositionAndOrientation(self.box_id,self.box_position,self.box_orientation)

        self.robot.gripperControl(0)

        qlist = np.load( os.path.join(self.robot_recordings_dir, "47-4/q.npy"))
        glist = np.load( os.path.join(self.robot_recordings_dir, "47-4/gripper.npy"))
        num_q = len(qlist[0])
        self.fix_orn = np.load (os.path.join (self.configs_dir, 'init', 'orn.npy'))
        self.null_q = qlist[180]
        self.robot.setJointValue(qlist[40],glist[40])
        gripper_v = 120

        for i in range(40,180,1):
            glist[i] = min(gripper_v,glist[i])
            self.robot.jointPositionControl(qlist[i],gripper=glist[i])

        for _ in range(2):
          pos = self.robot.getEndEffectorPos()
          pos[2] += 0.05
          orn = self.robot.getEndEffectorOrn()
          for i in range(30):
            self.robot.positionControl(pos,orn,null_pose=self.null_q,gripperPos=gripper_v)

        self.start_pos = self.p.getLinkState (self.robotId, 7)[0]

        ######## obj is lifted

        transl = np.random.uniform(-0.05,0.05,size=(2,))
        self.box_new_position = np.array(self.box_position)
        self.box_new_position[2] *= -1.0
        self.box_new_position[:2] += transl

        r = R.from_quat(self.box_orientation)
        HTrans = np.zeros((4,4))
        HTrans[:3,:3] = r.as_dcm()

        rotation_degree = np.random.uniform(-math.pi,math.pi)
        addRot = R.from_rotvec(rotation_degree * np.array([0,0,1]))
        addHTrans = np.zeros((4,4))
        addHTrans[:3,:3] = addRot.as_dcm()
        NewHTrans = addHTrans.dot(HTrans)

        self.box_new_orientation = R.from_dcm(NewHTrans[:3,:3]).as_quat()

        self.p.resetBasePositionAndOrientation(self.box_id,self.box_new_position,self.box_new_orientation)

        ##### obj is reset
        cur_joint = self.robot.getJointValue()
        cur_pos = np.array(self.box_new_position)#self.robot.getendeffectorpos()
        cur_pos[:2]  = cur_pos[:2] * 0.5 + self.robot.getEndEffectorPos()[:2] * 0.5
        cur_orn = self.robot.getEndEffectorOrn()
        cur_pos[2] = self.robot.getEndEffectorPos()[2]
        print("before initialiaztion")
        for i in range(19):
          self.robot.positionControl(cur_pos,cur_orn,null_pose=cur_joint,gripperPos=gripper_v)

    #    input("raw")
        ####
    #    cur_joint = self.robot.getJointValue()
    #    cur_pos = np.array(self.box_new_position) + self.robot.getendeffectorpos()
    #    cur_orn = self.robot.getEndEffectorOrn()
    #    pos_diff = np.random.uniform(-0.1,0.1,size=(2,))
    #    cur_pos[:2] += pos_diff
    #    cur_pos[2] = self.robot.getEndEffectorPos()[2] + np.random.uniform(-0.04,0.01)
    #    print("before initialiaztion")
    #    for i in range(19):
    #      self.robot.positionControl(cur_pos,cur_orn,null_pose=cur_joint,gripperPos=gripper_v)

#        cur_pos[2] += np.random.uniform(-0.05,0.05)

        for _ in range(20):
          self.p.stepSimulation()

    def init_grasp_(self):

        self.box_position[2] = -0.33
        self.p.resetBasePositionAndOrientation(self.box_id,self.box_position,self.box_orientation)

        self.robot.gripperControl(0)

        qlist = np.load( os.path.join(self.robot_recordings_dir, "47-4/q.npy"))
        glist = np.load( os.path.join(self.robot_recordings_dir, "47-4/gripper.npy"))
        num_q = len(qlist[0])
        self.fix_orn = np.load (os.path.join (self.configs_dir, 'init', 'orn.npy'))
        self.null_q = qlist[180]
        self.robot.setJointValue(qlist[40],glist[40])
        for i in range(40,180,1):
            glist[i] = min(100,glist[i])
            self.robot.jointPositionControl(qlist[i],gripper=glist[i])

        pos = self.robot.getEndEffectorPos()
        pos[2] += 0.05
        orn = self.robot.getEndEffectorOrn()

        for i in range(109):
           self.robot.positionControl(pos,orn,null_pose=self.null_q,gripperPos=100)
        self.start_pos = self.p.getLinkState (self.robotId, 7)[0]

        box_diff = np.random.uniform(-0.1,0.1,size=(2,))
        self.box_position[2] *= -1.0
        self.box_pos =  np.array([self.box_position[0]+box_diff[0],self.box_position[1]+box_diff[1],self.box_position[2]])
        self.p.resetBasePositionAndOrientation(self.box_id,self.box_pos,self.box_orientation)

        cur_joint = self.robot.getJointValue()
        cur_pos = self.robot.getEndEffectorPos()
        cur_orn = self.robot.getEndEffectorOrn()
        pos_diff = np.random.uniform(-0.1,0.1,size=(2,))
        cur_pos[0] =  self.box_position[0]
        cur_pos[1] = self.box_position[1]
        cur_pos[:2] = cur_pos[:2] + pos_diff
        cur_pos[2] += 0.02
        for i in range(19):
           self.robot.positionControl(cur_pos,cur_orn,null_pose=cur_joint,gripperPos=120)

    def get_success(self,seg=None):
        box = self.p.getAABB (self.box_id, -1)
        box_center = [(x + y) * 0.5 for x, y in zip (box[0], box[1])]
        box_center = np.array(box_center)
        obj = self.p.getAABB (self.obj_id, -1)
        obj_center = [(x + y) * 0.5 for x, y in zip (obj[0], obj[1])]
        obj_center = np.array(obj_center)
        dist = np.linalg.norm(box_center - obj_center)
        # check whether the object is still in the gripper
        left_closet_info = self.p.getContactPoints (self.robotId, self.obj_id, self.robot.gripper_left_tip_index, -1)
        right_closet_info = self.p.getContactPoints (self.robotId, self.obj_id, self.robot.gripper_right_tip_index, -1)
        box_obj_closet_info = self.p.getContactPoints(self.obj_id, self.box_id)
        #print(len (left_closet_info),len (right_closet_info),obj[0][0], box[1][0])
        if len (left_closet_info) > 0 and len (right_closet_info) > 0:
          if dist < 0.05 or len(box_obj_closet_info) > 0:
            return True
          else:
            return False
        else:
          return False


#encoding:utf-8
'''
用于解码原始SSD预测输出的自定义Keras层。 对应于SSD的原始Caffe实现中的“DetectionOutput”层。

Copyright (C) 2018 Pierluigi Ferrari

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.
'''

from __future__ import division
import numpy as np
import tensorflow as tf
import keras.backend as K
from keras.engine.topology import InputSpec
from keras.engine.topology import Layer

class DecodeDetections(Layer):
    '''
    一个用来解码预测数据的keras层
    输入shape: 3D  `(batch_size, n_boxes, n_classes + 12)`.
    输出shape: 3D  `(batch_size, top_k, 6)`.
    '''

    def __init__(self,
                 confidence_thresh=0.01,
                 iou_threshold=0.45,
                 top_k=200,
                 nms_max_output_size=400,
                 coords='centroids',
                 normalize_coords=True,
                 img_height=None,
                 img_width=None,
                 **kwargs):

        if K.backend() != 'tensorflow':
            raise TypeError("This layer only supports TensorFlow at the moment, but you are using the {} backend.".format(K.backend()))

        if normalize_coords and ((img_height is None) or (img_width is None)):
            raise ValueError("If relative box coordinates are supposed to be converted to absolute coordinates, the decoder needs the image size in order to decode the predictions, but `img_height == {}` and `img_width == {}`".format(img_height, img_width))

        if coords != 'centroids':
            raise ValueError("The DetectionOutput layer currently only supports the 'centroids' coordinate format.")

        self.confidence_thresh = confidence_thresh
        self.iou_threshold = iou_threshold
        self.top_k = top_k
        self.normalize_coords = normalize_coords
        self.img_height = img_height
        self.img_width = img_width
        self.coords = coords
        self.nms_max_output_size = nms_max_output_size

        self.tf_confidence_thresh = tf.constant(self.confidence_thresh, name='confidence_thresh')
        self.tf_iou_threshold = tf.constant(self.iou_threshold, name='iou_threshold')
        self.tf_top_k = tf.constant(self.top_k, name='top_k')
        self.tf_normalize_coords = tf.constant(self.normalize_coords, name='normalize_coords')
        self.tf_img_height = tf.constant(self.img_height, dtype=tf.float32, name='img_height')
        self.tf_img_width = tf.constant(self.img_width, dtype=tf.float32, name='img_width')
        self.tf_nms_max_output_size = tf.constant(self.nms_max_output_size, name='nms_max_output_size')

        super(DecodeDetections, self).__init__(**kwargs)

    def build(self, input_shape):
        self.input_spec = [InputSpec(shape=input_shape)]
        super(DecodeDetections, self).build(input_shape)

    def call(self, y_pred, mask=None):
        '''
        Returns:(batch_size, top_k, 6). The second axis is zero-padded
            to always yield `top_k` predictions per batch item. The last axis contains
            the coordinates for each predicted box in the format
            `[class_id, confidence, xmin, ymin, xmax, ymax]`.
        '''

        #####################################################################################
        # 1. 将预测偏移值转成绝对坐标
        #    (batch, n_boxes_total, n_classes + 4 + 8)
        #    [cx_pred,cy_pred,w_pred,y_pred,cx_anchor,cy_anchor,w_anchor,y_anchor,cx_variance,cy_variance,variance_w,variance_h]
        #####################################################################################

        # 将预测的偏移值转成坐标
        cx = y_pred[...,-12] * y_pred[...,-4] * y_pred[...,-6] + y_pred[...,-8] # cx = cx_pred * cx_variance * w_anchor + cx_anchor
        cy = y_pred[...,-11] * y_pred[...,-3] * y_pred[...,-5] + y_pred[...,-7] # cy = cy_pred * cy_variance * h_anchor + cy_anchor
        w = tf.exp(y_pred[...,-10] * y_pred[...,-2]) * y_pred[...,-6] # w = exp(w_pred * variance_w) * w_anchor
        h = tf.exp(y_pred[...,-9] * y_pred[...,-1]) * y_pred[...,-5] # h = exp(h_pred * variance_h) * h_anchor

        #  'centroids' -> 'corners'.
        xmin = cx - 0.5 * w
        ymin = cy - 0.5 * h
        xmax = cx + 0.5 * w
        ymax = cy + 0.5 * h

        def normalized_coords():
            xmin1 = tf.expand_dims(xmin * self.tf_img_width, axis=-1)
            ymin1 = tf.expand_dims(ymin * self.tf_img_height, axis=-1)
            xmax1 = tf.expand_dims(xmax * self.tf_img_width, axis=-1)
            ymax1 = tf.expand_dims(ymax * self.tf_img_height, axis=-1)
            return xmin1, ymin1, xmax1, ymax1
        def non_normalized_coords():
            return tf.expand_dims(xmin, axis=-1), tf.expand_dims(ymin, axis=-1), tf.expand_dims(xmax, axis=-1), tf.expand_dims(ymax, axis=-1)

        xmin, ymin, xmax, ymax = tf.cond(self.tf_normalize_coords, normalized_coords, non_normalized_coords)

        y_pred = tf.concat(values=[y_pred[...,:-12], xmin, ymin, xmax, ymax], axis=-1)

        #####################################################################################
        # 2. - confidence thresholding
        # - non-maximum suppression (NMS)
        # - top-k filtering
        #####################################################################################

        batch_size = tf.shape(y_pred)[0] # Output dtype: tf.int32
        n_boxes = tf.shape(y_pred)[1]
        n_classes = y_pred.shape[2] - 4
        class_indices = tf.range(1, n_classes)

        def filter_predictions(batch_item):
            def filter_single_class(index):
                # 迭代每一个类别，从(n_boxes, n_classes + 4 coordinates)提取(n_boxes, 1 + 4 coordinates)

                confidences = tf.expand_dims(batch_item[..., index], axis=-1)
                class_id = tf.fill(dims=tf.shape(confidences), value=tf.to_float(index))
                box_coordinates = batch_item[...,-4:]

                single_class = tf.concat([class_id, confidences, box_coordinates], axis=-1)
                threshold_met = single_class[:,1] > self.tf_confidence_thresh
                single_class = tf.boolean_mask(tensor=single_class, mask=threshold_met)

                # 大于阈值的boxes需要NMS处理.
                def perform_nms():
                    scores = single_class[...,1]
                    xmin = tf.expand_dims(single_class[...,-4], axis=-1)
                    ymin = tf.expand_dims(single_class[...,-3], axis=-1)
                    xmax = tf.expand_dims(single_class[...,-2], axis=-1)
                    ymax = tf.expand_dims(single_class[...,-1], axis=-1)
                    boxes = tf.concat(values=[ymin, xmin, ymax, xmax], axis=-1)

                    maxima_indices = tf.image.non_max_suppression(boxes=boxes,
                                                                  scores=scores,
                                                                  max_output_size=self.tf_nms_max_output_size,
                                                                  iou_threshold=self.iou_threshold,
                                                                  name='non_maximum_suppresion')
                    maxima = tf.gather(params=single_class,
                                       indices=maxima_indices,
                                       axis=0)
                    return maxima

                def no_confident_predictions():
                    return tf.constant(value=0.0, shape=(1,6))

                single_class_nms = tf.cond(tf.equal(tf.size(single_class), 0), no_confident_predictions, perform_nms)

                # 确保 `single_class` 是 `self.nms_max_output_size` 这么多的元素
                padded_single_class = tf.pad(tensor=single_class_nms,
                                             paddings=[[0, self.tf_nms_max_output_size - tf.shape(single_class_nms)[0]], [0, 0]],
                                             mode='CONSTANT',
                                             constant_values=0.0)

                return padded_single_class

            # 迭代所有类别
            filtered_single_classes = tf.map_fn(fn=lambda i: filter_single_class(i),
                                                elems=tf.range(1,n_classes),
                                                dtype=tf.float32,
                                                parallel_iterations=128,
                                                back_prop=False,
                                                swap_memory=False,
                                                infer_shape=True,
                                                name='loop_over_classes')

            # 将预测结果整合为一个张量
            filtered_predictions = tf.reshape(tensor=filtered_single_classes, shape=(-1,6))

            # Perform top-k filtering for this batch item or pad it in case there are
            # fewer than `self.top_k` boxes left at this point. Either way, produce a
            # tensor of length `self.top_k`. By the time we return the final results tensor
            # for the whole batch, all batch items must have the same number of predicted
            # boxes so that the tensor dimensions are homogenous. If fewer than `self.top_k`
            # predictions are left after the filtering process above, we pad the missing
            # predictions with zeros as dummy entries.
            def top_k():
                return tf.gather(params=filtered_predictions,
                                 indices=tf.nn.top_k(filtered_predictions[:, 1], k=self.tf_top_k, sorted=True).indices,
                                 axis=0)
            def pad_and_top_k():
                padded_predictions = tf.pad(tensor=filtered_predictions,
                                            paddings=[[0, self.tf_top_k - tf.shape(filtered_predictions)[0]], [0, 0]],
                                            mode='CONSTANT',
                                            constant_values=0.0)
                return tf.gather(params=padded_predictions,
                                 indices=tf.nn.top_k(padded_predictions[:, 1], k=self.tf_top_k, sorted=True).indices,
                                 axis=0)

            top_k_boxes = tf.cond(tf.greater_equal(tf.shape(filtered_predictions)[0], self.tf_top_k), top_k, pad_and_top_k)

            return top_k_boxes

        # Iterate `filter_predictions()` over all batch items.
        output_tensor = tf.map_fn(fn=lambda x: filter_predictions(x),
                                  elems=y_pred,
                                  dtype=None,
                                  parallel_iterations=128,
                                  back_prop=False,
                                  swap_memory=False,
                                  infer_shape=True,
                                  name='loop_over_batch')

        return output_tensor

    def compute_output_shape(self, input_shape):
        batch_size, n_boxes, last_axis = input_shape
        return (batch_size, self.tf_top_k, 6) # 最后一维: (class_ID, confidence, 4 box coordinates)

    def get_config(self):
        config = {
            'confidence_thresh': self.confidence_thresh,
            'iou_threshold': self.iou_threshold,
            'top_k': self.top_k,
            'nms_max_output_size': self.nms_max_output_size,
            'coords': self.coords,
            'normalize_coords': self.normalize_coords,
            'img_height': self.img_height,
            'img_width': self.img_width,
        }
        base_config = super(DecodeDetections, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


#  Created byMartin.cz
#  Copyright (c) Martin Strohalm. All rights reserved.

import numpy

from pero.enums import *
from pero.properties import *
from pero import Profile as ProfileGlyph
from pero import Path, Marker, MarkerLegend

from . series import Series
from . import utils

class Profile(Series):
    """
    This type of series plots raw x-sorted data as continuous line. Data can be
    provided either directly by specifying the 'x' and 'y' properties or as a
    sequence of raw 'data' points together with 'x' and 'y' coordinates
    selectors. All the coordinates as well as the 'base' are expected to be in
    real data units.

    To plot data as a line the 'show_line' property must be set to True.
    Similarly, to display individual points, the 'show_points' property must be
    set to True. If set to UNDEF, and line is enabled, the points are
    automatically visible if they are separated enough, as defined by the
    'spacing' property.

    Any property of the 'marker' can be dynamic (including the 'marker' property
    itself), expecting the raw data point as a 'source'. By this, its color,
    fill and other properties, as well as the marker itself can be set
    independently for each data point. However, be sure that all dynamic
    properties return reasonable value for UNDEF to be used for legend. If raw
    'data' property is not specified a sequence of internal raw data is created
    as ((x,y),) coordinates.

    Optionally the area under the curve can also be displayed if the 'show_area'
    property is set to True. In such case the line is additionally drawn as a
    filled polygon either by connecting first and last points directly or using
    strait line defined by the 'base' property.

    Properties:

        show_line: bool
            Specifies whether the profile line should be displayed.

        show_points: bool or UNDEF
            Specifies whether individual points should be displayed. If set to
            UNDEF, the points are displayed automatically as long as their
            distance is big enough.

        show_area: bool
            Specifies whether the area under profile line should be displayed.

        data: tuple, list, numpy.ndarray or UNDEF
            Specifies the sequence of the raw data points.

        x: int, float, tuple, list, numpy.ndarray, callable, None or UNDEF
            Specifies the sequence of x-coordinates in real data units or a
            function to retrieve the coordinates from the raw data.

        y: int, float, tuple, list, numpy.ndarray, callable, None or UNDEF
            Specifies the sequence of y-coordinates in real data units or a
            function to retrieve the coordinates from the raw data.

        base: int, float, None or UNDEF
            Specifies the area base value in real data units.

        marker: pero.MARKER, pero.Path, callable, None or UNDEF
            Specifies the marker to draw actual data points with. The value
            can be specified by any item from the pero.MARKER enum or as
            pero.Path.

        marker_size: int, float or callable
            Specifies the marker size.

        marker_line properties:
            Includes pero.LineProperties to specify the marker line.

        marker_fill properties:
            Includes pero.FillProperties to specify the marker fill.

        steps: pero.LINE_STEP
            Specifies the way stepped profile should be drawn as any item from
            the pero.LINE_STEP enum.

        spacing: int, float
            Specifies the minimum x-distance between points in device units to
            enable automatic points display.

        line properties:
            Includes pero.LineProperties to specify the line.

        fill properties:
            Includes pero.FillProperties to specify the area fill.
    """

    show_line = BoolProperty(True, dynamic=False)
    show_points = BoolProperty(UNDEF, dynamic=False)
    show_area = BoolProperty(False, dynamic=False)

    data = SequenceProperty(UNDEF, dynamic=False)
    x = Property(UNDEF)
    y = Property(UNDEF)
    base = NumProperty(UNDEF, dynamic=False, nullable=True)
    steps = EnumProperty(None, enum=LINE_STEP, nullable=True)
    spacing = NumProperty(20, dynamic=False)

    marker = Property(MARKER_CIRCLE, types=(str, Path, Marker), nullable=True)
    marker_size = NumProperty(6)
    marker_line = Include(LineProperties, prefix='marker_', line_color=UNDEF)
    marker_fill = Include(FillProperties, prefix='marker_', fill_color=UNDEF)

    line = Include(LineProperties, line_color=UNDEF, dynamic=False)
    fill = Include(FillProperties, fill_color=UNDEF, dynamic=False)

    def __init__(self, **overrides):
        """Initializes a new instance of Profile series."""

        # init legend
        if 'legend' not in overrides:
            overrides['legend'] = MarkerLegend(
                text = lambda d: d.title,
                show_marker = lambda d: d.show_points,
                show_line = lambda d: d.show_line,
                marker = lambda d: d.marker,
                marker_line_color = lambda d: d.color.darker(0.2),
                marker_fill_color = lambda d: d.color,
                line_color = lambda d: d.color,
                line_width = lambda d: d.line_width,
                line_style = lambda d: d.line_style,
                line_dash = lambda d: d.line_dash)

        # init base
        super().__init__(**overrides)

        # init profile glyph
        self._glyph = ProfileGlyph()

        # init buffers
        self._x_data = []
        self._y_data = []
        self._raw_data = []
        self._limits = None

        # extract data
        self.extract_data()

    def get_limits(self, x_range=None, y_range=None, exact=False):
        """Gets current data limits using whole range or specified crops."""

        # check data
        if self._limits is None:
            return None

        # init limits
        limits = self._limits

        # apply crop
        if x_range:

            limits = utils.calc_limits_sorted(
                data=(self._x_data, self._y_data),
                crop=x_range,
                extend=False,
                interpolate=True)

            if self.base not in (UNDEF, None) and limits[1] is not None:
                limits[1][0] = min(self.base, limits[1][0])
                limits[1][1] = max(self.base, limits[1][1])

        # finalize limits
        return self.finalize_limits(limits, exact)

    def get_labels(self, canvas=None, source=UNDEF, **overrides):
        """Gets series labels."""

        return self.prepare_labels(self._x_data, self._y_data, self._raw_data)

    def get_tooltip(self, x, y, limit):
        """Gets nearest data point tooltip."""

        return self.prepare_tooltip(self._x_data, self._y_data, self._raw_data, x, y, limit)

    def extract_data(self):
        """Extracts coordinates from raw data."""

        # reset buffers
        self._x_data = []
        self._y_data = []
        self._raw_data = []
        self._limits = None

        # get data size
        size = utils.extract_data_size(self, 'data', 'x', 'y')

        # extract data
        self._x_data, x_raw = utils.extract_data(self, 'x', self.data, size, self.x_mapper)
        self._y_data, y_raw = utils.extract_data(self, 'y', self.data, size, self.y_mapper)

        # set raw data
        if self.data is not UNDEF:
            self._raw_data = numpy.array(self.data)
        else:
            self._raw_data = numpy.array([x_raw, y_raw]).T

        # check data
        if not utils.is_sorted(self._x_data):
            raise ValueError("X-coordinates must be sorted!")

        # init full limits
        if len(self._raw_data) > 0:
            self._limits = (
                (self._x_data.min(), self._x_data.max()),
                (self._y_data.min(), self._y_data.max()))

    def draw(self, canvas, source=UNDEF, **overrides):
        """Uses given canvas to draw the series."""

        # check if visible
        if not self.is_visible(source, overrides):
            return

        # get properties
        tag = self.get_property('tag', source, overrides)
        x_scale = self.get_property('x_scale', source, overrides)
        y_scale = self.get_property('y_scale', source, overrides)
        base = self.get_property('base', source, overrides)
        color = self.get_property('color', source, overrides)

        # set overrides to skip
        skip = {'x', 'y', 'base'}

        # crop data
        i1, i2 = utils.crop_indices(self._x_data, x_scale.in_range, True)
        if i1 == i2:
            return

        x_data = self._x_data[i1:i2]
        y_data = self._y_data[i1:i2]
        raw_data = self._raw_data[i1:i2]

        # scale coords
        x_data = x_scale.scale(x_data)
        y_data = y_scale.scale(y_data)

        # scale base
        if base is not UNDEF:
            base = y_scale.scale(base)

        # start drawing group
        with canvas.group(tag, "series"):

            # update glyph
            self._glyph.set_properties_from(self, source=source, overrides=overrides, skip=skip)

            # get glyph colors
            line_color = self._glyph.get_property('line_color')
            if line_color is UNDEF:
                line_color = color

            fill_color = self._glyph.get_property('fill_color')
            if fill_color is UNDEF:
                fill_color = color.trans(0.4)

            marker_line_color = self._glyph.get_property('marker_line_color')
            if marker_line_color is UNDEF:
                marker_line_color = color.darker(0.2)

            marker_fill_color = self._glyph.get_property('marker_fill_color')
            if marker_fill_color is UNDEF:
                marker_fill_color = color

            # set overrides
            glyph_overrides = overrides.copy()
            glyph_overrides['x'] = x_data
            glyph_overrides['y'] = y_data
            glyph_overrides['base'] = base
            glyph_overrides['line_color'] = line_color
            glyph_overrides['fill_color'] = fill_color
            glyph_overrides['marker_line_color'] = marker_line_color
            glyph_overrides['marker_fill_color'] = marker_fill_color

            # draw profile
            self._glyph.draw(canvas, raw_data, **glyph_overrides)


import numpy as np
import tensorflow as tf

def get_batches(X, y, batch_size):
    for i in range(0, len(X), batch_size):
        yield X[i:i+batch_size], y[i:i+batch_size]

def encoded_to_sentence(stn, tokenizer):
    return ' '.join([tokenizer.index_word[w] for w in stn if w!=0])

def sample_from_model(inputs, target, model, texts_tokenizer, headlines_tokenizer):
    idx = np.random.choice(np.arange(len(inputs)), size=1)[0]
    inp = inputs[idx:idx+1]
    tar = target[idx:idx+1]

    predictions = model(inp)
    predictions = [np.random.choice(np.arange(len(headlines_tokenizer.word_counts)+1), p=p.reshape(-1)) for p in predictions[0, ...].numpy()]

    print("Original text: {}".format(encoded_to_sentence(inp.reshape(-1).tolist(), texts_tokenizer)))
    print("Original headline: {}".format(encoded_to_sentence(tar.reshape(-1).tolist(), headlines_tokenizer)))
    print("Generated headline: {}".format(encoded_to_sentence(predictions, headlines_tokenizer)))


#! /usr/bin/env python

import tensorflow as tf
import numpy as np
import os
import time
import datetime
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from builddata_ecir import *
from capsuleNet_SEARCH17 import CapsE

np.random.seed(1234)
tf.set_random_seed(1234)

# Parameters
# ==================================================
parser = ArgumentParser("CapsE", formatter_class=ArgumentDefaultsHelpFormatter, conflict_handler='resolve')

parser.add_argument("--data", default="./data/", help="Data sources.")
parser.add_argument("--run_folder", default="./", help="Data sources.")
parser.add_argument("--name", default="SEARCH17", help="Name of the dataset.")

parser.add_argument("--embedding_dim", default=200, type=int, help="Dimensionality of character embedding (fixed: 200)")
parser.add_argument("--filter_size", default=1, type=int, help="Comma-separated filter sizes (default: '3,4,5')")
parser.add_argument("--num_filters", default=400, type=int, help="Number of filters per filter size (default: 128)")
parser.add_argument("--learning_rate", default=0.00001, type=float, help="Learning rate")
parser.add_argument("--batch_size", default=128, type=int, help="Batch Size")
parser.add_argument("--neg_ratio", default=1.0, help="Number of negative triples generated by positive (default: 1.0)")
parser.add_argument("--useInitialization", default=True, type=bool, help="Using the pretrained embeddings")
parser.add_argument("--num_epochs", default=100, type=int, help="Number of training epochs")
parser.add_argument("--savedEpochs", default=10, type=int, help="")
parser.add_argument("--allow_soft_placement", default=True, type=bool, help="Allow device soft device placement")
parser.add_argument("--log_device_placement", default=False, type=bool, help="Log placement of ops on devices")
parser.add_argument("--model_name", default='search17model', help="")
parser.add_argument("--useConstantInit", action='store_true')

parser.add_argument('--iter_routing', default=1, type=int, help='number of iterations in routing algorithm')
parser.add_argument('--num_outputs_secondCaps', default=1, type=int, help='')
parser.add_argument('--vec_len_secondCaps', default=10, type=int, help='')

args = parser.parse_args()
print(args)
# Load data
# Load data
print("Loading data...")

train_triples, train_rank_triples, train_val_triples, valid_triples, valid_rank_triples, valid_val_triples, \
            test_triples, test_rank_triples, test_val_triples, query_indexes, user_indexes, doc_indexes, \
            indexes_query, indexes_user, indexes_doc = build_data_ecir()
data_size = len(train_triples)
train_batch = Batch_Loader_ecir(train_triples, train_val_triples, batch_size=args.batch_size)

assert args.embedding_dim % 200 == 0

pretrained_query = init_dataset_ecir(args.data + args.name + '/query2vec.200.init')
pretrained_user = init_dataset_ecir(args.data + args.name + '/user2vec.200.init')
pretrained_doc = init_dataset_ecir(args.data + args.name + '/doc2vec.200.init')

print("Using pre-trained initialization.")

lstEmbedQuery = assignEmbeddings(pretrained_query, query_indexes)
lstEmbedUser = assignEmbeddings(pretrained_user, user_indexes)
lstEmbedDoc = assignEmbeddings(pretrained_doc, doc_indexes)

lstEmbedQuery = np.array(lstEmbedQuery, dtype=np.float32)
lstEmbedUser = np.array(lstEmbedUser, dtype=np.float32)
lstEmbedDoc = np.array(lstEmbedDoc, dtype=np.float32)

print("Loading data... finished!")

# Training
# ==================================================
with tf.Graph().as_default():
    session_conf = tf.ConfigProto(allow_soft_placement=args.allow_soft_placement, log_device_placement=args.log_device_placement)
    session_conf.gpu_options.allow_growth = True
    sess = tf.Session(config=session_conf)
    with sess.as_default():
        global_step = tf.Variable(0, name="global_step", trainable=False)
        capse = CapsE(sequence_length=3,
                      batch_size=20 * args.batch_size,
                      initialization=[lstEmbedQuery, lstEmbedUser, lstEmbedDoc],
                            embedding_size=200,
                            filter_size=args.filter_size,
                            num_filters=args.num_filters,
                            iter_routing=args.iter_routing,
                            num_outputs_secondCaps=args.num_outputs_secondCaps,
                            vec_len_secondCaps=args.vec_len_secondCaps,
                            useConstantInit=args.useConstantInit
                            )

        # Define Training procedure
        #optimizer = tf.contrib.opt.NadamOptimizer(1e-3)
        optimizer = tf.train.AdamOptimizer(learning_rate=args.learning_rate)
        #optimizer = tf.train.RMSPropOptimizer(learning_rate=args.learning_rate)
        #optimizer = tf.train.GradientDescentOptimizer(learning_rate=args.learning_rate)
        grads_and_vars = optimizer.compute_gradients(capse.total_loss)
        train_op = optimizer.apply_gradients(grads_and_vars, global_step=global_step)

        out_dir = os.path.abspath(os.path.join(args.run_folder, "runs_CapsE_SEARCH17", args.model_name))
        print("Writing to {}\n".format(out_dir))

        checkpoint_dir = os.path.abspath(os.path.join(out_dir, "checkpoints"))
        checkpoint_prefix = os.path.join(checkpoint_dir, "model")
        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
        # Initialize all variables
        sess.run(tf.global_variables_initializer())

        def train_step(x_batch, y_batch):
            """
            A single training step
            """
            feed_dict = {
                capse.input_x: x_batch,
                capse.input_y: y_batch
            }
            _, step, loss = sess.run([train_op, global_step, capse.total_loss], feed_dict)
            return loss

        # Predict function to predict scores for test data
        def predict(x_batch, y_batch):
            feed_dict = {
                capse.input_x: x_batch,
                capse.input_y: y_batch,
                capse.dropout_keep_prob: 1.0
            }
            scores = sess.run([capse.predictions], feed_dict)
            return scores

        def test_prediction(x_batch, y_batch, lstOriginalRank):

            new_x_batch = np.concatenate(x_batch)
            new_y_batch = np.concatenate(y_batch, axis=0)

            while len(new_x_batch) % (args.batch_size * 20) != 0:
                new_x_batch = np.append(new_x_batch, np.array([new_x_batch[-1]]), axis=0)
                new_y_batch = np.append(new_y_batch, np.array([new_y_batch[-1]]), axis=0)

            results = []
            listIndexes = range(0, len(new_x_batch), 20 * args.batch_size)
            for tmpIndex in range(len(listIndexes) - 1):
                results = np.append(results,
                                    predict(new_x_batch[listIndexes[tmpIndex]:listIndexes[tmpIndex + 1]],
                                            new_y_batch[listIndexes[tmpIndex]:listIndexes[tmpIndex + 1]]))
            results = np.append(results,
                                predict(new_x_batch[listIndexes[-1]:], new_y_batch[listIndexes[-1]:]))

            lstresults = []
            _start = 0
            for tmp in lstOriginalRank:
                _end = _start + len(tmp)
                lstsorted = np.argsort(results[_start:_end])
                lstresults.append(np.where(lstsorted == 0)[0] + 1)
                _start = _end

            return lstresults

        wri = open(checkpoint_prefix + '.cls.' + '.txt', 'w')

        lstvalid_mrr = []
        lsttest_mrr = []
        num_batches_per_epoch = int((data_size - 1) / (args.batch_size)) + 1
        for epoch in range(args.num_epochs):
            for batch_num in range(num_batches_per_epoch):
                x_batch, y_batch = train_batch()
                train_step(x_batch, y_batch)
                current_step = tf.train.global_step(sess, global_step)

            valid_results = test_prediction(valid_triples, valid_val_triples, valid_rank_triples)
            test_results = test_prediction(test_triples, test_val_triples, test_rank_triples)
            valid_mrr = computeMRR(valid_results)
            test_mrr = computeMRR(test_results)
            test_p1 = computeP1(test_results)
            lstvalid_mrr.append(valid_mrr)
            lsttest_mrr.append([test_mrr, test_p1])

            wri.write("epoch " + str(epoch) + ": " + str(valid_mrr) + " " + str(test_mrr) + " " + str(test_p1) + "\n")

        index_valid_max = np.argmax(lstvalid_mrr)
        wri.write("\n--------------------------\n")
        wri.write("\nBest mrr in valid at epoch " + str(index_valid_max) + ": " + str(lstvalid_mrr[index_valid_max]) + "\n")
        wri.write("\nMRR and P1 in test: " + str(lsttest_mrr[index_valid_max][0]) + " " + str(lsttest_mrr[index_valid_max][1]) + "\n")
        wri.close()