# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

from __future__ import absolute_import

import subprocess
import sys
import io
import os
import boto3
import time
import json
import uuid

import mxnet as mx
import numpy as np
from mxnet import gluon,nd
from sagemaker_inference import content_types, default_inference_handler, errors
from io import BytesIO
from datetime import datetime

import gluoncv
from gluoncv.data.transforms import video
from gluoncv.data import VideoClsCustom
from gluoncv.utils.filesystem import try_import_decord

ctx = mx.gpu(0) if mx.context.num_gpus() > 0 else mx.cpu()
#UCF101 classes
classes = ['ApplyEyeMakeup'
,'ApplyLipstick'
, 'Archery'
, 'BabyCrawling'
, 'BalanceBeam'
, 'BandMarching'
, 'BaseballPitch'
, 'Basketball'
, 'BasketballDunk'
, 'BenchPress'
, 'Biking'
, 'Billiards'
, 'BlowDryHair'
, 'BlowingCandles'
, 'BodyWeightSquats'
, 'Bowling'
, 'BoxingPunchingBag'
, 'BoxingSpeedBag'
, 'BreastStroke'
, 'BrushingTeeth'
, 'CleanAndJerk'
, 'CliffDiving'
, 'CricketBowling'
, 'CricketShot'
, 'CuttingInKitchen'
, 'Diving'
, 'Drumming'
, 'Fencing'
, 'FieldHockeyPenalty'
, 'FloorGymnastics'
, 'FrisbeeCatch'
, 'FrontCrawl'
, 'GolfSwing'
, 'Haircut'
, 'Hammering'
, 'HammerThrow'
, 'HandstandPushups'
, 'HandstandWalking'
, 'HeadMassage'
, 'HighJump'
, 'HorseRace'
, 'HorseRiding'
, 'HulaHoop'
, 'IceDancing'
, 'JavelinThrow'
, 'JugglingBalls'
, 'JumpingJack'
, 'JumpRope'
, 'Kayaking'
, 'Knitting'
, 'LongJump'
, 'Lunges'
, 'MilitaryParade'
, 'Mixing'
, 'MoppingFloor'
, 'Nunchucks'
, 'ParallelBars'
, 'PizzaTossing'
, 'PlayingCello'
, 'PlayingDaf'
, 'PlayingDhol'
, 'PlayingFlute'
, 'PlayingGuitar'
, 'PlayingPiano'
, 'PlayingSitar'
, 'PlayingTabla'
, 'PlayingViolin'
, 'PoleVault'
, 'PommelHorse'
, 'PullUps'
, 'Punch'
, 'PushUps'
, 'Rafting'
, 'RockClimbingIndoor'
, 'RopeClimbing'
, 'Rowing'
, 'SalsaSpin'
, 'ShavingBeard'
, 'Shotput'
, 'SkateBoarding'
, 'Skiing'
, 'Skijet'
, 'SkyDiving'
, 'SoccerJuggling'
, 'SoccerPenalty'
, 'StillRings'
, 'SumoWrestling'
, 'Surfing'
, 'Swing'
, 'TableTennisShot'
, 'TaiChi'
, 'TennisSwing'
, 'ThrowDiscus'
, 'TrampolineJumping'
, 'Typing'
, 'UnevenBars'
, 'VolleyballSpiking'
, 'WalkingWithDog'
, 'WallPushups'
, 'WritingOnBoard'
, 'YoYo']
dict_classes = dict(zip(range(len(classes)), classes))
# ------------------------------------------------------------ #
# Hosting methods                                              #
# ------------------------------------------------------------ #

def model_fn(model_dir):
    print('here')
    print(ctx)
    symbol = mx.sym.load('%s/model-symbol.json' % model_dir)
    outputs = mx.symbol.softmax(data=symbol, name='softmax_label')
    inputs = mx.sym.var('data')
    net = gluon.SymbolBlock(outputs, inputs)
    net.load_parameters('%s/model-0000.params' % model_dir, ctx=ctx)
    return net

#transform function that uses json (s3 path) as input and output
def transform_fn(net, data, input_content_type, output_content_type):
    print('transform_fn here')
    start = time.time()
    data = json.loads(data)
    video_data = read_video_data(data['S3_VIDEO_PATH'])
    print(time.time())
    video_input = video_data.as_in_context(ctx)
    probs = net(video_input.astype('float32', copy=False))
    print(time.time())
    predicted = mx.nd.argmax(probs, axis=1).asnumpy().tolist()[0]
    probability = mx.nd.max(probs, axis=1).asnumpy().tolist()[0]

    probability = '{:.4f}'.format(probability)
    predicted_name = dict_classes[int(predicted)]
    total_prediction = time.time()-start
    total_prediction = '{:.4f}'.format(total_prediction)
    print(probability)
    print(predicted_name)
    print('Model prediction time: ', total_prediction)

    now = datetime.utcnow()
    time_format = '%Y-%m-%d %H:%M:%S %Z%z'
    now = now.strftime(time_format)

    response = {
        'S3Path': {'S': data['S3_VIDEO_PATH']},
        'Predicted': {'S': predicted_name},
        'Probability': {'S': probability},
        'DateCreatedUTC': {'S': now},
    }

    return json.dumps(response), output_content_type

def get_bucket_and_key(s3_path):
    """Get the bucket name and key from the given path.
    Args:
        s3_path(str): Input S3 path
    """
    s3_path = s3_path.replace('s3://', '')
    s3_path = s3_path.replace('S3://', '') #Both cases
    bucket, key = s3_path.split('/', 1)
    return bucket, key

def read_video_data(s3_video_path, num_frames=32):
    """Read and preprocess video data from the S3 bucket."""
    print('read and preprocess video data here ')
    s3_client = boto3.client('s3')
    #print(uuid.uuid4())
    fname = s3_video_path.replace('s3://', '')
    fname = fname.replace('S3://', '')
    fname = fname.replace('/', '')
    #download_path = '/tmp/{}-{}'.format(uuid.uuid4(), fname)
    #video_list_path = '/tmp/{}-{}'.format(uuid.uuid4(), 'video_list.txt')
    download_path = '/tmp/' + fname
    video_list_path = '/tmp/video_list' + str(uuid.uuid4()) + '.txt'
    bucket, key = get_bucket_and_key(s3_video_path)
    s3_client.download_file(bucket, key, download_path)

    #update download_path filename to be unique
    filename,ext = os.path.splitext(download_path)    # save the file extension
    filename = filename + str(uuid.uuid4())
    os.rename(download_path, filename+ext)
    download_path = filename+ext

    #Dummy duration and label with each video path
    video_list = '{} {} {}'.format(download_path, 10, 1)
    with open(video_list_path, 'w') as fopen:
        fopen.write(video_list)

    #Constants
    data_dir = '/tmp/'
    num_segments = 1
    new_length = num_frames
    new_step =1
    use_decord = True
    video_loader = True
    slowfast = False
    #Preprocessing params

    #The transformation function does three things: center crop the image to 224x224 in size, transpose it to num_channels,num_frames,height*width, and normalize with mean and standard deviation calculated across all ImageNet images.

    #Use the general gluoncv dataloader VideoClsCustom to load the data with num_frames = 32 as the length.
    input_size = 224
    mean = [0.485, 0.456, 0.406]
    std=[0.229, 0.224, 0.225]

    transform = video.VideoGroupValTransform(size=input_size, mean=mean, std=std)
    video_utils = VideoClsCustom(root=data_dir,
                                 setting=video_list_path,
                                 num_segments=num_segments,
                                 new_length=new_length,
                                 new_step=new_step,
                                 video_loader=video_loader,
                                 use_decord=use_decord,
                                 slowfast=slowfast)

    #Read for the video list
    video_name = video_list.split()[0]

    decord = try_import_decord()
    decord_vr = decord.VideoReader(video_name)
    duration = len(decord_vr)

    skip_length = new_length * new_step
    segment_indices, skip_offsets = video_utils._sample_test_indices(duration)

    if video_loader:
        if slowfast:
            clip_input = video_utils._video_TSN_decord_slowfast_loader(video_name, decord_vr,
                                                                       duration, segment_indices, skip_offsets)
        else:
            clip_input = video_utils._video_TSN_decord_batch_loader(video_name, decord_vr,
                                                                    duration, segment_indices, skip_offsets)
    else:
        raise RuntimeError('We only support video-based inference.')

    clip_input = transform(clip_input)

    if slowfast:
        sparse_sampels = len(clip_input) // (num_segments * num_crop)
        clip_input = np.stack(clip_input, axis=0)
        clip_input = clip_input.reshape((-1,) + (sparse_sampels, 3, input_size, input_size))
        clip_input = np.transpose(clip_input, (0, 2, 1, 3, 4))
    else:
        clip_input = np.stack(clip_input, axis=0)
        clip_input = clip_input.reshape((-1,) + (new_length, 3, input_size, input_size))
        clip_input = np.transpose(clip_input, (0, 2, 1, 3, 4))

    if new_length == 1:
        clip_input = np.squeeze(clip_input, axis=2)    # this is for 2D input case

    clip_input = nd.array(clip_input)

    #Cleanup temp files
    os.remove(download_path)
    os.remove(video_list_path)
    #os.system('rm {}'.format(download_path))
    #os.system('rm {}'.format(video_list_path))

    return clip_input


import json
import os
import pathlib
import pickle
from typing import List, Tuple

import numpy as np
import tensorflow as tf
from tensorflow.keras import models
from tensorflow.python.keras.backend import set_session
from tensorflow.python.keras.preprocessing.sequence import pad_sequences

from credsweeper.common.constants import ThresholdPreset
from credsweeper.credentials import Candidate
from credsweeper.credentials.line_data import LineData
from credsweeper.logger.logger import logging
from credsweeper.ml_model import features

class MlValidator:
    @classmethod
    def __init__(cls, threshold_preset: ThresholdPreset = ThresholdPreset.balanced) -> None:
        tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)  # To make TF logger quiet
        config = tf.compat.v1.ConfigProto()
        config.gpu_options.allow_growth = True  # dynamically grow the memory used on the GPU
        config.log_device_placement = True  # to log device placement (on which device the operation ran)
        sess = tf.compat.v1.Session(config=config)
        set_session(sess)
        dir_path = os.path.dirname(os.path.realpath(__file__))
        model_file_path = os.path.join(dir_path, "ml_model.h5")
        index_file_path = os.path.join(dir_path, "char_to_index.pkl")
        cls.model = models.load_model(model_file_path)
        with open(index_file_path, "rb") as index_file:
            cls.char_to_index = pickle.load(index_file)

        model_detail_path = f"{pathlib.Path(__file__).parent.absolute()}/model_config.json"
        with open(model_detail_path) as f:
            model_details = json.load(f)
        if "thresholds" in model_details:
            cls.threshold = model_details["thresholds"][threshold_preset.value]
        else:
            cls.threshold = 0.5
        cls.maxlen = model_details.get("max_len", 50)
        cls.common_feature_list = []
        cls.unique_feature_list = []
        logging.info(f'Init ML validator, model file path: {model_file_path} \tindex file path: {index_file_path}')
        logging.debug(f'ML validator details: {model_details}')
        for feature_definition in model_details["features"]:
            feature_class = feature_definition["type"]
            kwargs = feature_definition.get("kwargs", {})
            feature_constructor = getattr(features, feature_class, None)
            if feature_constructor is None:
                raise ValueError(f'Error while parsing model details. Cannot create feature "{feature_class}"')
            try:
                feature = feature_constructor(**kwargs)
            except TypeError:
                raise TypeError(
                    f'Error while parsing model details. Cannot create feature "{feature_class}" with kwargs "{kwargs}"'
                )
            if feature_definition["type"] in ["RuleName"]:
                cls.unique_feature_list.append(feature)
            else:
                cls.common_feature_list.append(feature)

    @classmethod
    def encode(cls, line, char_to_index) -> np.ndarray:
        encoded = []
        for c in line:
            if c in char_to_index:
                encoded.append(char_to_index[c])
            else:
                encoded.append(char_to_index['NON_ASCII'])
        return pad_sequences([encoded], padding='post', maxlen=cls.maxlen)

    @classmethod
    def extract_common_features(cls, candidates: List[Candidate]) -> np.ndarray:
        """Extract features that are guaranteed to be the same for all candidates on the same line with same value"""
        features = np.array([], dtype=float)
        # Extract features from credential candidate
        default_candidate = candidates[0]
        for feature in cls.common_feature_list:
            new_feature = feature([default_candidate])[0]
            if not isinstance(new_feature, np.ndarray):
                new_feature = np.array([new_feature])
            features = np.append(features, new_feature)
        return features

    @classmethod
    def extract_unique_features(cls, candidates: List[Candidate]) -> np.ndarray:
        """Extract features that can by different between candidates. Join them with or operator"""
        features = np.array([], dtype=bool)
        default_candidate = candidates[0]
        for feature in cls.unique_feature_list:
            new_feature = feature([default_candidate])[0]
            if not isinstance(new_feature, np.ndarray):
                new_feature = np.array([new_feature])
            features = np.append(features, new_feature)
        for candidate in candidates[1:]:
            for feature in cls.unique_feature_list:
                new_feature = feature([candidate])[0]
                if not isinstance(new_feature, np.ndarray):
                    new_feature = np.array([new_feature])
                features = features | new_feature
        return features

    @classmethod
    def validate(cls, line_data: LineData, candidate: Candidate) -> bool:
        sample_as_batch = [(line_data.value, [candidate])]
        is_cred_batch = cls.validate_groups(sample_as_batch, 1)
        return is_cred_batch[0]

    @classmethod
    def get_group_features(cls, value: str, candidates: List[Candidate]) -> Tuple[np.ndarray, np.ndarray]:
        line_input = np.array(cls.encode(value, cls.char_to_index)).reshape(1, -1)

        common_features = cls.extract_common_features(candidates)
        unique_features = cls.extract_unique_features(candidates)
        features = np.hstack([common_features, unique_features])
        features = np.array([features])
        return line_input, features

    @classmethod
    def validate_groups(cls, group_list: List[Tuple[str, List[Candidate]]], batch_size: int) -> np.ndarray:
        """Use ml model on list of candidate groups

         Args:
            group_list: List of tuples (value, group)
            batch_size: ML model batch

        Return:
            Numpy array with same length as group_list
        """
        line_input_list = []
        features_list = []
        for (value, candidates) in group_list:
            line_input, features = cls.get_group_features(value, candidates)
            line_input_list.append(line_input)
            features_list.append(features)

        pred = np.zeros(len(features_list))
        for i in range(0, len(features_list), batch_size):
            line_inputs = line_input_list[i:i + batch_size]
            line_inputs = np.vstack(line_inputs)
            features = features_list[i:i + batch_size]
            features = np.vstack(features)
            pred[i:i + batch_size] = cls.model([line_inputs, features])[:, 0]
        is_cred = pred > cls.threshold
        for i in range(len(is_cred)):
            logging.debug(
                f"ML decision: {is_cred[i]} with prediction: {round(pred[i], 3)} for value: {group_list[i][0]}")
        return is_cred


'''
Entangled: The world's first Quantum dating algorithm!
By Alex Breitweiser for PennApps XVI, Fall 2017
'''

###IBM QX Imports###

# Checking the version of PYTHON; we only support 3 at the moment
import sys
if sys.version_info < (3,0):
    raise Exception('Please use Python version 3 or greater.')

# useful additional packages
import matplotlib.pyplot as plt
#%matplotlib inline
import numpy as np

import sys
sys.path.append('/Users/sabreitweiser/qiskit-sdk-py/') #Change to point to your qiskit root

# importing the QISKit
from qiskit import QuantumCircuit, QuantumProgram
import Qconfig

# import basic plot tools
from qiskit.tools.visualization import plot_histogram

###Tinder API imports###
sys.path.append('/Users/sabreitweiser/Tinder/') #Change to point to your Tinder api root
from features import get_match_info, pause
from tinder_api import authverif, like, dislike, get_recommendations, get_person

###My imports###
from queue import Queue
from pprint import pprint
from urllib.request import urlopen
from PIL import Image
from time import sleep
from datetime import date
from random import random

###Helper code###
class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

###Setting up Tinder###
authverif()

###Setting up IBMQX###
device = 'ibmqx2' # the device to run on
device = 'local_qasm_simulator' # uncomment to run on the simulator
batch_size = 50 # Number of decisions in one request

QPS_SPECS = {
    "name": "decision",
    "circuits": [{
        "name": "decision_circuit",
        "quantum_registers": [{
            "name":"q",
            "size":1
        }],
        "classical_registers": [{
            "name":"c",
            "size":1
        }]}]
}

Q_program = QuantumProgram(specs=QPS_SPECS)
Q_program.set_api(Qconfig.APItoken, Qconfig.config["url"])

# Quantum circuits to generate decisions
circuits = ["decision"+str(i) for i in range(batch_size)]
# NB: Can't have more than one measurement per circuit
for circuit in circuits:
    q = Q_program.get_quantum_register("q")
    c = Q_program.get_classical_register('c')
    decision_circuit = Q_program.create_circuit(circuit, [q], [c])
    decision_circuit.h(q[0]) #Turn the qubit into |0> + |1>
    decision_circuit.measure(q[0], c[0])
_ = Q_program.get_qasms(circuits) # Suppress the output

###Main Loop###
recommendations = Queue() #Queue to hold tinder recommendations
decisions = Queue() #Queue to hold QX decisions

while True:
	#Replenish the queues
	if recommendations.empty():
		results = get_recommendations()
		results = results['results']
		for key in results:
			recommendations.put(key['_id'])

	if decisions.empty():
		results = Q_program.execute(circuits, device, shots=1,
						max_credits=5, wait=10,
						timeout=240)
		for ciruit in circuits:
			for key in results.get_counts(ciruit):
				decisions.put(int(key))

	#Match a decision to a recommendation
	rec = recommendations.get()
	dec = decisions.get()

	print('ID: ' + rec)
	person = get_person(rec)
	person = person['results']
	print('NAME: ' + person['name'])
	if person.get('jobs'):
		for job in person['jobs']:
			if job.get('company'):
				print('COMPANY: ' + job['company']['name'])
			if job.get('title'):
				print('TITLE: ' + job['title']['name'])

	if person.get('schools'):
		for school in person['schools']:
			print('SCHOOL: ' + school['name'])

	if person.get('bio'):
		print('BIO: ' + person['bio'])

	if person.get('distance_mi'):
		print('DISTANCE: ' + str(person['distance_mi']) + ' miles')

	for photo in person['photos']:
		URL = photo['url']
		with urlopen(URL) as url:
			with open('temp.jpg', 'wb') as f:
				f.write(url.read())

		img = Image.open('temp.jpg')
		img.show()

	if dec:
		like(rec)
		print(bcolors.OKGREEN + 'LIKED' + bcolors.ENDC)
	else:
		dislike(rec)
		print(bcolors.FAIL + 'DISLIKED' + bcolors.ENDC)

	for _ in range(5):
		sleep(4*random())
		print('...')


#!python3
# -*- coding: utf-8 -*-
# pylint: disable=W0312, line-too-long, C0103
# pylama: noqa
# flake8: noqa

import sys
import time

sys.path.append('waveglow-mirror/')

import librosa
import numpy as np
import torch

from audio_processing import griffin_lim
from hparams import create_hparams
from layers import STFT, TacotronSTFT
from model import Tacotron2
from text import text_to_sequence
from train import load_model

torch.set_num_threads(4)

hparams = create_hparams()
hparams.sampling_rate = 22050

checkpoint_path = "tacotron2_statedict.pt"
model = load_model(hparams)
model.load_state_dict(torch.load(checkpoint_path, map_location='cpu')['state_dict'])
_ = model.cpu().eval()

waveglow_path = 'waveglow_256channels.pt'
waveglow = torch.load(waveglow_path, map_location='cpu')['model']
waveglow.cpu().eval()
# for k in waveglow.convinv:
#    k.float()

for m in waveglow.modules():
    if 'Conv' in str(type(m)):
        setattr(m, 'padding_mode', 'zeros')
        # print(m)

def generate_texts(texts):
    audios = None

    for text in texts:
        start = time.time()
        print("Calculating %s" % (text[:10]))
        sequence = np.array(text_to_sequence(text, ['english_cleaners']))[None, :]
        sequence = torch.autograd.Variable(torch.from_numpy(sequence)).cpu().long()

        mel_outputs, mel_outputs_postnet, _, alignments = model.inference(sequence)
        tacotron_time = time.time() - start
        start = time.time()

        with torch.no_grad():
            audio = waveglow.infer(mel_outputs_postnet, sigma=0.666)
        waveglow_time = time.time() - start

        if audios is not None:
            audios = np.append(audios, audio[0].data.cpu().numpy())
        else:
            audios = audio[0].data.cpu().numpy()

        audio_len = len(audio[0]) / 22050

        print("Total calculation time %.2lf, Length %.2lf, Tacotron: %.2lfx-realtime, Waveglow: %.2lfx-realtime" % (tacotron_time + waveglow_time, audio_len, audio_len / tacotron_time, audio_len / waveglow_time))

    return audios

texts = ["A very warm welcome to my dear colleagues from Original Equipment Innovation.",
         "I am a new text to speech engine that is perceived to be really natural.",
         "I am based upon the so-called Tacotron Two model and a Waveglow synthesizer.",
         "I, hereby, apply as your new voice for nice chat!",
         "Currently I'm running offline on Alex's computer on an Intel Core i5 CPU, but I'm quite confident that I can also run on your smaller aircraft N U Cs.",
         "Alex Yasha and Eyekey will be happy to tell you more about me and my current capabilities!.",
         "Thank you!"]
audios = generate_texts(texts)

librosa.output.write_wav("output.wav", audios, hparams.sampling_rate)


from PIL import ImageFilter, Image, ImageDraw, ImageOps
from .Texturer import Texturer
import numpy

class Artifacter:
  def __init__(self):
    self.predefined_overlays = [
      (0, 0, 0, 255),
      (242, 246, 208, 255),
      (143, 191, 224, 255),
      (255, 255, 255, 255)
    ]
    self.texturer = Texturer()
    pass

  def blur(self, canvas, blur_radius=1):
    blur_radius = numpy.random.random() * blur_radius
    return canvas.filter(ImageFilter.GaussianBlur(blur_radius))

  def rotate(self, canvas, max_degree=4):
    size = canvas.size
    sign = numpy.random.choice([-1, 1])
    degree = sign * numpy.random.random() * max_degree
    return canvas.rotate(degree, expand=1).resize(size)

  def add_overlay(self, canvas):
    color_index = numpy.random.randint(len(self.predefined_overlays))
    color = self.predefined_overlays[color_index]
    overlay = Image.new('RGBA', canvas.size, color)
    alpha = numpy.random.random() * 0.08
    return Image.blend(canvas, overlay, alpha)

  def add_texture(self, canvas, probability=0.7):
    return self.texturer.texture(canvas, 0.7)


import numpy as np

def map_17_to_16(joint_17):
    joint_16 = np.zeros((16, 3), dtype=np.float32)
    dict_map = {0:16, 1:14, 2:12, 3:11, 4:13, 5:15, 6:[11, 12], 7:[5,6], 9:[1,2], 10:10, 11:8, 12:6, 13:5, 14:7, 15:9}
    for idx in range(16):
        if idx == 8:
            continue # deal thrx joint later
        elif idx in [6, 7, 9]:
            #calc Pelv joint from two hip joint
            #calc neck joint from two shoulder joint
            #calc head joint from two eye joint
            joint_16[idx, 0] = (joint_17[dict_map[idx][0], 0] + joint_17[dict_map[idx][1], 0]) * 0.5
            joint_16[idx, 1] = (joint_17[dict_map[idx][0], 1] + joint_17[dict_map[idx][1], 1]) * 0.5
            joint_16[idx, 2] = (joint_17[dict_map[idx][0], 2] + joint_17[dict_map[idx][1], 2]) * 0.5
        else:
            joint_16[idx] = joint_17[dict_map[idx]]
    #calc thrx joint from head joint and neck, assume the distance is 1:3
    joint_16[8, 0] = joint_16[7, 0] * 0.75  + joint_16[9, 0] * 0.25
    joint_16[8, 1] = joint_16[7, 1] * 0.75  + joint_16[9, 1] * 0.25
    joint_16[8, 2] = joint_16[7, 2] * 0.75  + joint_16[9, 2] * 0.25

    return joint_16

def output_part_box(joint, img_bbox):

    flag_bad_joint = 0

    # 16 part names correspond to the center of 16 input joint
    part_size = [1.2, 1, 1, 1.2, 1.2, 1.2, 0.9, 1, 1, 0.9]
    part = [0, 1, 4, 5, 6, 9, 10, 12, 13, 15]

    height = get_distance(joint, 6, 8)

    if (joint[8, 2] < 0.2) or (joint[6, 2] < 0.2):
        flag_bad_joint = 1

    group_score_head      = (joint[7, 2] + joint[8, 2] + joint[9, 2]) / 3
    group_score_left_arm  = (joint[13, 2] + joint[14, 2] + joint[15, 2]) / 3
    group_score_right_arm = (joint[10, 2] + joint[11, 2] + joint[12, 2]) / 3
    group_score_left_leg  = (joint[3, 2] + joint[4, 2] + joint[5, 2]) / 3
    group_score_right_leg = (joint[0, 2] + joint[1, 2] + joint[2, 2]) / 3

    # 'Pelv'&'Neck' scaling by the distance of Pelv and Neck
    bbox = np.zeros((10, 4), dtype=np.float32)
    for i in range(10):
        score_joint = joint[part[i], 2]

        # the keypoint is not reliable/ cannot be seen / do not exist
        if (score_joint < 0.2):
            bbox[i, 0] = img_bbox[0]
            bbox[i, 1] = img_bbox[1]
            bbox[i, 2] = img_bbox[2]
            bbox[i, 3] = img_bbox[3]

        # the keypoint is reliable, but the distance cannot be measured by distance between pelv and neck
        elif (score_joint >= 0.2) and (flag_bad_joint == 1):
            if i == 5: # head group
                if group_score_head > 0.2:
                    half_box_width = get_distance(joint, 7, 9)
                    pbox = get_part_box(i, joint, half_box_width)
                    bbox[i, 0] = pbox[0]
                    bbox[i, 1] = pbox[1]
                    bbox[i, 2] = pbox[2]
                    bbox[i, 3] = pbox[3]
                else:
                    bbox[i, 0] = img_bbox[0]
                    bbox[i, 1] = img_bbox[1]
                    bbox[i, 2] = img_bbox[2]
                    bbox[i, 3] = img_bbox[3]
            elif i in [6,7]: # right arm group
                if group_score_right_arm > 0.2:
                    half_box_width = get_distance(joint, 10, 12)
                    pbox = get_part_box(i, joint, half_box_width)
                    bbox[i, 0] = pbox[0]
                    bbox[i, 1] = pbox[1]
                    bbox[i, 2] = pbox[2]
                    bbox[i, 3] = pbox[3]
                else:
                    bbox[i, 0] = img_bbox[0]
                    bbox[i, 1] = img_bbox[1]
                    bbox[i, 2] = img_bbox[2]
                    bbox[i, 3] = img_bbox[3]
            elif i in [8,9]: # left arm group
                if group_score_left_arm > 0.2:
                    half_box_width = get_distance(joint, 13, 15)
                    pbox = get_part_box(i, joint, half_box_width)
                    bbox[i, 0] = pbox[0]
                    bbox[i, 1] = pbox[1]
                    bbox[i, 2] = pbox[2]
                    bbox[i, 3] = pbox[3]
                else:
                    bbox[i, 0] = img_bbox[0]
                    bbox[i, 1] = img_bbox[1]
                    bbox[i, 2] = img_bbox[2]
                    bbox[i, 3] = img_bbox[3]
            elif i in [0,1]: # right leg group
                if group_score_right_leg > 0.2:
                    half_box_width = get_distance(joint, 0, 2)
                    pbox = get_part_box(i, joint, half_box_width)
                    bbox[i, 0] = pbox[0]
                    bbox[i, 1] = pbox[1]
                    bbox[i, 2] = pbox[2]
                    bbox[i, 3] = pbox[3]
                else:
                    bbox[i, 0] = img_bbox[0]
                    bbox[i, 1] = img_bbox[1]
                    bbox[i, 2] = img_bbox[2]
                    bbox[i, 3] = img_bbox[3]
            elif i in [2,3]: # left leg group
                if group_score_left_leg > 0.2:
                    half_box_width = get_distance(joint, 3, 5)
                    pbox = get_part_box(i, joint, half_box_width)
                    bbox[i, 0] = pbox[0]
                    bbox[i, 1] = pbox[1]
                    bbox[i, 2] = pbox[2]
                    bbox[i, 3] = pbox[3]
                else:
                    bbox[i, 0] = img_bbox[0]
                    bbox[i, 1] = img_bbox[1]
                    bbox[i, 2] = img_bbox[2]
                    bbox[i, 3] = img_bbox[3]
            else: # pelv keypoint
                bbox[i, 0] = img_bbox[0]
                bbox[i, 1] = img_bbox[1]
                bbox[i, 2] = img_bbox[2]
                bbox[i, 3] = img_bbox[3]

        else: # the keypoint is reliable and the distance can be measured by distance between pelv and neck
            half_box_width = height * part_size[i] / 2
            pbox = get_part_box(i, joint, half_box_width)
            bbox[i, 0] = pbox[0]
            bbox[i, 1] = pbox[1]
            bbox[i, 2] = pbox[2]
            bbox[i, 3] = pbox[3]
    return np.concatenate([np.zeros((10, 1)), bbox], axis=1)

def get_part_box(i, joint, half_box_width):
    part = [0, 1, 4, 5, 6, 9, 10, 12, 13, 15]
    center_x = joint[part[i], 0]
    center_y = joint[part[i], 1]
    return center_x - half_box_width, center_y - half_box_width, center_x + half_box_width, center_y + half_box_width

def get_distance(joint, keypoint1, keypoint2):
    height_y = joint[keypoint1, 1] - joint[keypoint2, 1]
    height_x = joint[keypoint1, 0] - joint[keypoint2, 0]
    return np.sqrt(height_x ** 2 + height_y ** 2)

def check_iou(human_bbox_pkl, human_bbox_pose):

    x1, y1, x2, y2 = human_bbox_pose
    x1d, y1d, x2d, y2d = human_bbox_pkl

    xa = max(x1, x1d)
    ya = max(y1, y1d)
    xb = min(x2, x2d)
    yb = min(y2, y2d)

    iw1 = xb - xa + 1
    iw2 = yb - ya + 1

    if iw1 > 0 and iw2 > 0:
        inter_area = iw1 * iw2
        a_area = (x2 - x1) * (y2 - y1)
        b_area = (x2d - x1d) * (y2d - y1d)
        union_area = a_area + b_area - inter_area
        return inter_area / float(union_area)
    else:
        return 0

if __name__ == '__main__':
    # Given the human bounding box and the corresponding human pose, Generate the correponding part box
    hbox = np.zeros((4))
    joint_17 = np.zeros((17, 3))               # the pose
    joint_16 = map_17_to_16(joint_17)
    part_box = output_part_box(joint_16, hbox) # return the part bounding box


"""
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
from keras.models import Model
from keras.layers import Input, Lambda, Activation, Conv2D, \
    DepthwiseConv2D, Reshape, Concatenate, BatchNormalization, ReLU
import keras.backend as K
from layers.AnchorBoxesLayer import AnchorBoxes
from layers.DecodeDetectionsLayer import DecodeDetections
from layers.DecodeDetectionsFastLayer import DecodeDetectionsFast
from models.graphs.peleenet_graph import peleenet_graph, res_block_graph, extend_conv2d_graph

def predict_block(inputs, out_channel, sym, id):
    name = 'ssd_' + sym + '{}'.format(id)
    x = DepthwiseConv2D(kernel_size=3, strides=1,
                           activation=None, use_bias=False, padding='same', name=name + '_dw_conv')(inputs)
    x = BatchNormalization(epsilon=1e-3, momentum=0.999, name=name + '_dw_bn')(x)
    x = ReLU(6., name=name + '_dw_relu')(x)

    x = Conv2D(out_channel, kernel_size=1, padding='same', use_bias=False,
                  activation=None, name=name + 'conv2')(x)
    x = BatchNormalization(epsilon=1e-3, momentum=0.999, name=name + 'conv2_bn')(x)
    return x

def pelee(image_size,
            n_classes,
            mode='training',
            l2_regularization=0.0005,
            min_scale=None,
            max_scale=None,
            scales=None,
            aspect_ratios_global=None,
            aspect_ratios_per_layer=[[1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
                                     [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
                                     [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
                                     [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
                                     [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0]],
            two_boxes_for_ar1=True,
            steps=[8, 16, 32, 64, 100, 300],
            offsets=None,
            clip_boxes=False,
            variances=[0.1, 0.1, 0.2, 0.2],
            coords='centroids',
            normalize_coords=True,
            subtract_mean=[123, 117, 104],
            divide_by_stddev=None,
            swap_channels=[2, 1, 0],
            confidence_thresh=0.01,
            iou_threshold=0.45,
            top_k=200,
            nms_max_output_size=400,
            return_predictor_sizes=False):
    """
    Build a Keras model with SSD300 architecture, see references.
    The base network is a reduced atrous VGG-16, extended by the SSD architecture,
    as described in the paper.
    Most of the arguments that this function takes are only needed for the anchor
    box layers. In case you're training the network, the parameters passed here must
    be the same as the ones used to set up `SSDBoxEncoder`. In case you're loading
    trained weights, the parameters passed here must be the same as the ones used
    to produce the trained weights.
    Some of these arguments are explained in more detail in the documentation of the
    `SSDBoxEncoder` class.
    Note: Requires Keras v2.0 or later. Currently works only with the
    TensorFlow backend (v1.0 or later).
    Arguments:
        image_size (tuple): The input image size in the format `(height, width, channels)`.
        n_classes (int): The number of positive classes, e.g. 20 for Pascal VOC, 80 for MS COCO.
        mode (str, optional): One of 'training', 'inference' and 'inference_fast'. In 'training' mode,
            the model outputs the raw prediction tensor, while in 'inference' and 'inference_fast' modes,
            the raw predictions are decoded into absolute coordinates and filtered via confidence thresholding,
            non-maximum suppression, and top-k filtering. The difference between latter two modes is that
            'inference' follows the exact procedure of the original Caffe implementation, while
            'inference_fast' uses a faster prediction decoding procedure.
        l2_regularization (float, optional): The L2-regularization rate. Applies to all convolutional layers.
            Set to zero to deactivate L2-regularization.
        min_scale (float, optional): The smallest scaling factor for the size of the anchor boxes as a fraction
            of the shorter side of the input images.
        max_scale (float, optional): The largest scaling factor for the size of the anchor boxes as a fraction
            of the shorter side of the input images. All scaling factors between the smallest and the
            largest will be linearly interpolated. Note that the second to last of the linearly interpolated
            scaling factors will actually be the scaling factor for the last predictor layer, while the last
            scaling factor is used for the second box for aspect ratio 1 in the last predictor layer
            if `two_boxes_for_ar1` is `True`.
        scales (list, optional): A list of floats containing scaling factors per convolutional predictor layer.
            This list must be one element longer than the number of predictor layers. The first `k` elements are the
            scaling factors for the `k` predictor layers, while the last element is used for the second box
            for aspect ratio 1 in the last predictor layer if `two_boxes_for_ar1` is `True`. This additional
            last scaling factor must be passed either way, even if it is not being used. If a list is passed,
            this argument overrides `min_scale` and `max_scale`. All scaling factors must be greater than zero.
        aspect_ratios_global (list, optional): The list of aspect ratios for which anchor boxes are to be
            generated. This list is valid for all prediction layers.
        aspect_ratios_per_layer (list, optional): A list containing one aspect ratio list for each prediction layer.
            This allows you to set the aspect ratios for each predictor layer individually, which is the case for the
            original SSD300 implementation. If a list is passed, it overrides `aspect_ratios_global`.
        two_boxes_for_ar1 (bool, optional): Only relevant for aspect ratio lists that contain 1. Will be ignored
            otherwise. If `True`, two anchor boxes will be generated for aspect ratio 1. The first will be generated
            using the scaling factor for the respective layer, the second one will be generated using
            geometric mean of said scaling factor and next bigger scaling factor.
        steps (list, optional): `None` or a list with as many elements as there are predictor layers. The elements
            can be either ints/floats or tuples of two ints/floats. These numbers represent for each predictor layer
            how many pixels apart the anchor box center points should be vertically and horizontally along the spatial
            grid over the image. If the list contains ints/floats, then that value will be used for both spatial
            dimensions. If the list contains tuples of two ints/floats, then they represent `(step_height, step_width)`.
            If no steps are provided, then they will be computed such that the anchor box center points will form an
            equidistant grid within the image dimensions.
        offsets (list, optional): `None` or a list with as many elements as there are predictor layers. The elements
            can be either floats or tuples of two floats. These numbers represent for each predictor layer how many
            pixels from the top and left boarders of the image the top-most and left-most anchor box center points
            should be as a fraction of `steps`. The last bit is important: The offsets are not absolute pixel values,
            but fractions of the step size specified in the `steps` argument. If the list contains floats, then that
            value will be used for both spatial dimensions. If the list contains tuples of two floats, then they
            represent `(vertical_offset, horizontal_offset)`. If no offsets are provided, then they will default
            to 0.5 of the step size.
        clip_boxes (bool, optional): If `True`, clips the anchor box coordinates to stay within image boundaries.
        variances (list, optional): A list of 4 floats >0. The anchor box offset for each coordinate will be divided by
            its respective variance value.
        coords (str, optional): The box coordinate format to be used internally by the model (i.e. this is not the
            input format of the ground truth labels). Can be either 'centroids' for the format `(cx, cy, w, h)`
            (box center coordinates, width, and height), 'minmax' for the format `(xmin, xmax, ymin, ymax)`,
            or 'corners' for the format `(xmin, ymin, xmax, ymax)`.
        normalize_coords (bool, optional): Set to `True` if the model is supposed to use relative instead of
            absolute coordinates, i.e. if the model predicts box coordinates within [0,1] instead of absolute
            coordinates.
        subtract_mean (array-like, optional): `None` or an array-like object of integers or floating point values
            of any shape that is broadcast-compatible with the image shape. The elements of this array will be
            subtracted from the image pixel intensity values. For example, pass a list of three integers
            to perform per-channel mean normalization for color images.
        divide_by_stddev (array-like, optional): `None` or an array-like object of non-zero integers or
            floating point values of any shape that is broadcast-compatible with the image shape. The image pixel
            intensity values will be divided by the elements of this array. For example, pass a list
            of three integers to perform per-channel standard deviation normalization for color images.
        swap_channels (list, optional): Either `False` or a list of integers representing the desired order in which
            the input image channels should be swapped.
        confidence_thresh (float, optional): A float in [0,1), the minimum classification confidence in a specific
            positive class in order to be considered for the non-maximum suppression stage for the respective class.
            A lower value will result in a larger part of the selection process being done by the non-maximum
            suppression stage, while a larger value will result in a larger part of the selection process happening
            in the confidence thresholding stage.
        iou_threshold (float, optional): A float in [0,1]. All boxes that have a Jaccard similarity of greater
            than `iou_threshold` with a locally maximal box will be removed from the set of predictions for a given
            class, where 'maximal' refers to the box's confidence score.
        top_k (int, optional): The number of highest scoring predictions to be kept for each batch item after the
            non-maximum suppression stage.
        nms_max_output_size (int, optional): The maximal number of predictions that will be left over after the
            NMS stage.
        return_predictor_sizes (bool, optional): If `True`, this function not only returns the model, but also
            a list containing the spatial dimensions of the predictor layers. This isn't strictly necessary since
            you can always get their sizes easily via the Keras API, but it's convenient and less error-prone
            to get them this way. They are only relevant for training anyway (SSDBoxEncoder needs to know the
            spatial dimensions of the predictor layers), for inference you don't need them.
    Returns:
        model: The Keras SSD300 model.
        predictor_sizes (optional): A Numpy array containing the `(height, width)` portion
            of the output tensor shape for each convolutional predictor layer. During
            training, the generator function needs this in order to transform
            the ground truth labels into tensors of identical structure as the
            output tensors of the model, which is in turn needed for the cost
            function.
    References:
        https://arxiv.org/abs/1512.02325v5
    """

    n_predictor_layers = 5  # The number of predictor conv layers in the network is 6 for the original SSD300.
    n_classes += 1  # Account for the background class.
    l2_reg = l2_regularization  # Make the internal name shorter.
    img_height, img_width, img_channels = image_size[0], image_size[1], image_size[2]

    ############################################################################
    # Get a few exceptions out of the way.
    ############################################################################

    if aspect_ratios_global is None and aspect_ratios_per_layer is None:
        raise ValueError(
            "`aspect_ratios_global` and `aspect_ratios_per_layer` \
            cannot both be None. At least one needs to be specified.")
    if aspect_ratios_per_layer:
        if len(aspect_ratios_per_layer) != n_predictor_layers:
            raise ValueError(
                "It must be either aspect_ratios_per_layer is None or len(aspect_ratios_per_layer) == {}, \
                but len(aspect_ratios_per_layer) == {}.".format(
                    n_predictor_layers, len(aspect_ratios_per_layer)))

    if (min_scale is None or max_scale is None) and scales is None:
        raise ValueError("Either `min_scale` and `max_scale` or `scales` need to be specified.")
    if scales:
        if len(scales) != n_predictor_layers + 1:
            raise ValueError("It must be either scales is None or len(scales) == {}, but len(scales) == {}.".format(
                n_predictor_layers + 1, len(scales)))
    # If no explicit list of scaling factors was passed,
    # compute the list of scaling factors from `min_scale` and `max_scale`
    else:
        scales = np.linspace(min_scale, max_scale, n_predictor_layers + 1)

    if len(variances) != 4:
        raise ValueError("4 variance values must be pased, but {} values were received.".format(len(variances)))
    variances = np.array(variances)
    if np.any(variances <= 0):
        raise ValueError("All variances must be >0, but the variances given are {}".format(variances))

    if (not (steps is None)) and (len(steps) != n_predictor_layers):
        raise ValueError("You must provide at least one step value per predictor layer.")

    if (not (offsets is None)) and (len(offsets) != n_predictor_layers):
        raise ValueError("You must provide at least one offset value per predictor layer.")

    ############################################################################
    # Compute the anchor box parameters.
    ############################################################################

    # Set the aspect ratios for each predictor layer. These are only needed for the anchor box layers.
    if aspect_ratios_per_layer:
        aspect_ratios = aspect_ratios_per_layer
    else:
        aspect_ratios = [aspect_ratios_global] * n_predictor_layers

    # Compute the number of boxes to be predicted per cell for each predictor layer.
    # We need this so that we know how many channels the predictor layers need to have.
    if aspect_ratios_per_layer:
        n_boxes = []
        for ar in aspect_ratios_per_layer:
            if (1 in ar) & two_boxes_for_ar1:
                n_boxes.append(len(ar) + 1)  # +1 for the second box for aspect ratio 1
            else:
                n_boxes.append(len(ar))
    # If only a global aspect ratio list was passed,
    # then the number of boxes is the same for each predictor layer
    else:
        if (1 in aspect_ratios_global) & two_boxes_for_ar1:
            n_boxes = len(aspect_ratios_global) + 1
        else:
            n_boxes = len(aspect_ratios_global)
        n_boxes = [n_boxes] * n_predictor_layers

    if steps is None:
        steps = [None] * n_predictor_layers
    if offsets is None:
        offsets = [None] * n_predictor_layers

    ############################################################################
    # Define functions for the Lambda layers below.
    ############################################################################

    def identity_layer(tensor):
        return tensor

    def input_mean_normalization(tensor):
        return tensor - np.array(subtract_mean)

    def input_stddev_normalization(tensor):
        return tensor / np.array(divide_by_stddev)

    def input_channel_swap(tensor):
        if len(swap_channels) == 3:
            return K.stack(
                [tensor[..., swap_channels[0]], tensor[..., swap_channels[1]], tensor[..., swap_channels[2]]], axis=-1)
        elif len(swap_channels) == 4:
            return K.stack([tensor[..., swap_channels[0]], tensor[..., swap_channels[1]], tensor[..., swap_channels[2]],
                            tensor[..., swap_channels[3]]], axis=-1)

    ############################################################################
    # Build the network.
    ############################################################################

    x = Input(shape=(img_height, img_width, img_channels))

    # The following identity layer is only needed so that the subsequent lambda layers can be optional.
    x1 = Lambda(identity_layer, output_shape=(img_height, img_width, img_channels), name='identity_layer')(x)

    if not (subtract_mean is None):
        x1 = Lambda(input_mean_normalization, output_shape=(img_height, img_width, img_channels),
                    name='input_mean_normalization')(x1)
    if not (divide_by_stddev is None):
        x1 = Lambda(input_stddev_normalization, output_shape=(img_height, img_width, img_channels),
                    name='input_stddev_normalization')(x1)
    if swap_channels:
        x1 = Lambda(
            input_channel_swap, output_shape=(img_height, img_width, img_channels), name='input_channel_swap')(x1)

    link1, link2 = peleenet_graph(x1)

    link3 = extend_conv2d_graph(
        link2, out_channels=256, mid_chennel=256, strides=2, padding='same', name='ssd_feature_extend1')
    link4 = extend_conv2d_graph(
        link3, out_channels=256, mid_chennel=128, strides=1, padding='valid', name='ssd_feature_extend2')
    link5 = extend_conv2d_graph(
        link4, out_channels=256, mid_chennel=128, strides=1, padding='valid', name='ssd_feature_extend3')

    link1 = res_block_graph(link1, name='ssd_res_block1')
    link2 = res_block_graph(link2, name='ssd_res_block2')
    link3 = res_block_graph(link3, name='ssd_res_block3')
    link4 = res_block_graph(link4, name='ssd_res_block4')
    link5 = res_block_graph(link5, name='ssd_res_block5')

    link1_cls = predict_block(link1, n_boxes[0] * n_classes, 'cls', 1)
    link2_cls = predict_block(link2, n_boxes[1] * n_classes, 'cls', 2)
    link3_cls = predict_block(link3, n_boxes[2] * n_classes, 'cls', 3)
    link4_cls = predict_block(link4, n_boxes[3] * n_classes, 'cls', 4)
    link5_cls = predict_block(link5, n_boxes[4] * n_classes, 'cls', 5)

    link1_box = predict_block(link1, n_boxes[0] * 4, 'box', 1)
    link2_box = predict_block(link2, n_boxes[1] * 4, 'box', 2)
    link3_box = predict_block(link3, n_boxes[2] * 4, 'box', 3)
    link4_box = predict_block(link4, n_boxes[3] * 4, 'box', 4)
    link5_box = predict_block(link5, n_boxes[4] * 4, 'box', 5)

    priorbox1 = AnchorBoxes(img_height, img_width, this_scale=scales[0], next_scale=scales[1],
                                             aspect_ratios=aspect_ratios[0],
                                             two_boxes_for_ar1=two_boxes_for_ar1, this_steps=steps[0],
                                             this_offsets=offsets[0], clip_boxes=clip_boxes,
                                             variances=variances, coords=coords, normalize_coords=normalize_coords,
                                             name='ssd_priorbox_1')(link1_box)
    priorbox2 = AnchorBoxes(img_height, img_width, this_scale=scales[1], next_scale=scales[2],
                                    aspect_ratios=aspect_ratios[1],
                                    two_boxes_for_ar1=two_boxes_for_ar1, this_steps=steps[1], this_offsets=offsets[1],
                                    clip_boxes=clip_boxes,
                                    variances=variances, coords=coords, normalize_coords=normalize_coords,
                                    name='ssd_priorbox_2')(link2_box)
    priorbox3 = AnchorBoxes(img_height, img_width, this_scale=scales[2], next_scale=scales[3],
                                        aspect_ratios=aspect_ratios[2],
                                        two_boxes_for_ar1=two_boxes_for_ar1, this_steps=steps[2],
                                        this_offsets=offsets[2], clip_boxes=clip_boxes,
                                        variances=variances, coords=coords, normalize_coords=normalize_coords,
                                        name='ssd_priorbox_3')(link3_box)
    priorbox4 = AnchorBoxes(img_height, img_width, this_scale=scales[3], next_scale=scales[4],
                                        aspect_ratios=aspect_ratios[3],
                                        two_boxes_for_ar1=two_boxes_for_ar1, this_steps=steps[3],
                                        this_offsets=offsets[3], clip_boxes=clip_boxes,
                                        variances=variances, coords=coords, normalize_coords=normalize_coords,
                                        name='ssd_priorbox_4')(link4_box)
    priorbox5 = AnchorBoxes(img_height, img_width, this_scale=scales[4], next_scale=scales[5],
                                        aspect_ratios=aspect_ratios[4],
                                        two_boxes_for_ar1=two_boxes_for_ar1, this_steps=steps[4],
                                        this_offsets=offsets[4], clip_boxes=clip_boxes,
                                        variances=variances, coords=coords, normalize_coords=normalize_coords,
                                        name='ssd_priorbox_5')(link5_box)

    # Reshape
    cls1_reshape = Reshape((-1, n_classes), name='ssd_cls1_reshape')(link1_cls)
    cls2_reshape = Reshape((-1, n_classes), name='ssd_cls2_reshape')(link2_cls)
    cls3_reshape = Reshape((-1, n_classes), name='ssd_cls3_reshape')(link3_cls)
    cls4_reshape = Reshape((-1, n_classes), name='ssd_cls4_reshape')(link4_cls)
    cls5_reshape = Reshape((-1, n_classes), name='ssd_cls5_reshape')(link5_cls)

    box1_reshape = Reshape((-1, 4), name='ssd_box1_reshape')(link1_box)
    box2_reshape = Reshape((-1, 4), name='ssd_box2_reshape')(link2_box)
    box3_reshape = Reshape((-1, 4), name='ssd_box3_reshape')(link3_box)
    box4_reshape = Reshape((-1, 4), name='ssd_box4_reshape')(link4_box)
    box5_reshape = Reshape((-1, 4), name='ssd_box5_reshape')(link5_box)

    priorbox1_reshape = Reshape((-1, 8), name='ssd_priorbox1_reshape')(priorbox1)
    priorbox2_reshape = Reshape((-1, 8), name='ssd_priorbox2_reshape')(priorbox2)
    priorbox3_reshape = Reshape((-1, 8), name='ssd_priorbox3_reshape')(priorbox3)
    priorbox4_reshape = Reshape((-1, 8), name='ssd_priorbox4_reshape')(priorbox4)
    priorbox5_reshape = Reshape((-1, 8), name='ssd_priorbox5_reshape')(priorbox5)

    cls = Concatenate(axis=1, name='ssd_cls')(
        [cls1_reshape, cls2_reshape, cls3_reshape, cls4_reshape, cls5_reshape])

    box = Concatenate(axis=1, name='ssd_box')(
        [box1_reshape, box2_reshape, box3_reshape, box4_reshape, box5_reshape])

    priorbox = Concatenate(axis=1, name='ssd_priorbox')(
        [priorbox1_reshape, priorbox2_reshape, priorbox3_reshape,
         priorbox4_reshape, priorbox5_reshape])

    cls = Activation('softmax', name='ssd_mbox_conf_softmax')(cls)

    predictions = Concatenate(axis=2, name='ssd_predictions')([cls, box, priorbox])

    if mode == 'training':
        model = Model(inputs=x, outputs=predictions)
    elif mode == 'inference':
        decoded_predictions = DecodeDetections(confidence_thresh=confidence_thresh,
                                               iou_threshold=iou_threshold,
                                               top_k=top_k,
                                               nms_max_output_size=nms_max_output_size,
                                               coords=coords,
                                               normalize_coords=normalize_coords,
                                               img_height=img_height,
                                               img_width=img_width,
                                               name='ssd_decoded_predictions')(predictions)
        model = Model(inputs=x, outputs=decoded_predictions)
    elif mode == 'inference_fast':
        decoded_predictions = DecodeDetectionsFast(confidence_thresh=confidence_thresh,
                                                   iou_threshold=iou_threshold,
                                                   top_k=top_k,
                                                   nms_max_output_size=nms_max_output_size,
                                                   coords=coords,
                                                   normalize_coords=normalize_coords,
                                                   img_height=img_height,
                                                   img_width=img_width,
                                                   name='ssd_decoded_predictions')(predictions)
        model = Model(inputs=x, outputs=decoded_predictions)
    else:
        raise ValueError(
            "`mode` must be one of 'training', 'inference' or 'inference_fast', but received '{}'.".format(mode))

    if return_predictor_sizes:
        predictor_sizes = np.array(
            [cls[0]._keras_shape[1:3], cls[1]._keras_shape[1:3], cls[2]._keras_shape[1:3],
             cls[3]._keras_shape[1:3], cls[4]._keras_shape[1:3]])
        return model, predictor_sizes
    else:
        return model


import os
import numpy as np
import argparse
from sklearn.metrics.pairwise import cosine_similarity
from glob import glob

CODE_DIR = "/projects/tir5/users/jwieting/multi-single-crosslingual3"
TATOEBA_DATA_DIR = "/projects/tir5/users/jwieting/LASER/data/tatoeba/v1/*"

def embed_load(fname, dim=1024, verbose=False):
    x = np.fromfile(fname, dtype=np.float32, count=-1)
    x.resize(x.shape[0] // dim, dim)
    if verbose:
        print(' - Embeddings: {:s}, {:d}x{:d}'.format(fname, x.shape[0], dim))
    return x

def compute_accuracy(scores):
    n = scores.shape[0]
    correct = 0
    for i in range(n):
        _scores = scores[i]
        max_idx = np.argmax(_scores)
        if max_idx == i:
            correct += 1
    return correct / n * 100

parser = argparse.ArgumentParser()

parser.add_argument("--data", help="training data")
parser.add_argument("--path", help="training data")
parser.add_argument("--sentencepiece", help="")
parser.add_argument("--lang-pairs", help="")
parser.add_argument("--task", help="")
parser.add_argument("--compressed", type=int, help="")
parser.add_argument("--tokenize", type=int, help="")
parser.add_argument("--add-lang-tokens", type=int, default=0, help="")
parser.add_argument("--add-encoder-tokens", type=int, default=0, help="")
parser.add_argument("--eval-type", choices=["tri", "multi"], help="")

args = parser.parse_args()

args.model = args.path.split("/")[-1]

flis = glob(TATOEBA_DATA_DIR)
flis.sort()
lang_lis = []
total = 0
num = 0

lo = []
hi = []

for n in range(0, len(flis), 2):
    i = flis[n]
    if "README" not in i:
        key = i[-3:]
        eng = None
        foreign = None
        lang = None

        if key == "eng":
            eng = i
            foreign = flis[n - 1]
        else:
            eng = flis[n - 1]
            foreign = i

        lang = foreign[-3:]

        if args.eval_type == "tri":
            if lang == "glg" or lang == "aze" or lang == "ara" or lang == "spa" or lang == "tur":
                pass
            else:
                continue
        elif args.eval_type == "multi":
            if lang == "glg" or lang == "aze" or lang == "ukr" or lang == "ara" or lang == "spa" or \
                    lang == "tur" or lang == "fra" or lang == "deu" or lang == "rus":
                pass
            else:
                continue

        os.chdir(CODE_DIR)

        cmd = "python -u embed_bucc.py {0} --path {1} --sentencepiece {2} --task multi_vae --lang-pairs {3} " \
              "--compressed {4}  --add-lang-tokens {5} --add-encoder-tokens {6} --tokenize {9} --bucc-file {7} " \
              "--output-file analysis/embeddings-english-{8}.np".format(args.data,
                                                               args.path, args.sentencepiece, args.lang_pairs,
                                                               args.compressed, args.add_lang_tokens,
                                                               args.add_encoder_tokens, eng, args.model, args.tokenize)
        os.system(cmd)

        cmd = "python -u embed_bucc.py {0} --path {1} --sentencepiece {2} --task multi_vae --lang-pairs {3} " \
              "--compressed {4}  --add-lang-tokens {5} --add-encoder-tokens {6} --tokenize {9} --bucc-file {7} " \
              "--output-file analysis/embeddings-foreign-{8}.np".format(args.data,
                                                               args.path, args.sentencepiece, args.lang_pairs,
                                                               args.compressed, args.add_lang_tokens,
                                                               args.add_encoder_tokens, foreign, args.model, args.tokenize)
        os.system(cmd)

        eng_emb = embed_load("analysis/embeddings-english-{0}.np".format(args.model))
        foreign_emb = embed_load("analysis/embeddings-foreign-{0}.np".format(args.model))
        scores = cosine_similarity(eng_emb, foreign_emb)

        english_acc = compute_accuracy(scores)
        foreign_acc = compute_accuracy(np.transpose(scores))

        total += english_acc + foreign_acc
        num += 2

        print(lang, english_acc, foreign_acc)

        if lang == "glg" or lang == "aze" or lang == "ukr":
            lo.append(english_acc)
            lo.append(foreign_acc)
        else:
            hi.append(english_acc)
            hi.append(foreign_acc)

print("low resource: ", np.mean(lo))
print("high resource: ", np.mean(hi))

print(total / num)


import numpy as np
import tensorflow as tf
import flask
import io

from keras.models import load_model
from keras.applications import imagenet_utils
from keras.preprocessing.image import img_to_array

from PIL import Image

# initialize our Flask application
app = flask.Flask(__name__)

# CONSTANTS
# Store list of lesion names
lesion_names = ['Melanoma', 'Nevus', 'Seborrheic Keratosis']

# USAGE
# Start the server:
# 	python launch_prediction_server.py
# Submit a request via cURL:
# 	curl -X POST -F image=@image.jpg "http://{IP-address}:5000/predict"

# PREDICTION methods
def load_cnn_model():
    """ Loads the model using the keras hdf5 file saved to disk """
    global model
    model = load_model('saved_models/final_trained_model.h5')

    global graph  # Attempt to make keras run on threads
    graph = tf.get_default_graph()

def prepare_image(image, target):
    """ Converts the image to the required format for classification """

    # If the image mode is not RGB, convert it
    if image.mode != "RGB":
        image = image.convert("RGB")

    # resize the input image and preprocess it
    image = image.resize(target)
    image = img_to_array(image)
    image = np.expand_dims(image, axis=0)
    image = imagenet_utils.preprocess_input(image)

    # return the processed image
    return image

@app.route("/predict", methods=["POST"])
def classify_image():
    """ Classifies an image provided by the user
     NB: this service is exposed as an endpoint """

    # Iinitialize the data dictionary that will be returned from the view
    data = {"success": False}

    # Ensure an image was properly uploaded to our endpoint
    if flask.request.method == "POST":
        if flask.request.files.get("image"):
            # read the image in PIL format
            image = flask.request.files["image"].read()
            image = Image.open(io.BytesIO(image))

            # Preprocess the image and prepare it for classification
            image = prepare_image(image, target=(224, 224))

            # classify the input image and then initialize the list
            # of predictions to return to the client

            # Code for threadsafing (required to make code run on AWS EC2)
            global graph
            with graph.as_default():
                preds = model.predict(image)[0]

            # Return most probable
            data["prediction"] = lesion_names[np.argmax(preds)]

            # Return all predictions
            preds = preds.tolist()
            data["predictions"] = {name: preds[idx] for idx, name in enumerate(lesion_names)}

            # loop over the results and add them to the list of
            # returned predictions
            # for (imagenetID, label, prob) in results[0]:
            #     r = {"label": label, "probability": float(prob)}
            #     data["predictions"].append(r)

            # indicate that the request was a success
            data["success"] = True

    # return the data dictionary as a JSON response
    return flask.jsonify(data)

# if this is the main thread of execution first load the model and
# then start the server
if __name__ == "__main__":
    print(("* Loading Keras model and Flask starting server..."
           "please wait until server has fully started"))
    load_cnn_model()
    app.run(host='0.0.0.0')  # add host arg so it's available on external addresses


import os
from curses import wrapper
import csv
import numpy as np
from settings import settings

from utils import data
from action_utils import parse_action_args
from args import get_args
from comm import CommNetMLP
from evaluator import Evaluator
from nns.models import *
from utils.util_fns import *
from utils.game_tracker import GameTracker

from nns.probe import Probe

torch.utils.backcompat.broadcast_warning.enabled = True
torch.utils.backcompat.keepdim_warning.enabled = True
torch.set_default_tensor_type('torch.DoubleTensor')

def load_parent_child_probes(probe_seed, probe_dropout_rate):
    tracker_path = os.path.join(args.load, args.env_name, args.exp_name, "seed" + str(args.seed), 'tracker.pkl')
    old_tracker = GameTracker.from_file(tracker_path)
    c_dim = old_tracker.data[0][2][0].detach().numpy().shape[1]
    probe_pred_dim = 49
    c_probes = [Probe(c_dim, probe_pred_dim, num_layers=3) for _ in range(2)]
    [c_probe.load_state_dict(torch.load(os.path.join(args.load, args.env_name, args.exp_name, "seed" +
                                                     str(args.seed), 'probes', 'seed' + str(probe_seed),
                                                     'c_probe_' + str(i) + '_dropout_' + str(
                                                         probe_dropout_rate) + '.pth'))) for i, c_probe in
     enumerate(c_probes)]
    [c_probe.eval() for c_probe in c_probes]

    h_probes = [Probe(c_dim, probe_pred_dim, num_layers=3) for _ in range(self.num_agents)]
    [h_probe.load_state_dict(torch.load(
        os.path.join(args.load, args.env_name, args.exp_name, "seed" + str(args.seed), 'probes',
                     'seed' + str(probe_seed), 'h_probe_' +
                     str(i) + '_dropout_' + str(probe_dropout_rate) + '.pth'))) for i, h_probe in
     enumerate(h_probes)]
    [h_probe.eval() for h_probe in h_probes]
    return c_probes, h_probes

def run_eval(_):
    def load(path):
        load_path = os.path.join(args.load, args.env_name, args.exp_name, "seed" + str(args.seed), "models")
        print(f"load directory is {load_path}")
        log_path = os.path.join(args.load, args.env_name, args.exp_name, "seed" + str(args.seed), "logs")
        print(f"log dir directory is {log_path}")

        assert 'model.pt' in os.listdir(load_path), "No model to load!?"
        model_path = os.path.join(load_path, "model.pt")
        d = torch.load(model_path)
        policy_net.load_state_dict(d['policy_net'])

    if args.ic3net:
        args.commnet = 1
        args.hard_attn = 1
        args.mean_ratio = 0

        # For TJ set comm action to 1 as specified in paper to showcase
        # importance of individual rewards even in cooperative games
        if args.env_name == "traffic_junction":
            args.comm_action_one = True
    # Enemy comm
    args.nfriendly = args.nagents
    if hasattr(args, 'enemy_comm') and args.enemy_comm:
        if hasattr(args, 'nenemies'):
            args.nagents += args.nenemies
        else:
            raise RuntimeError("Env. needs to pass argument 'nenemy'.")

    env = data.init(args.env_name, args, False)

    num_inputs = env.observation_dim
    args.num_actions = env.num_actions

    # Multi-action
    if not isinstance(args.num_actions, (list, tuple)): # single action case
        args.num_actions = [args.num_actions]
    args.dim_actions = env.dim_actions
    args.num_inputs = num_inputs

    # Hard attention
    if args.hard_attn and args.commnet:
        # add comm_action as last dim in actions
        args.num_actions = [*args.num_actions, 2]
        args.dim_actions = env.dim_actions + 1

    # Recurrence
    if args.commnet and (args.recurrent or args.rnn_type == 'LSTM'):
        args.recurrent = True
        args.rnn_type = 'LSTM'

    parse_action_args(args)

    if args.seed == -1:
        args.seed = np.random.randint(0, 10000)
    torch.manual_seed(args.seed)

    print(args)
    print(args.seed)

    if args.commnet:
        print("Creating commnet mlp")
        policy_net = CommNetMLP(args, num_inputs, train_mode=False)
    elif args.random:
        policy_net = Random(args, num_inputs)

    # this is what we are working with for IC3 Net predator prey.
    elif args.recurrent:
        print("Creating an RNN!")
        policy_net = RNN(args, num_inputs)
    else:
        policy_net = MLP(args, num_inputs)

    load(args.load)

    if not args.display:
        display_models([policy_net])

    # share parameters among threads, but not gradients
    for p in policy_net.parameters():
        p.data.share_memory_()

    if args.use_tracker:
        evaluator = Evaluator(args, policy_net, data.init(args.env_name, args), 0, 0,
                              -1)
        all_stats = []
        for i in range(1000):
            ep, stat, all_comms = evaluator.run_episode()
            all_stats.append(stat)
            if i % 50 == 0:
                print("Episode number", i)
        tracker_path = os.path.join(args.load, args.env_name, args.exp_name, "seed" + str(args.seed), "tracker.pkl")
        evaluator.tracker.to_file(tracker_path)
        return

    # num_steps = [i for i in range(1, 2)]
    # For traffic, always intervene
    num_steps = [10]
    dropout_rates = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
    # dropout_rates = [0.1]
    # dropout_rates = [0.0, 0.1, 0.2, 0.3, 0.4]
    # dropout_rates = [0.0]
    probe_seeds = [i for i in range(0, 5)]
    # probe_seeds = [0]
    # xfact_steps = [50, 250, 500, 750, 1000, 2000, 3000, 4000, 5000]
    xfact_steps = [200]
    success_table = np.zeros((len(dropout_rates), len(num_steps), 2))
    collision_table = np.zeros((len(dropout_rates), len(num_steps), 2))
    time_table = np.zeros((len(dropout_rates), len(num_steps), 2))
    for xfact_step in xfact_steps:
        settings.NUM_XFACT_STEPS = xfact_step
        for inter_idx, num_intervention_steps in enumerate(num_steps):
            for dropout_idx, probe_dropout_rate in enumerate(dropout_rates):
                succ_for_seed = []
                collision_for_seed = []
                time_for_seed = []
                for probe_seed in probe_seeds:
                    print("Eval for dropout", probe_dropout_rate, "for", num_intervention_steps, "steps for probe seed", probe_seed)
                    evaluator = Evaluator(args, policy_net, data.init(args.env_name, args), probe_dropout_rate, probe_seed, num_intervention_steps)
                    st_time = time.time()
                    all_stats = []
                    for i in range(100):
                        ep, stat, all_comms = evaluator.run_episode()
                        all_stats.append(stat)
                        if i % 20 == 0:
                            print("Episode", i)
                    if args.use_tracker:
                        tracker_path = os.path.join(args.load, args.env_name, args.exp_name, "seed" + str(args.seed), "tracker.pkl")
                        evaluator.tracker.to_file(tracker_path)

                    total_episode_time = time.time() - st_time
                    average_stat = {}
                    for key in all_stats[0].keys():
                        average_stat[key] = np.mean([stat.get(key) for stat in all_stats])
                    print("average stats is: ", average_stat)
                    succ_for_seed.append(average_stat.get('success'))
                    if average_stat.get('collisions') is not None:
                        collision_for_seed.append(average_stat.get('collisions'))
                    time_for_seed.append(total_episode_time/average_stat['num_steps'])
                success_table[dropout_idx, inter_idx, 0] = np.mean(succ_for_seed)
                success_table[dropout_idx, inter_idx, 1] = np.std(succ_for_seed)
                collision_table[dropout_idx, inter_idx, 0] = np.mean(collision_for_seed)
                collision_table[dropout_idx, inter_idx, 1] = np.std(collision_for_seed)
                time_table[dropout_idx, inter_idx, 0] = np.mean(time_for_seed)
                time_table[dropout_idx, inter_idx, 1] = np.std(time_for_seed)
                # print("Success table\n", success_table[:, :, 0])
                # print("Collision table\n", collision_table[:, :, 0])
                # print("Success std table\n", success_table[:, :, 1])
                # print("Collision std table\n", collision_table[:, :, 1])
                with open(os.path.join(args.load, args.env_name, args.exp_name, "seed" + str(args.seed),
                                       "success_table_" + str(xfact_step) + ".csv"), 'w') as f:
                    f.write("Success mean\n")
                    writer = csv.writer(f)
                    writer.writerows(success_table[:, :, 0])
                    f.write("Success std\n")
                    writer.writerows(success_table[:, :, 1])
                    f.write("Collision mean\n")
                    writer.writerows(collision_table[:, :, 0])
                    f.write("Collision std\n")
                    writer.writerows(collision_table[:, :, 1])
                    f.write("Time mean\n")
                    writer.writerows(time_table[:, :, 0])
                    f.write("Time std\n")
                    writer.writerows(time_table[:, :, 1])

if __name__ == '__main__':
    # Wrap entire execution in curses wrapper to protect terminal against ugly end states.
    parser = get_args()
    init_args_for_env(parser)
    args = parser.parse_args()
    if args.display:
        wrapper(run_eval)
    else:
        run_eval(None)


# Code to check if left or right mouse buttons were pressed
import win32api
import time
import pyautogui
import numpy
from win32gui import WindowFromPoint, GetWindowRect

pyautogui.PAUSE = 0
pyautogui.FAILSAVE = False

state_left = win32api.GetKeyState(0x01)  # Left button down = 0 or 1. Button up = -127 or -128
state_right = win32api.GetKeyState(0x02)  # Right button down = 0 or 1. Button up = -127 or -128

large_movement_picture = None
movement_start_time = None
previous_position = pyautogui.position()
window_bounds = GetWindowRect( WindowFromPoint( (0, 0) ) )
window_bounds_text = '[' + ','.join(str(x) for x in window_bounds) + ']'

while True:
    a = win32api.GetKeyState(0x01)
    b = win32api.GetKeyState(0x02)
    position = pyautogui.position()
    if a != state_left:  # Button state changed
        state_left = a
        pic = pyautogui.screenshot()

        if a < 0:
		    # Keep the window bounds only when holding down the mouse button, because the windows size can change based on releasing the mouse button
            window_bounds = GetWindowRect( WindowFromPoint( position ) )
            window_bounds_text = '[' + ','.join(str(x) for x in window_bounds) + ']'
            pic.save('data/raw/' + str( int( time.time() * 100 ) ) + '--(' + str(position[0]) + '-' + str(position[1]) + ')--' + window_bounds_text + '--mousedown.png')
            print('Saving mousedown screenshot')
        else:
            pic.save('data/raw/' + str( int( time.time() * 100 ) ) + '--(' + str(position[0]) + '-' + str(position[1]) + ')--' + window_bounds_text + '--mouseup.png')
            print( "Saving mouseup screenshot" )

        if large_movement_picture is not None:
            large_movement_picture.save('data/raw/' + str( int( movement_start_time * 100 ) ) + '--(' + str(position[0]) + '-' + str(position[1]) + ')--' + window_bounds_text + '--mousemove.png')
            print( "Saving mousemovement screenshot" )
            large_movement_picture = None
            movement_start_time = None

    if b != state_right:  # Button state changed
        state_right = b
        #print(b)
        #if b < 0:
            #print('Right Button Pressed')
        #else:
            #print('Right Button Released')

    # Large movement detection
    xDistance = numpy.linalg.norm(previous_position[0]-position[0])
    yDistance = numpy.linalg.norm(previous_position[1]-position[1])
    if( xDistance + yDistance > 10 ):
        large_movement_picture = pyautogui.screenshot()
        movement_start_time = time.time()
        print( "Detecting large movement - " + str( xDistance + yDistance ) )

    previous_position = position