# construct tube for box_level tracking

# -*- coding: UTF-8 -*-

import json
import copy
import numpy as np
import cv2
from multiprocessing import Pool, current_process
import pickle
import os

score_threshold = 0.5 # 对于检测分数选择阈值
mask_threshold = 0.5 # 对于道路主体mask的阈值
link_iou_threshold = 0.4 # 链接起来形成tube的阈值
iou_threshold = 0.8 #停止的车辆至少有多少的iou交集
iou_threshold_merge = 0.1 #停止的车辆至少有多少的iou交集
time_threshold = 100 #跨过多少帧就进入候选考虑范围
skip_threshold = 10 #允许跳过几个帧，弥补检测性能的缺失
span_threshold = 1000 #扩散的阈值，50s
merge_threshold = 3500 #前一个结束与后一个开始的允许阈值
stop_freqency_threshold = 0.5
# for idx in range(100):
KEYS = ['x1', 'y1', 'x_scale', 'y_scale']
im_w = float(800)
im_h = float(410)

worker_cnt = 10
full_rate = 30
rate = 30
Sam_interval = int(full_rate/rate)
print("Sam_interval", Sam_interval)

def compute_IoU(box1, box2):
    if isinstance(box1, list):
        box1 = {key: val for key, val in zip(KEYS, box1)}
    if isinstance(box2, list):
        box2 = {key: val for key, val in zip(KEYS, box2)}

    box1['x2'] = box1['x1'] + box1['x_scale']
    box1['y2'] = box1['y1'] + box1['y_scale']

    box2['x2'] = box2['x1'] + box2['x_scale']
    box2['y2'] = box2['y1'] + box2['y_scale']
    width = max(min(box1['x2'] / im_w, box2['x2'] / im_w) - max(box1['x1'] / im_w, box2['x1'] / im_w), 0)
    height = max(min(box1['y2'] / im_h, box2['y2'] / im_h) - max(box1['y1'] / im_h, box2['y1'] / im_h), 0)
    intersection = width * height
    box1_area = (box1['x2'] / im_w - box1['x1'] / im_w) * (box1['y2'] / im_h - box1['y1'] / im_h)
    box2_area = (box2['x2'] / im_w - box2['x1'] / im_w) * (box2['y2'] / im_h - box2['y1'] / im_h)

    union = box1_area + box2_area - intersection

    return float(intersection) / float(union)

def get_top_tubes(det_list_org, link_iou_threshold, skip_threshold, time_threshold):
    # tube construction process
    det_list = copy.deepcopy(det_list_org)
    tubes = []
    tube_scores = []
    continue_flg = True

    while continue_flg:
        # find the top score in the current det_list
        max_score = 0
        max_timestep = 0
        max_boxid = 0
        for i in range(len(det_list)):
            for j in range(len(det_list[i]["detection_result"])):
                if det_list[i]["detection_result"][j][-1] > max_score:
                    max_score = det_list[i]["detection_result"][j][-1]
                    max_timestep = i
                    max_boxid = j

        tube_id = []
        tube_id.append((max_timestep, max_boxid, 1))

        # forward
        timestep = max_timestep
        continue_flag = True
        while timestep + 1 < len(det_list) and continue_flag == True:
            curbox_coods = det_list[tube_id[-1][0]]["detection_result"][tube_id[-1][1]][0]
            n_curbox = len(det_list[timestep + 1]["detection_result"])
            add_flag = False
            if n_curbox > 0:
                max_link_score = 0
                for i_nextbox in range(n_curbox):
                    nextbox_coods = det_list[timestep + 1]["detection_result"][i_nextbox][0]
                    link_score = compute_IoU(nextbox_coods, curbox_coods)

                    if link_score > link_iou_threshold and link_score > max_link_score:
                        max_boxid = i_nextbox
                        max_link_score = link_score
                        add_flag = True

            if timestep + 1 - tube_id[-1][0] > skip_threshold:
                continue_flag = False
            elif add_flag == True:
                tube_id.append((timestep + 1, max_boxid, max_link_score))

            timestep += 1

            if timestep + 1 == len(det_list):
                continue_flag = False

        # backward
        timestep = max_timestep
        continue_flag = True
        while timestep - 1 >= 0 and continue_flag == True:
            curbox_coods = det_list[tube_id[0][0]]["detection_result"][tube_id[0][1]][0]
            n_curbox = len(det_list[timestep - 1]["detection_result"])
            add_flag = False
            if n_curbox > 0:
                max_link_score = 0
                for i_prevbox in range(n_curbox):
                    prevbox_coods = det_list[timestep - 1]["detection_result"][i_prevbox][0]
                    link_score = compute_IoU(prevbox_coods, curbox_coods)

                    if link_score > iou_threshold and link_score > max_link_score:
                        max_boxid = i_prevbox
                        max_link_score = link_score
                        add_flag = True

            if tube_id[0][0] - timestep - 1 > skip_threshold:
                continue_flag = False
            elif add_flag == True:
                tube_id.insert(0, (timestep - 1, max_boxid, max_link_score))

            timestep -= 1

            if timestep <= 0:
                continue_flag = False

        # get path and remove used boxes
        cur_tube = []
        tube_score = 0
        for i, tube_path in enumerate(tube_id):

            inform = {}
            inform["frame_id"] = tube_path[0]
            inform["detection_result"] = det_list[tube_path[0]]["detection_result"][tube_path[1]]
            inform["link_score"] = tube_path[2]
            tube_score += tube_path[2]

            cur_tube.append(inform)
            det_list[tube_path[0]]["detection_result"] = np.delete(det_list[tube_path[0]]["detection_result"],
                                                                   tube_path[1], 0)

            sum_detection = 0
            for kk in range(len(det_list)):
                sum_detection += len(det_list[kk]["detection_result"])
            if sum_detection == 0:
                continue_flg = False

        if len(tube_id) > time_threshold:
            tube_score = tube_score / len(tube_id)
            tubes.append(cur_tube)
            tube_scores.append(tube_score)

    return tubes, tube_scores

def process(idx):
    print('searching video number %s' % (idx))

    idx = int(idx)
    video_data = []
    mask_path = "./intermediate_result/mask/mask-fuse/mask_%03d.jpg" %idx
    mask_img = cv2.imread(mask_path)
    imgray = cv2.cvtColor(mask_img, cv2.COLOR_BGR2GRAY)

    thresh = 128
    ret, binary_mask = cv2.threshold(imgray, thresh, 1, cv2.THRESH_BINARY)

    for num in range(1,30000):
        try:
            current_id = num * Sam_interval
            Det_path = "./intermediate_result/det_result/result_30frame_detmodel2_5w/%d/test_%d_%05d.jpg.npy" % (idx, idx, current_id)
            det_res = np.load(Det_path, allow_pickle=True)
            inform = {}
            inform["image_id"] = num
            len_box = len(det_res)
            det_filter = []
            for i in range(len_box):
                if det_res[i][2] > score_threshold:
                    x1 = int(det_res[i][0][0])
                    y1 = int(det_res[i][0][1])
                    x2 = int(det_res[i][0][0] + det_res[i][0][2])
                    y2 = int(det_res[i][0][1] + det_res[i][0][3])
                    mask_area = (binary_mask[y1:y2, x1:x2]).sum()
                    area = (y2 - y1) * (x2 - x1)
                    mask_iou = mask_area / float(area)
                    if mask_iou >= mask_threshold:
                        det_filter.append((det_res[i][0], det_res[i][1], det_res[i][2]))
            inform["detection_result"] = det_filter
            video_data.append(inform)
        except:
            pass

    tube_all = []
    tubes, tube_scores = get_top_tubes(video_data,  link_iou_threshold, skip_threshold, time_threshold)
    tube_all.append((tubes, tube_scores))

    save_dir = "./intermediate_result/box_track/video_data5_30_mask_fuse_5"
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    json_dir = "./intermediate_result/box_track/video_data5_30_mask_fuse_5/" + str(idx) + ".pkl"

    with open(json_dir, "wb") as file:
        pickle.dump(tube_all, file)

threads = []

idx_list = []
happen_list = []

for idx in range(1,101):
    idx_list.append(str(idx))

pool = Pool(processes=worker_cnt)
pool.map(process, idx_list)


import numpy
import os
import sys

from nematus.nmt import train

#VOCAB_SIZE = 90000
# are src and trg vocab sizes required?

# For AmuNMT WMT 2016 pretrained models
SRC_VOCAB_SIZE=40587
MT_VOCAB_SIZE=40822
PE_VOCAB_SIZE=41101

PATIENCE = 99999999999 #never early stop
SRC = "en"
TGT = "de"

#VOCAB_DIR = "/media/1tb_drive/nematus_ape_experiments/pretrained_wmt16_models/en-de"
VOCAB_DIR = "/media/1tb_drive/nematus_ape_experiments/amunmt_ape_pretrained/system/models"
SRC_VOCAB = os.path.join(VOCAB_DIR, 'src-pe/vocab.src.json')
MT_VOCAB = os.path.join(VOCAB_DIR, 'mt-pe/vocab.mt.json')
PE_VOCAB = os.path.join(VOCAB_DIR, 'mt-pe/vocab.pe.json')

#TRAIN_DATA_DIR = "/media/1tb_drive/Dropbox/data/qe/amunmt_artificial_ape_2016/data/concat_500k_with_wmt16/factored_ape_corpus"
TRAIN_DATA_DIR = "/media/1tb_drive/Dropbox/data/qe/amunmt_artificial_ape_2016/data/4M"
SRC_TRAIN = os.path.join(TRAIN_DATA_DIR, 'train.mt.factor_corpus')
TRG_TRAIN = os.path.join(TRAIN_DATA_DIR, 'train.pe.prepped')

# WMT 16 EN-DE QE/APE DEV Data
QE_DATA_DIR = "/media/1tb_drive/Dropbox/data/qe/ape/concat_wmt_2016_2017"
SRC_DEV = os.path.join(QE_DATA_DIR, 'dev.mt.factor_corpus')
TRG_DEV = os.path.join(QE_DATA_DIR, 'dev.pe.prepped')

if __name__ == '__main__':
    validerr = train(saveto='model/model.npz',
                    reload_=True,
                    dim_word=256,
                    dim=512,
                    n_words=PE_VOCAB_SIZE,
                    n_words_src=SRC_VOCAB_SIZE,
                    decay_c=0.,
                    clip_c=1.,
                    lrate=0.0001,
                    optimizer='adadelta',
                    maxlen=50,
                    batch_size=32,
                    valid_batch_size=32,
                    datasets=[SRC_TRAIN, TRG_TRAIN],
                    valid_datasets=[SRC_DEV, TRG_DEV],
                    dictionaries=[MT_VOCAB, SRC_VOCAB, PE_VOCAB],
                    tie_encoder_decoder_embeddings=False,
                    factors=2,
                    dim_per_factor=[128, 128],
                    validFreq=5000,
                    dispFreq=500,
                    saveFreq=5000,
                    sampleFreq=5000,
                    use_dropout=False,
                    dropout_embedding=0.2, # dropout for input embeddings (0: no dropout)
                    dropout_hidden=0.2, # dropout for hidden layers (0: no dropout)
                    dropout_source=0.1, # dropout source words (0: no dropout)
                    dropout_target=0.1, # dropout target words (0: no dropout)
                    overwrite=False,
                    external_validation_script='./validate.sh')
    print validerr


"""
Tabulate the computing resource usage for each DRP pipetask.
"""
import os
import time
from collections import defaultdict
import json
import numpy as np
import pandas as pd
from lsst.daf.butler import Butler, DimensionUniverse
from lsst.pipe.base.graph import QuantumGraph

__all__ = ['get_pipetask_resource_funcs', 'tabulate_pipetask_resources',
           'tabulate_data_product_sizes', 'total_node_hours']

class PipetaskFunc:
    """
    Functor class to package outputs of cpu_time and maxRSS functions.
    """
    def __init__(self, cpu_time, maxRSS):
        """
        Parameters
        ----------
        cpu_time : function
            Function to return cpu time in units of hours as a function
            of number of visits.
        maxRSS : function
            Function to return maximum RSS in units of GB as a function
            of number of visits.
        """
        self.cpu_time = cpu_time
        self.maxRSS = maxRSS

    def __call__(self, num_visits=1):
        """
        Parameters
        ----------
        num_visits : int [1]
            Number of visits to consider.  Not all functions depend
            meaningfully on the number of visits, e.g., the single
            frame processing, so a default of 1 is used.

        Returns
        -------
        tuple of (cpu_time (h), memory (GB))
        """
        return self.cpu_time(num_visits), self.maxRSS(num_visits)

def get_pipetask_resource_funcs(json_file, cpu_time_label='cpu_time (m)',
                                maxRSS_label='maxRSS (GB)',
                                cpu_time_factor=1./60., maxRSS_factor=1):
    """
    Read in the json file containing the functions describing the
    computing resource usage anticipated for each task as a function
    of number of visits.  These functions should return a tuple of cpu
    time (hours) and max RSS usage (GB).

    Parameters
    ----------
    json_file : str
        json file containing parameters describing the anticipated
        cpu time and memory used by the DRP pipetasks.
    cpu_time_label : str ['cpu_time (m)']
        Key name pointing to cpu time parameters.
    maxRSS_label : str ['maxRSS (GB)']
        Key name pointing to maximum RSS parameters.
    cpu_time_factor : float [1/60]
        Conversion factor to multiply the native units to convert to hours.
    maxRSS_factor : float [1]
        Conversion factor to multiply the native units to convert to GB.

    Returns
    -------
    dict of PipetaskFunc objects, keyed by task type.
    """
    with open(json_file) as fd:
        model_params = json.load(fd)
    pipetask_funcs = dict()
    for task, pars in model_params.items():
        cpu_time = np.poly1d(cpu_time_factor*np.array(pars[cpu_time_label]))
        maxRSS = np.poly1d(maxRSS_factor*np.array(pars[maxRSS_label]))
        pipetask_funcs[task] = PipetaskFunc(cpu_time, maxRSS)
    return pipetask_funcs

def tabulate_pipetask_resources(coadd_df, task_counts, pipetask_funcs,
                                num_visit_col='num_visits', verbose=False):
    """
    Tabulate the computing resources (cpu time, memory) for each
    of the pipetasks given a dataframe with the overlaps information.

    Parameters
    ----------
    coadd_df : pandas.DataFrame
        Dataframe with the number of visits for each band-tract-patch
        combination.
    task_counts : dict
        Dictionary of number of instances per task. Only counts for 'isr',
        'makeWarp', and 'assembleCoadd' are used.
    pipetask_funcs : dict
        Dictionary of functions, keyed by task type.  Each function should
        return a tuple (cpu_time in hours, memory usage in GB) taking
        the number of visits as the argument.
    num_visit_col : str ['num_visits']
        Column name in coadd_df that contains the number of visits per
        band-tract-patch coadd.
    verbose : bool [False]
        Verbosity flag.

    Returns
    -------
    pandas.DataFrame with the number of instances, total cpu time,
    and maximum and average memory used per pipetask.
    """
    pt_data = defaultdict(list)

    # sensor-visits:
    num_ccd_visits = task_counts['isr']
    for task_name in 'isr characterizeImage calibrate'.split():
        if verbose:
            print("processing", task_name)
        pt_data['pipetask'].append(task_name)
        pt_data['num_instances'].append(num_ccd_visits)
        cpu_hours, mem_GB = pipetask_funcs[task_name]()
        pt_data['cpu_hours'].append(cpu_hours*num_ccd_visits)
        pt_data['max_GB'].append(mem_GB)
        pt_data['avg_GB'].append(mem_GB)

    # warps:
    task_name = 'makeWarp'
    num_warps = task_counts[task_name]
    if verbose:
        print("processing", task_name)
    pt_data['pipetask'].append(task_name)
    pt_data['num_instances'].append(num_warps)
    cpu_hours, mem_GB = pipetask_funcs[task_name]()
    pt_data['cpu_hours'].append(cpu_hours*num_warps)
    pt_data['max_GB'].append(mem_GB)
    pt_data['avg_GB'].append(mem_GB)

    for task_name in ('assembleCoadd', 'detection', 'measure',
                      'forcedPhotCoadd'):
        if verbose:
            print("processing", task_name, end=' ')
        pt_data['pipetask'].append(task_name)
        pt_data['num_instances'].append(len(coadd_df))
        cpu_hours_total = 0
        memory = []
        t0 = time.time()
        for _, row in coadd_df.iterrows():
            num_visits = row[num_visit_col]
            cpu_hours, mem_GB = pipetask_funcs[task_name](num_visits)
            cpu_hours_total += cpu_hours
            memory.append(mem_GB)
        pt_data['cpu_hours'].append(cpu_hours_total)
        pt_data['max_GB'].append(np.max(memory))
        pt_data['avg_GB'].append(np.mean(memory))
        print(time.time() - t0)

    for task_name in ('mergeCoaddDetections', 'deblend',
                      'mergeCoaddMeasurements'):
        try:
            cpu_hours, mem_GB = pipetask_funcs[task_name]()
        except KeyError:
            continue
        if verbose:
            print("processing", task_name)
        pt_data['pipetask'].append(task_name)
        num_instances = int(len(coadd_df)/6)
        pt_data['num_instances'].append(num_instances)
        pt_data['cpu_hours'].append(num_instances*cpu_hours)
        pt_data['max_GB'].append(mem_GB)
        pt_data['avg_GB'].append(mem_GB)

    return pd.DataFrame(data=pt_data)

def tabulate_data_product_sizes(qgraph_file, repo, collection):
    """
    Tabulate the mean sizes of data products listed in a QuantumGraph
    using files in a given repo and collection.

    Parameters
    ----------
    qgraph_file : str
        QuantumGraph file produced by `pipetask qgraph`.
    repo : str
        Path to data repository.
    collection : str
        Collection in repo to use for finding example data products.

    Returns
    -------
    dict(dict(tuple)) Outer dict keyed by task label, inner dicts keyed
    by dataset type with tuple of (mean file size (GB), std file sizes (GB),
    number of files in examples).
    """
    qgraph = QuantumGraph.loadUri(qgraph_file, DimensionUniverse())

    butler = Butler(repo, collections=[collection])
    registry = butler.registry

    # Traverse the QuantumGraph finding the dataset types associated
    # with each task type.
    dstypes = defaultdict(set)
    for node in qgraph:
        task = node.taskDef.label
        for dstype in node.quantum.outputs:
            dstypes[task].add(dstype.name)

    # Loop over task types and query for each dataset types and
    # compute mean and stdev file sizes for each dataset type.
    data = defaultdict(dict)
    for task, dstypes in dstypes.items():
        for dstype in dstypes:
            file_sizes = [os.stat(butler.getURI(_).path).st_size/1024**3
                          for _ in registry.queryDatasets(dstype)]
            data[task][dstype] = (np.nanmean(file_sizes), np.nanstd(file_sizes),
                                  len(file_sizes))
    return data

def total_node_hours(pt_df, cpu_factor=8, cores_per_node=68,
                     memory_per_node=96, memory_min=10):
    """
    Estimate the total number of node hours to do an image processing
    run.

    Parameters
    ----------
    pt_df : pandas.DataFrame
        DataFrame containing the number of instances, total cpu time,
        and maximum and average memory used per pipetask.  This is
        the output of `tabulate_pipetask_resources`.
    cpu_factor : float [8]
        Slow down factor to apply to the pipetask cpu times.  The
        cpu times were derived from Cori-Haswell runs and the default
        value of 8 is the empirically observed slow-down for running
        the same cpu-bound job on a Cori-KNL node.
    cores_per_node : int [68]
        Number of cores per node.  68 is the Cori-KNL value.
    memory_per_node : int [96]
        Memory per node in GB.  96 is the Cori-KNL value.
    memory_min : int [10]
        Memory in GB to reserve per node as a safety factor.  10GB is
        a conservative number for these jobs.

    Returns
    -------
    tuple(float, float): The first entry, `node_hours`, is computed
    using the maximum memory estimate per process to determine the
    number of cores per node for a given pipe task ; the second entry,
    `node_hours_opt`, is an optimistic estimate using the average
    memory per process.
    """
    available_memory = memory_per_node - memory_min
    node_hours = 0
    node_hours_opt = 0
    for _, row in pt_df.iterrows():
        ncores = min(cores_per_node, int(available_memory/row['max_GB']))
        ncores_avg = min(cores_per_node, int(available_memory/row['avg_GB']))
        node_hours += row['cpu_hours']*cpu_factor/ncores
        node_hours_opt += row['cpu_hours']*cpu_factor/ncores_avg
    return node_hours, node_hours_opt


# last edited: 29/07/2016, 15:00

import os,sys

sys.path.append(os.path.join(os.getenv('XChemExplorer_DIR'),'lib'))

import XChemMain

import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
import numpy as np

class summary_plot(object):
    def __init__(self,datasource,overview_axes):
        self.datasource=datasource
        self.overview_axes=overview_axes

    def update_overview(self):

        db_dict=XChemMain.get_datasource_summary(self.datasource)

        category = (    'Crystals mounted',
                        'Data Collection',
                        'Maps',
                        'PANDDA',
                        'Refinement',
                        'Comp Chem'             )

        N = len(category)
        Success = np.array([    db_dict['nr_samples'],
                                db_dict['nr_data_collection_success'],
                                db_dict['nr_initial_maps_available'],
                                db_dict['nr_pandda_hits'],
                                5,
                                6   ])

        Processed = np.array([  0,
                                0,
                                0,
                                db_dict['nr_pandda_processed'],
                                0,
                                0   ])

        Pending = np.array([    0,
                                db_dict['nr_data_collection_pending'],
                                db_dict['nr_initial_maps_pending'],
                                db_dict['nr_pandda_pending'],
                                2,
                                1   ])

        Failure = np.array([    db_dict['nr_samples_failed_to_mount'],
                                db_dict['nr_data_collection_failed'],
                                db_dict['nr_initial_maps_fail'],
                                db_dict['nr_pandda_reject'],
                                2,
                                3   ])

        ind = np.arange(N)  # the x locations for the groups
        width = 0.35       # the width of the bars

        p0 = self.overview_axes.bar(ind, Success,   width, color='g')
        p1 = self.overview_axes.bar(ind, Processed, width, color='b', bottom=Success)
        p2 = self.overview_axes.bar(ind, Pending,   width, color='y', bottom=Success+Processed)
        p3 = self.overview_axes.bar(ind, Failure,   width, color='r', bottom=Pending+Success+Processed)

        # add some text for labels, title and axes ticks
        self.overview_axes.set_ylabel('N')
        self.overview_axes.set_title('Overview')
        self.overview_axes.set_xticks(ind + width)
        self.overview_axes.set_xticklabels(category)

        self.overview_axes.legend((p0[0], p1[0], p2[0], p3[0]), ('Success', 'Processed', 'Pending', 'Failure'))


from envs.yumi.yumi_door_env import YumiDoorEnv
from gym import utils, spaces
import numpy as np
import os, copy

class YumiDoorButtonEnv(YumiDoorEnv, utils.EzPickle):
    """
    Mujoco yumi environment for button pressing task.
    The robot needs to open the door first then press the button and close the door.
    The extrinsic reward will be given when the button is pressed and the door is closed.

    :param goal_dim: the dimension of the goal state
    :param reward_type: the type for reward computing. The reward type can be "dense", "sparse", "density" and "rnd"
    :param mocap_high: the up-limit of the mocap position
    :param mocap_low: the low-limit of the mocap position
    :param mocap_init_pos: the initial position of mocap
    :param mocap_init_quat: the initial orientation of mocap
    """
    def __init__(self, reward_type='sparse', goal_dim=6,
                    mocap_high=np.hstack((0.0, 0.13, 0.35)), mocap_low=np.hstack((-0.02, -0.155, 0.19)),
                    mocap_init_pos=np.array([-0.04670722, -0.150812, 0.265042291]),
                    mocap_init_quat=np.array([-0.70699531, 0.09357381, -0.4721928, 0.51810765])):

        YumiDoorEnv.__init__(self, xml_name='yumi_door_button.xml', reward_type=reward_type, goal_dim=goal_dim,
                            mocap_high=mocap_high, mocap_low=mocap_low,
                            mocap_init_pos=mocap_init_pos, mocap_init_quat=mocap_init_quat)
        utils.EzPickle.__init__(self)
        self.name = 'yumi_box_pick'
        x_scale, y_scale, z_scale, gripper_scale = 0.015, 0.015, 0.01, 1000
        self.action_scale = np.array([x_scale, y_scale, z_scale, gripper_scale])

        self.initial_qpos = np.array([ 2.08826485e+00,  1.17524374e-01, -2.81825192e+00,  7.20298118e-01,
                                       4.08614131e+00,  1.79107153e-01, -2.93887965e+00,  2.03631476e-06,
                                      -1.09979637e-02, -4.01938280e-06,  0.0])

        self.initial_qvel = np.zeros(11)
        self.reset()
        self.always_render = False

    def compute_extrinsic_reward(self, achieved_goal):
        task_angle1 = 0
        task_angle2 = 0.02
        task_weight = 10.0
        x1 = achieved_goal[-2] - task_angle1
        x2 = achieved_goal[-1] - task_angle2

        delta_alpha = 0.5
        angle_error1 = 0.1
        angle_error2 = 0.012
        if np.abs(x1) < angle_error1 and np.abs(x2) < angle_error2:
            task_reward1 = np.exp(-(x1/delta_alpha)**2)/(np.abs(delta_alpha)*np.sqrt(np.pi))
            task_reward2 = np.exp(-(x2/delta_alpha)**2)/(np.abs(delta_alpha)*np.sqrt(np.pi))
            task_reward = (task_reward1 + task_reward2)/2
        else:
            task_reward = 0

        return task_reward * task_weight

    def set_extra_task_params(self):
        """
        Thsi function is used for reading any task-related information from xml fills.(i.e. requires element id in xml file.)
        """
        self.initial_state = copy.deepcopy(self.sim.get_state())
        self.end_effector = 'gripper_r_base'
        # self.gripper_joints = [self.sim.model.get_joint_qpos_addr('gripper_r_joint'), self.sim.model.get_joint_qpos_addr('gripper_r_joint_m')]
        self.gripper_joints = [self.sim.model.get_joint_qpos_addr('gripper_r_joint')]
        self.door_joint = [self.sim.model.get_joint_qpos_addr('door_l_joint'),
                            self.sim.model.get_joint_qpos_addr('door_r_joint')]

        xyz_range = self.mocap_high - self.mocap_low

        #finger_range = self.sim.model.jnt_qposadr('gripper_r_joint')
        gripper_r_joint_id =  self.sim.model.joint_name2id('gripper_r_joint')
        left_limit = self.sim.model.jnt_range[gripper_r_joint_id]

        gripper_r_joint_id_m =  self.sim.model.joint_name2id('gripper_r_joint_m')
        right_limit = self.sim.model.jnt_range[gripper_r_joint_id_m]

        # gripper_range = np.array([left_limit[1] - left_limit[0],
        #                 right_limit[1] - right_limit[0]])
        gripper_range = [left_limit[1] - left_limit[0]]

        door_l_joint_id =  self.sim.model.joint_name2id('door_l_joint')
        door_l_limit = self.sim.model.jnt_range[door_l_joint_id]
        door_range = [door_l_limit[1] - door_l_limit[0]]

        door_r_joint_id =  self.sim.model.joint_name2id('door_r_joint')
        door_r_limit = self.sim.model.jnt_range[door_r_joint_id]

        # if self.args.use_auto_scale:
        self.obs_scale = np.ones(self.goal_dim)
        # else:
        # self.obs_scale = np.concatenate(([1,1,1,], np.max(xyz_range)/gripper_range,
                                        # np.max(xyz_range)/door_range, np.max(xyz_range)/door_range))
        # self.obs_scale = np.array([1,1,1,1,1])

        self.xyz_start = self.mocap_low * self.obs_scale[0:3]
        self.xyz_end = self.mocap_high * self.obs_scale[0:3]
        self.gripper_start = left_limit[0] * self.obs_scale[3]
        self.gripper_end = left_limit[1] * self.obs_scale[3]
        self.door_l_start = door_l_limit[0] * self.obs_scale[4]
        self.door_l_end = door_l_limit[1] * self.obs_scale[4]
        self.door_r_start = door_r_limit[0] * self.obs_scale[5]
        self.door_r_end = door_r_limit[1] * self.obs_scale[5]


import math
import imageio
import skvideo.io
import numpy as np
import cv2

from tqdm import tqdm
from PIL import Image
from timeit import default_timer as timer
from util import DataLoader

def selectBetterBitrate(height, fps):
    #print(height,fps)
    if (140 < height) and (height < 200):
        bitrate = "280k"
    elif (200 < height) and (height < 250):
        bitrate = "400k"
    elif((300 < height) and (height < 400)) and ((20 < fps) and (fps < 40)):
        bitrate = "1M"
    elif((300 < height) and (height < 400)) and ((40 < fps) and (fps < 70)):
        bitrate = "1.5M"
    elif((400 < height) and (height < 500)) and ((20 < fps) and (fps < 40)):
        bitrate = "2.5M"
    elif((400 < height) and (height < 500)) and ((40 < fps) and (fps < 70)):
        bitrate = "4M"
    elif((700 < height) and (height < 800)) and ((20 < fps) and (fps < 40)):
        bitrate = "5M"
    elif((700 < height) and (height < 800)) and ((40 < fps) and (fps < 70)):
        bitrate = "7.5M"
    elif((1000 < height) and (height < 1100)) and ((20 < fps) and (fps < 40)):
        bitrate = "8M"
    elif((1000 < height) and (height < 1100)) and ((40 < fps) and (fps < 70)):
        bitrate = "12M"
    elif((1300 < height) and (height < 1600)) and ((20 < fps) and (fps < 40)):
        bitrate = "16M"
    elif((1300 < height) and (height < 1600)) and ((40 < fps) and (fps < 70)):
        bitrate = "24M"
    elif((1800 < height) and (height < 2400)) and ((20 < fps) and (fps < 40)):
        bitrate = "40M"
    elif((1800 < height) and (height < 2400)) and ((40 < fps) and (fps < 70)):
        bitrate = "55M"
    else:
        print(">> Unknow resolution.")
        exit()
    print(">> BITRATE: ",bitrate)
    return bitrate

def scale_lr_imgs(imgs):
    """Scale low-res images prior to passing to SRGAN"""
    return imgs / 255.

def unscale_hr_imgs(imgs):
    """Un-Scale high-res images"""
    imgs = imgs * 255
    imgs = np.clip(imgs, 0., 255.)
    return imgs.astype('uint8')

def sr_genarator(model,img_lr,scale):
    """Predict sr frame given a LR frame"""
    # Predict high-resolution version (add batch dimension to image)
    img_lr = cv2.resize(img_lr,(img_lr.shape[1]*scale,img_lr.shape[0]*scale), interpolation = cv2.INTER_CUBIC)
    img_lr = scale_lr_imgs(img_lr)
    img_sr = model.predict(np.expand_dims(img_lr, 0))
    img_sr = unscale_hr_imgs(img_sr)
    # Remove batch dimension
    img_sr = img_sr.reshape(img_sr.shape[1], img_sr.shape[2], img_sr.shape[3])
    return img_sr

def write_srvideo(model=None,lr_videopath=None,sr_videopath=None,scale=None,print_frequency=False,crf=15,fps=None,gpu=False):
    """Generate SR video given LR video """
    videogen = skvideo.io.FFmpegReader(lr_videopath)
    t_frames, height, width, _  = videogen.getShape()
    print(">> Inputshape: ",videogen.getShape())
    metadata = skvideo.io.ffprobe(lr_videopath)
    #print(json.dumps(metadata["video"], indent=4))
    _fps = metadata['video']['@r_frame_rate'] if (fps == None) else str(fps)
    codec = 'h264_nvenc' if (gpu == 'True') else 'libx264'
    writer = skvideo.io.FFmpegWriter(sr_videopath,
    inputdict={'-r': _fps, '-width': str(width*scale-12), '-height': str(height*scale-12)},
    outputdict={'-vcodec': codec, '-r': _fps, '-crf': str(crf), '-pix_fmt': 'yuv420p',
                '-b:v': selectBetterBitrate(height*scale,int(_fps.split('/')[0])/int(_fps.split('/')[1]))})
    count = 0
    time_elapsed = []
    print(">> Writing video...")
    for frame in tqdm(videogen):
        start = timer()
        img_sr = sr_genarator(model,frame,scale=scale)
        writer.writeFrame(img_sr)
        end = timer()
        time_elapsed.append(end - start)
        count +=1
        if (print_frequency):
            if(count % print_frequency == 0):
                print('... Time per Frame: '+str(np.mean(time_elapsed))+'s')
                print('... Estimated time: '+str(np.mean(time_elapsed)*(t_frames-count)/60.)+'min')
    writer.close()
    videogen = skvideo.io.FFmpegReader(sr_videopath)
    print(">> Outputshape: ",videogen.getShape())
    print('>> Video resized in '+str(np.sum(time_elapsed))+'s')
    return time_elapsed

def write_sr_images(model=None, lr_imagepath=None, sr_imagepath=None,scale=None):
    print(">> Writing image...")
    time_elapsed = []
    # Load the images to perform test on images
    img_lr = DataLoader.load_img(lr_imagepath,colorspace='RGB')

    # Create super resolution images
    start = timer()
    img_sr = sr_genarator(model,img_lr,scale)
    end = timer()
    time_elapsed.append(end - start)

    img_sr = Image.fromarray(img_sr.astype(np.uint8))
    img_sr.save(sr_imagepath)
    print('>> Image resized in '+str(np.mean(time_elapsed))+'s')
    return time_elapsed