# -*- coding: utf-8 -*-
"Validate FCOS end to end."
import os
import argparse
# disable autotune
os.environ['MXNET_CUDNN_AUTOTUNE_DEFAULT'] = '0'
import glob
import logging
import time
import numpy as np
import mxnet as mx
from tqdm import tqdm
from mxnet import gluon
from mxnet import autograd
import gluoncv as gcv
from gluoncv import data as gdata
from gluoncv import utils as gutils
from gluoncv.model_zoo import get_model
from gluoncv.data import batchify
from gluoncv.data.transforms.presets.fcos import FCOSDefaultValTransform
from gluoncv.utils.metrics.voc_detection import VOC07MApMetric
from gluoncv.utils.metrics.coco_detection import COCODetectionMetric
from IPython import embed

def parse_args():
    parser = argparse.ArgumentParser(description='Validate FCOS networks.')
    parser.add_argument('--network', type=str, default='resnet50_v1b',
                        help="Base feature extraction network name")
    parser.add_argument('--dataset', type=str, default='coco',
                        help='Validation dataset.')
    parser.add_argument('--num-workers', '-j', dest='num_workers', type=int,
                        default=4, help='Number of data workers')
    parser.add_argument('--gpus', type=str, default='0',
                        help='Training with GPUs, you can specify 1,3 for example.')
    parser.add_argument('--pretrained', type=str, default='True',
                        help='Load weights from previously saved parameters.')
    parser.add_argument('--save-prefix', type=str, default='',
                        help='Saving parameter prefix')
    parser.add_argument('--save-json', action='store_true',
                        help='Save coco output json')
    parser.add_argument('--eval-all', action='store_true',
                        help='Eval all models begins with save prefix. Use with pretrained.')

    args = parser.parse_args()
    return args

def get_dataset(dataset, args):
    if dataset.lower() == 'voc':
        val_dataset = gdata.VOCDetection(
            splits=[(2007, 'test')])
        val_metric = VOC07MApMetric(iou_thresh=0.5, class_names=val_dataset.classes)
    elif dataset.lower() == 'coco':
        val_dataset = gdata.COCODetection(splits='instances_val2017', skip_empty=False)
        val_metric = COCODetectionMetric(val_dataset, args.save_prefix + '_eval',
                                         cleanup=not args.save_json)
    else:
        raise NotImplementedError('Dataset: {} not implemented.'.format(dataset))
    return val_dataset, val_metric

def get_dataloader(net, val_dataset, val_transform, batch_size, num_workers):
    """Get dataloader."""
    val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(4)])
    short = net.short[-1] if isinstance(net.short, (tuple, list)) else net.short
    val_loader = mx.gluon.data.DataLoader(
        val_dataset.transform(val_transform(short, net.max_size, net.base_stride)),
        batch_size, False, batchify_fn=val_bfn, last_batch='keep', num_workers=num_workers)
    return val_loader

def split_and_load(batch, ctx_list):
    """Split data to 1 batch each device."""
    num_ctx = len(ctx_list)
    new_batch = []
    for i, data in enumerate(batch):
        new_data = [x.as_in_context(ctx) for x, ctx in zip(data, ctx_list)]
        new_batch.append(new_data)
    return new_batch

def validate(net, val_data, ctx, eval_metric, size):
    """Test on validation dataset."""
    clipper = gcv.nn.bbox.BBoxClipToImage()
    eval_metric.reset()
    net.hybridize(static_alloc=True)
    nms_thresh = net.nms_thresh
    nms_topk = net.nms_topk
    save_topk = net.save_topk
    with tqdm(total=size) as pbar:
        for batch in val_data:
            batch = split_and_load(batch, ctx_list=ctx)
            det_bboxes = []
            det_ids = []
            det_scores = []
            gt_bboxes = []
            gt_ids = []
            gt_difficults = []
            for x, y, cor, im_scale in zip(*batch):
                # get prediction results
                cls_probs, bboxes = net(x)
                cls_id = cls_probs.argmax(axis=-1)
                probs = mx.nd.pick(cls_probs, cls_id)
                bboxes = net.box_converter(bboxes, cor)
                bboxes = clipper(bboxes.squeeze(axis=0), x)
                im_scale = im_scale.reshape((-1)).asscalar()
                bboxes *= im_scale
                cls_id = cls_id.squeeze(axis=0)
                probs = probs.squeeze(axis=0)
                bboxes = bboxes.squeeze(axis=0)
                target = mx.nd.concat(cls_id.expand_dims(axis=1),
                            probs.expand_dims(axis=1), bboxes, dim=-1)
                keep = mx.nd.contrib.box_nms(target, overlap_thresh=nms_thresh, coord_start=2,
                                             topk=nms_topk, valid_thresh=0.00001, score_index=1,
                                             id_index=0, force_suppress=False,
                                             in_format='corner', out_format='corner')
                keep = keep[:save_topk].expand_dims(axis=0)
                embed()
                det_ids.append(keep.slice_axis(axis=-1, begin=0, end=1))
                det_scores.append(keep.slice_axis(axis=-1, begin=1, end=2))
                det_bboxes.append(keep.slice_axis(axis=-1, begin=2, end=None))
                # split ground truths
                gt_ids.append(y.slice_axis(axis=-1, begin=4, end=5))
                gt_bboxes.append(y.slice_axis(axis=-1, begin=0, end=4))
                gt_bboxes[-1] *= im_scale
                gt_difficults.append(
                        y.slice_axis(axis=-1, begin=5, end=6) if y.shape[-1] > 5 else None)
            # update metric
            for det_bbox, det_id, det_score, gt_bbox, gt_id, gt_diff in \
                zip(det_bboxes, det_ids, det_scores, gt_bboxes, gt_ids, gt_difficults):
                eval_metric.update(det_bbox, det_id, det_score, gt_bbox, gt_id, gt_diff)
            pbar.update(len(ctx))
    return eval_metric.get()

if __name__ == '__main__':
    args = parse_args()

    # validation contexts
    ctx = [mx.gpu(int(i)) for i in args.gpus.split(',') if i.strip()]
    ctx = ctx if ctx else [mx.cpu()]
    args.batch_size = len(ctx) # 1 batch per device

    # network
    kwargs = {}
    net_name = "_".join(("fcos", args.network, args.dataset))
    args.save_prefix += net_name
    if args.pretrained.lower() in ['true', '1', 'yes', 't']:
        net = gcv.model_zoo.get_model(net_name, pretrained=True, **kwargs)
    else:
        net = gcv.model_zoo.get_model(net_name, pretrained=False, **kwargs)
        net.load_parameters(args.pretrained.strip(), ignore_extra=True)
    net.collect_params().reset_ctx(ctx)
    embed()

    # validation data
    val_dataset, eval_metric = get_dataset(args.dataset, args)
    val_data = get_dataloader(
        net, val_dataset, FCOSDefaultValTransform, args.batch_size, args.num_workers)

    # validation
    if not args.eval_all:
        names, values = validate(net, val_data, ctx, eval_metric, len(val_dataset))
        for k, v in zip(names, values):
            print(k, v)
    else:
        saved_models = glob.glob(args.save_prefix + '*.params')
        for epoch. saved_model in enumerate(sorted(saved_models)):
            print('[Epoch {}] Validating from {}'.format(epoch, saved_model))
            net.load_parameters(saved_model)
            net.collect_params().reset_ctx(ctx)
            map_name, mean_ap = validate(net, val_data, ctx, eval_metric, len(val_dataset))
            val_msg = '\n'.join(['{}={}'.format(k, v) for k, v in zip(map_name, mean_ap)])
            print('[Epoch {}] Validation: \n{}'.format(epoch, val_msg))
            current_map = float(mean_ap[-1])
            with open(args.save_prefix+'_best_map.log', 'a') as f:
                f.write('\n{:04d}:\t{:.4f}'.format(epoch, current_map))


import numpy as np
from gym import utils
from gym.envs.mujoco import mujoco_env
from inverse_rl.envs.dynamic_mjc.model_builder import MJCModel
from rllab.misc import logger

def ant_env(gear=150, eyes=True):
    mjcmodel = MJCModel('ant_maze')
    mjcmodel.root.compiler(inertiafromgeom="true", angle="degree", coordinate="local")
    mjcmodel.root.option(timestep="0.01", integrator="RK4")
    mjcmodel.root.custom().numeric(data="0.0 0.0 0.55 1.0 0.0 0.0 0.0 0.0 1.0 0.0 -1.0 0.0 -1.0 0.0 1.0",name="init_qpos")
    asset = mjcmodel.root.asset()
    asset.texture(builtin="gradient",height="100",rgb1="1 1 1",rgb2="0 0 0",type="skybox",width="100")
    asset.texture(builtin="flat",height="1278",mark="cross",markrgb="1 1 1",name="texgeom",random="0.01",rgb1="0.8 0.6 0.4",rgb2="0.8 0.6 0.4",type="cube",width="127")
    asset.texture(builtin="checker",height="100",name="texplane",rgb1="0 0 0",rgb2="0.8 0.8 0.8",type="2d",width="100")
    asset.material(name="MatPlane",reflectance="0.5",shininess="1",specular="1",texrepeat="60 60",texture="texplane")
    asset.material(name="geom",texture="texgeom",texuniform="true")

    default = mjcmodel.root.default()
    default.joint(armature=1, damping=1, limited='true')
    default.geom(friction=[1.5,0.5,0.5], density=5.0, margin=0.01, condim=3, conaffinity=0, rgba="0.8 0.6 0.4 1")

    worldbody = mjcmodel.root.worldbody()
    worldbody.light(cutoff="100",diffuse=[.8,.8,.8],dir="-0 0 -1.3",directional="true",exponent="1",pos="0 0 1.3",specular=".1 .1 .1")
    worldbody.geom(conaffinity=1, condim=3, material="MatPlane",name="floor",pos="0 0 0",rgba="0.8 0.9 0.8 1",size="40 40 40",type="plane")

    ant = worldbody.body(name='torso', pos=[0, 0, 0.75])
    ant.geom(name='torso_geom', pos=[0, 0, 0], size="0.25", type="sphere")
    ant.joint(armature="0", damping="0", limited="false", margin="0.01", name="root", pos=[0, 0, 0], type="free")

    if eyes:
        eye_z = 0.1
        eye_y = -.21
        eye_x_offset = 0.07
        # eyes
        ant.geom(fromto=[eye_x_offset,0,eye_z,eye_x_offset,eye_y,eye_z], name='eye1', size='0.03', type='capsule', rgba=[1,1,1,1])
        ant.geom(fromto=[eye_x_offset,0,eye_z,eye_x_offset,eye_y-0.02,eye_z], name='eye1_', size='0.02', type='capsule', rgba=[0,0,0,1])
        ant.geom(fromto=[-eye_x_offset,0,eye_z,-eye_x_offset,eye_y,eye_z], name='eye2', size='0.03', type='capsule', rgba=[1,1,1,1])
        ant.geom(fromto=[-eye_x_offset,0,eye_z,-eye_x_offset,eye_y-0.02,eye_z], name='eye2_', size='0.02', type='capsule', rgba=[0,0,0,1])
        # eyebrows
        ant.geom(fromto=[eye_x_offset-0.03,eye_y, eye_z+0.07, eye_x_offset+0.03, eye_y, eye_z+0.1], name='brow1', size='0.02', type='capsule', rgba=[0,0,0,1])
        ant.geom(fromto=[-eye_x_offset+0.03,eye_y, eye_z+0.07, -eye_x_offset-0.03, eye_y, eye_z+0.1], name='brow2', size='0.02', type='capsule', rgba=[0,0,0,1])

    front_left_leg = ant.body(name="front_left_leg", pos=[0, 0, 0])
    front_left_leg.geom(fromto=[0.0, 0.0, 0.0, 0.2, 0.2, 0.0], name="aux_1_geom", size="0.08", type="capsule")
    aux_1 = front_left_leg.body(name="aux_1", pos=[0.2, 0.2, 0])
    aux_1.joint(axis=[0, 0, 1], name="hip_1", pos=[0.0, 0.0, 0.0], range=[-30, 30], type="hinge")
    aux_1.geom(fromto=[0.0, 0.0, 0.0, 0.2, 0.2, 0.0], name="left_leg_geom", size="0.08", type="capsule")
    ankle_1 = aux_1.body(pos=[0.2, 0.2, 0])
    ankle_1.joint(axis=[-1, 1, 0], name="ankle_1", pos=[0.0, 0.0, 0.0], range=[30, 70], type="hinge")
    ankle_1.geom(fromto=[0.0, 0.0, 0.0, 0.4, 0.4, 0.0], name="left_ankle_geom", size="0.08", type="capsule")

    front_right_leg = ant.body(name="front_right_leg", pos=[0, 0, 0])
    front_right_leg.geom(fromto=[0.0, 0.0, 0.0, -0.2, 0.2, 0.0], name="aux_2_geom", size="0.08", type="capsule")
    aux_2 = front_right_leg.body(name="aux_2", pos=[-0.2, 0.2, 0])
    aux_2.joint(axis=[0, 0, 1], name="hip_2", pos=[0.0, 0.0, 0.0], range=[-30, 30], type="hinge")
    aux_2.geom(fromto=[0.0, 0.0, 0.0, -0.2, 0.2, 0.0], name="right_leg_geom", size="0.08", type="capsule")
    ankle_2 = aux_2.body(pos=[-0.2, 0.2, 0])
    ankle_2.joint(axis=[1, 1, 0], name="ankle_2", pos=[0.0, 0.0, 0.0], range=[-70, -30], type="hinge")
    ankle_2.geom(fromto=[0.0, 0.0, 0.0, -0.4, 0.4, 0.0], name="right_ankle_geom", size="0.08", type="capsule")

    back_left_leg = ant.body(name="back_left_leg", pos=[0, 0, 0])
    back_left_leg.geom(fromto=[0.0, 0.0, 0.0, -0.2, -0.2, 0.0], name="aux_3_geom", size="0.08", type="capsule")
    aux_3 = back_left_leg.body(name="aux_3", pos=[-0.2, -0.2, 0])
    aux_3.joint(axis=[0, 0, 1], name="hip_3", pos=[0.0, 0.0, 0.0], range=[-30, 30], type="hinge")
    aux_3.geom(fromto=[0.0, 0.0, 0.0, -0.2, -0.2, 0.0], name="backleft_leg_geom", size="0.08", type="capsule")
    ankle_3 = aux_3.body(pos=[-0.2, -0.2, 0])
    ankle_3.joint(axis=[-1, 1, 0], name="ankle_3", pos=[0.0, 0.0, 0.0], range=[-70, -30], type="hinge")
    ankle_3.geom(fromto=[0.0, 0.0, 0.0, -0.4, -0.4, 0.0], name="backleft_ankle_geom", size="0.08", type="capsule")

    back_right_leg = ant.body(name="back_right_leg", pos=[0, 0, 0])
    back_right_leg.geom(fromto=[0.0, 0.0, 0.0, 0.2, -0.2, 0.0], name="aux_4_geom", size="0.08", type="capsule")
    aux_4 = back_right_leg.body(name="aux_4", pos=[0.2, -0.2, 0])
    aux_4.joint(axis=[0, 0, 1], name="hip_4", pos=[0.0, 0.0, 0.0], range=[-30, 30], type="hinge")
    aux_4.geom(fromto=[0.0, 0.0, 0.0, 0.2, -0.2, 0.0], name="backright_leg_geom", size="0.08", type="capsule")
    ankle_4 = aux_4.body(pos=[0.2, -0.2, 0])
    ankle_4.joint(axis=[1, 1, 0], name="ankle_4", pos=[0.0, 0.0, 0.0], range=[30, 70], type="hinge")
    ankle_4.geom(fromto=[0.0, 0.0, 0.0, 0.4, -0.4, 0.0], name="backright_ankle_geom", size="0.08", type="capsule")

    actuator = mjcmodel.root.actuator()
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_4", gear=gear)
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_4", gear=gear)
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_1", gear=gear)
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_1", gear=gear)
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_2", gear=gear)
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_2", gear=gear)
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_3", gear=gear)
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_3", gear=gear)
    return mjcmodel

def angry_ant_crippled(gear=150):
    mjcmodel = MJCModel('ant_maze')
    mjcmodel.root.compiler(inertiafromgeom="true", angle="degree", coordinate="local")
    mjcmodel.root.option(timestep="0.01", integrator="RK4")
    mjcmodel.root.custom().numeric(data="0.0 0.0 0.55 1.0 0.0 0.0 0.0 0.0 1.0 0.0 -1.0 0.0 -1.0 0.0 1.0",name="init_qpos")
    asset = mjcmodel.root.asset()
    asset.texture(builtin="gradient",height="100",rgb1="1 1 1",rgb2="0 0 0",type="skybox",width="100")
    asset.texture(builtin="flat",height="1278",mark="cross",markrgb="1 1 1",name="texgeom",random="0.01",rgb1="0.8 0.6 0.4",rgb2="0.8 0.6 0.4",type="cube",width="127")
    asset.texture(builtin="checker",height="100",name="texplane",rgb1="0 0 0",rgb2="0.8 0.8 0.8",type="2d",width="100")
    asset.material(name="MatPlane",reflectance="0.5",shininess="1",specular="1",texrepeat="60 60",texture="texplane")
    asset.material(name="geom",texture="texgeom",texuniform="true")

    default = mjcmodel.root.default()
    default.joint(armature=1, damping=1, limited='true')
    default.geom(friction=[1.5,0.5,0.5], density=5.0, margin=0.01, condim=3, conaffinity=0, rgba="0.8 0.6 0.4 1")

    worldbody = mjcmodel.root.worldbody()

    worldbody.geom(conaffinity=1, condim=3, material="MatPlane",name="floor",pos="0 0 0",rgba="0.8 0.9 0.8 1",size="40 40 40",type="plane")
    worldbody.light(cutoff="100",diffuse=[.8,.8,.8],dir="-0 0 -1.3",directional="true",exponent="1",pos="0 0 1.3",specular=".1 .1 .1")

    ant = worldbody.body(name='torso', pos=[0, 0, 0.75])
    ant.geom(name='torso_geom', pos=[0, 0, 0], size="0.25", type="sphere")
    ant.joint(armature="0", damping="0", limited="false", margin="0.01", name="root", pos=[0, 0, 0], type="free")

    eye_z = 0.1
    eye_y = -.21
    eye_x_offset = 0.07
    # eyes
    ant.geom(fromto=[eye_x_offset,0,eye_z,eye_x_offset,eye_y,eye_z], name='eye1', size='0.03', type='capsule', rgba=[1,1,1,1])
    ant.geom(fromto=[eye_x_offset,0,eye_z,eye_x_offset,eye_y-0.02,eye_z], name='eye1_', size='0.02', type='capsule', rgba=[0,0,0,1])
    ant.geom(fromto=[-eye_x_offset,0,eye_z,-eye_x_offset,eye_y,eye_z], name='eye2', size='0.03', type='capsule', rgba=[1,1,1,1])
    ant.geom(fromto=[-eye_x_offset,0,eye_z,-eye_x_offset,eye_y-0.02,eye_z], name='eye2_', size='0.02', type='capsule', rgba=[0,0,0,1])
    # eyebrows
    ant.geom(fromto=[eye_x_offset-0.03,eye_y, eye_z+0.07, eye_x_offset+0.03, eye_y, eye_z+0.1], name='brow1', size='0.02', type='capsule', rgba=[0,0,0,1])
    ant.geom(fromto=[-eye_x_offset+0.03,eye_y, eye_z+0.07, -eye_x_offset-0.03, eye_y, eye_z+0.1], name='brow2', size='0.02', type='capsule', rgba=[0,0,0,1])

    front_left_leg = ant.body(name="front_left_leg", pos=[0, 0, 0])
    front_left_leg.geom(fromto=[0.0, 0.0, 0.0, 0.2, 0.2, 0.0], name="aux_1_geom", size="0.08", type="capsule")
    aux_1 = front_left_leg.body(name="aux_1", pos=[0.2, 0.2, 0])
    aux_1.joint(axis=[0, 0, 1], name="hip_1", pos=[0.0, 0.0, 0.0], range=[-30, 30], type="hinge")
    aux_1.geom(fromto=[0.0, 0.0, 0.0, 0.2, 0.2, 0.0], name="left_leg_geom", size="0.08", type="capsule")
    ankle_1 = aux_1.body(pos=[0.2, 0.2, 0])
    ankle_1.joint(axis=[-1, 1, 0], name="ankle_1", pos=[0.0, 0.0, 0.0], range=[30, 70], type="hinge")
    ankle_1.geom(fromto=[0.0, 0.0, 0.0, 0.4, 0.4, 0.0], name="left_ankle_geom", size="0.08", type="capsule")

    front_right_leg = ant.body(name="front_right_leg", pos=[0, 0, 0])
    front_right_leg.geom(fromto=[0.0, 0.0, 0.0, -0.2, 0.2, 0.0], name="aux_2_geom", size="0.08", type="capsule")
    aux_2 = front_right_leg.body(name="aux_2", pos=[-0.2, 0.2, 0])
    aux_2.joint(axis=[0, 0, 1], name="hip_2", pos=[0.0, 0.0, 0.0], range=[-30, 30], type="hinge")
    aux_2.geom(fromto=[0.0, 0.0, 0.0, -0.2, 0.2, 0.0], name="right_leg_geom", size="0.08", type="capsule")
    ankle_2 = aux_2.body(pos=[-0.2, 0.2, 0])
    ankle_2.joint(axis=[1, 1, 0], name="ankle_2", pos=[0.0, 0.0, 0.0], range=[-70, -30], type="hinge")
    ankle_2.geom(fromto=[0.0, 0.0, 0.0, -0.4, 0.4, 0.0], name="right_ankle_geom", size="0.08", type="capsule")

    # Back left leg is crippled
    thigh_length = 0.1 #0.2
    ankle_length = 0.2 #0.4
    dark_red = [0.8,0.3,0.3,1.0]

    back_left_leg = ant.body(name="back_left_leg", pos=[0, 0, 0])
    back_left_leg.geom(fromto=[0.0, 0.0, 0.0, -0.2, -0.2, 0.0], name="aux_3_geom", size="0.08", type="capsule",
                       rgba=dark_red)
    aux_3 = back_left_leg.body(name="aux_3", pos=[-0.2, -0.2, 0])
    aux_3.joint(axis=[0, 0, 1], name="hip_3", pos=[0.0, 0.0, 0.0], range=[-30, 30], type="hinge")
    aux_3.geom(fromto=[0.0, 0.0, 0.0, -thigh_length, -thigh_length, 0.0], name="backleft_leg_geom", size="0.08", type="capsule",
               rgba=dark_red)
    ankle_3 = aux_3.body(pos=[-thigh_length, -thigh_length, 0])
    ankle_3.joint(axis=[-1, 1, 0], name="ankle_3", pos=[0.0, 0.0, 0.0], range=[-70, -30], type="hinge")
    ankle_3.geom(fromto=[0.0, 0.0, 0.0, -ankle_length, -ankle_length, 0.0], name="backleft_ankle_geom", size="0.08", type="capsule",
                 rgba=dark_red)

    back_right_leg = ant.body(name="back_right_leg", pos=[0, 0, 0])
    back_right_leg.geom(fromto=[0.0, 0.0, 0.0, 0.2, -0.2, 0.0], name="aux_4_geom", size="0.08", type="capsule",
                        rgba=dark_red)
    aux_4 = back_right_leg.body(name="aux_4", pos=[0.2, -0.2, 0])
    aux_4.joint(axis=[0, 0, 1], name="hip_4", pos=[0.0, 0.0, 0.0], range=[-30, 30], type="hinge")
    aux_4.geom(fromto=[0.0, 0.0, 0.0, thigh_length, -thigh_length, 0.0], name="backright_leg_geom", size="0.08", type="capsule",
               rgba=dark_red)
    ankle_4 = aux_4.body(pos=[thigh_length, -thigh_length, 0])
    ankle_4.joint(axis=[1, 1, 0], name="ankle_4", pos=[0.0, 0.0, 0.0], range=[30, 70], type="hinge")
    ankle_4.geom(fromto=[0.0, 0.0, 0.0, ankle_length, -ankle_length, 0.0], name="backright_ankle_geom", size="0.08", type="capsule",
                 rgba=dark_red)

    actuator = mjcmodel.root.actuator()
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_1", gear=gear)
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_1", gear=gear)
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_2", gear=gear)
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_2", gear=gear)
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_3", gear=1) # cripple the joints
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_3", gear=1) # cripple the joints
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_4", gear=1)
    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_4", gear=1)
    return mjcmodel

class CustomAntEnv(mujoco_env.MujocoEnv, utils.EzPickle):
    """
    A modified ant env with lower joint gear ratios so it flips less often and learns faster.
    """
    def __init__(self, max_timesteps=1000, disabled=False, gear=150):
        #mujoco_env.MujocoEnv.__init__(self, 'ant.xml', 5)
        utils.EzPickle.__init__(self)
        self.timesteps = 0
        self.max_timesteps=max_timesteps

        if disabled:
            model = angry_ant_crippled(gear=gear)
        else:
            model = ant_env(gear=gear)
        with model.asfile() as f:
            mujoco_env.MujocoEnv.__init__(self, f.name, 5)

    def step(self,a):
        return self._step(a)

    def _step(self, a):
        vel = self.data.qvel.flat[0]
        forward_reward = vel
        self.do_simulation(a, self.frame_skip)

        ctrl_cost = .01 * np.square(a).sum()
        contact_cost = 0.5 * 1e-3 * np.sum(
            np.square(np.clip(self.data.cfrc_ext, -1, 1)))
        state = self.state_vector()
        flipped = not (state[2] >= 0.2)
        flipped_rew = -1 if flipped else 0
        reward = forward_reward - ctrl_cost - contact_cost +flipped_rew

        self.timesteps += 1
        done = self.timesteps >= self.max_timesteps

        ob = self._get_obs()
        return ob, reward, done, dict(
            reward_forward=forward_reward,
            reward_ctrl=-ctrl_cost,
            reward_contact=-contact_cost,
            reward_flipped=flipped_rew)

    def _get_obs(self):
        return np.concatenate([
            self.data.qpos.flat[2:],
            self.data.qvel.flat,
            np.clip(self.data.cfrc_ext, -1, 1).flat,
        ])

    def reset_model(self):
        self.timesteps = 0
        qpos = self.init_qpos + self.np_random.uniform(size=self.model.nq, low=-.1, high=.1)
        qvel = self.init_qvel + self.np_random.randn(self.model.nv) * .1
        self.set_state(qpos, qvel)
        return self._get_obs()

    def viewer_setup(self):
        self.viewer.cam.distance = self.model.stat.extent * 0.5

    def log_diagnostics(self, paths):
        forward_rew = np.array([np.mean(traj['env_infos']['reward_forward']) for traj in paths])
        reward_ctrl = np.array([np.mean(traj['env_infos']['reward_ctrl']) for traj in paths])
        reward_cont = np.array([np.mean(traj['env_infos']['reward_contact']) for traj in paths])
        reward_flip = np.array([np.mean(traj['env_infos']['reward_flipped']) for traj in paths])

        logger.record_tabular('AvgRewardFwd', np.mean(forward_rew))
        logger.record_tabular('AvgRewardCtrl', np.mean(reward_ctrl))
        logger.record_tabular('AvgRewardContact', np.mean(reward_cont))
        logger.record_tabular('AvgRewardFlipped', np.mean(reward_flip))

if __name__ == "__main__":
    env = CustomAntEnv(disabled=True, gear=30)

    for _ in range(1000):
        env.render()
        env.step(env.action_space.sample())


import numpy as np

class PlateCompare:
    def __init__(self):
        self.cum_plateWithCharMatchCnt = 0
        self.cum_plateFrameMatchCnt = 0
        self.cum_plateCnt_gt = 0
        self.cum_plateCnt_pred = 0
        self.cum_charMatchCnt = 0
        self.cum_charCnt_gt = 0
        self.cum_charCnt_pred = 0

    # generate some stats for the images analysed
    # platesWithCharCorrect_recall - Plates detected in correct location and containing correct characters in the correct locations
    #                               divided by the number of ground truth plates
    # platesWithCharCorrect_precision - Plates detected in correct location and containing correct characters in the correct locations
    #                                  divided by the total number of plates predicted (ie true pos plus false pos)
    # plateFrames_recall - Plate frames in correct location (no checking of characters) divided by the
    #                     number of ground truth plates
    # plateFrames_precision - Plate frames in correct location (no checking of characters) divided by the
    #                        the total number of plates predicted (ie true pos plus false pos)
    # chars_recall - Characters detected in the correct place with the correct contents
    #               divided by the number of ground truth characters
    # chars_precision - Chars detected outside of the correct location, or the location is correct,
    #                  but the contents are wrong. Divided by the total number of plates predicted (ie true pos plus false pos)
    def calcStats(self):
        platesWithCharCorrect_recall = (
            self.cum_plateWithCharMatchCnt / self.cum_plateCnt_gt
        )
        platesWithCharCorrect_precision = (
            self.cum_plateWithCharMatchCnt / self.cum_plateCnt_pred
        )
        plateFrames_recall = self.cum_plateFrameMatchCnt / self.cum_plateCnt_gt
        plateFrames_precision = self.cum_plateFrameMatchCnt / self.cum_plateCnt_pred
        chars_recall = self.cum_charMatchCnt / self.cum_charCnt_gt
        chars_precision = self.cum_charMatchCnt / self.cum_charCnt_pred
        return (
            platesWithCharCorrect_recall,
            platesWithCharCorrect_precision,
            plateFrames_recall,
            plateFrames_precision,
            chars_recall,
            chars_precision,
        )

    # calculate the intersection over union of two boxes
    def intersectionOverUnion(self, box1, box2):
        (box1StartY, box1StartX, box1EndY, box1EndX) = box1
        (box2StartY, box2StartX, box2EndY, box2EndX) = box2
        # determine the (x, y)-coordinates of the intersection rectangle
        xA = max(box2StartX, box1StartX)
        yA = max(box2StartY, box1StartY)
        xB = min(box2EndX, box1EndX)
        yB = min(box2EndY, box1EndY)

        # if the boxes are intersecting, then compute the area of intersection rectangle
        if xB > xA and yB > yA:
            interArea = (xB - xA) * (yB - yA)
        else:
            interArea = 0.0

        # compute the area of the box1 and box2
        box1Area = (box1EndY - box1StartY) * (box1EndX - box1StartX)
        box2Area = (box2EndY - box2StartY) * (box2EndX - box2StartX)

        # compute the intersection area / box1 area
        iou = interArea / float(box1Area + box2Area - interArea)

        # return the intersection over area value
        return iou

    # compare ground truth and predicted results
    # return stats for this compare and keep cumulative stats in instance variables that can be retrieved
    # using calcStats
    def comparePlates(
        self,
        plateBoxes_gt,
        charBoxes_gt,
        charTexts_gt,
        plate_boxes_pred,
        charBoxes_pred,
        charTexts_pred,
    ):
        # set mask to all true
        maskPlate_pred = np.ones(len(plate_boxes_pred), dtype=bool)

        matchIndices = []
        plateFrameMatchCnt = 0
        charCntTotal_gt = 0
        # loop over the ground truth and predicted plate boxes
        for (i, plateBox_gt) in enumerate(plateBoxes_gt):
            # keep a count of the total number of gt chars
            charCntTotal_gt += len(charBoxes_gt[i])
            for (j, plateBox_pred) in enumerate(plate_boxes_pred):
                iou = self.intersectionOverUnion(plateBox_gt, plateBox_pred)
                if iou >= 0.5 and maskPlate_pred[j] == True:
                    # maintain a list of indices for matching gt and pred plate boxes
                    matchIndices.append((i, j))
                    maskPlate_pred[j] = False
                    plateFrameMatchCnt += 1
                    # we have found a match for this plate gt, so move onto the next one
                    continue

        # init the counters prior to loop
        plateWithCharMatchCnt = 0
        charCntTotal_pred = 0
        charMatchCntTotal = 0

        # loop over all the matching plate boxes found in the previous code block
        if len(plate_boxes_pred) != len(charBoxes_pred):
            print("mismatch")
        for (i, j) in matchIndices:
            # if j >= len(charBoxes_pred):
            #  continue
            # set mask to all true
            maskChar_pred = np.ones(len(charBoxes_pred[j]), dtype=bool)

            # grab all the char boxes and text for the current plate
            chBoxes_gt = charBoxes_gt[i]
            chBoxes_pred = charBoxes_pred[j]
            chTexts_gt = charTexts_gt[i]
            chTexts_pred = charTexts_pred[j]

            charMatchCnt = 0
            # loop over the ground truth and predicted char boxes
            for (k, chBox_gt) in enumerate(chBoxes_gt):
                for (l, chBox_pred) in enumerate(chBoxes_pred):
                    # If the iou is greater than 0.5 and predChar is not already matched, then check
                    # if the chars match and increment the match count
                    iou = self.intersectionOverUnion(chBox_gt, chBox_pred)
                    if iou >= 0.5 and maskChar_pred[l] == True:
                        maskChar_pred[l] = False
                        if chTexts_gt[k] == chTexts_pred[l]:
                            charMatchCnt += 1
                        # we have found a match for this gt char, so move onto the next gt char
                        continue

            # check for perfect plate match. ie plate boxes match, and char boxes and labels match
            if charMatchCnt == len(chBoxes_gt) and charMatchCnt == len(chBoxes_pred):
                plateWithCharMatchCnt += 1

            # update the counters
            charMatchCntTotal += charMatchCnt
            charCntTotal_pred += len(chBoxes_pred)

        plateCntTotal_gt = len(plateBoxes_gt)
        plateCntTotal_pred = len(plate_boxes_pred)

        # maintain the cumulative scores
        self.cum_plateWithCharMatchCnt += plateWithCharMatchCnt
        self.cum_plateFrameMatchCnt += plateFrameMatchCnt
        self.cum_plateCnt_gt += plateCntTotal_gt
        self.cum_plateCnt_pred += plateCntTotal_pred
        self.cum_charMatchCnt += charMatchCntTotal
        self.cum_charCnt_gt += charCntTotal_gt
        self.cum_charCnt_pred += charCntTotal_pred

        return (
            plateWithCharMatchCnt,
            plateFrameMatchCnt,
            plateCntTotal_gt,
            plateCntTotal_pred,
            charMatchCntTotal,
            charCntTotal_gt,
            charCntTotal_pred,
        )


import numpy as np
from flask import Flask, request, render_template
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.preprocessing.text import one_hot
from tensorflow.keras.models import load_model
import nltk
import re
from nltk.corpus import stopwords
nltk.download('stopwords')

from nltk.stem.porter import PorterStemmer
ps = PorterStemmer()

# load model
model = load_model('model.h5')

app = Flask(__name__)

@app.route('/')
def home():
    return render_template('index.html')
corpus = []
@app.route('/predict',methods=['POST'])
def predict():
    message=request.form['message']
    data=str(message)
    review=re.sub('[^a-zA-Z]',' ',data)
    review = review.lower()
    review = review.split()
    review = [ps.stem(word) for word in review if not word in stopwords.words('english')]
    review = ' '.join(review)
    corpus.append(review)
    onehot_rep = [one_hot(words,5000)for words in corpus]
    embedded_docs = pad_sequences(onehot_rep,padding='pre',maxlen=25)
    test_final = np.array(embedded_docs)
    my_prediction=model.predict_classes(test_final)
    return render_template('result.html',prediction=my_prediction)

if __name__ == "__main__":
    app.run(debug=True)


#! /usr/bin/env python3
import argparse
import gym
import time
import ray
import threading
import roboschool
import util
import scipy.signal

import numpy as np
import tensorflow as tf
import tensorlayer as tl

from tabulate import tabulate
from gym import wrappers
from collections import OrderedDict
from sklearn.utils import shuffle

parser = argparse.ArgumentParser(description='ppo adaptive KL algorithm')

parser.add_argument(
    '--gamma', default=0.995, type=float, help='gamma')

parser.add_argument(
    '--lambda_gae', default=0.98, type=float, help='lambda for GAE')

parser.add_argument(
    '--log_vars', default=0.0, type=float, help='init action log variance')

parser.add_argument(
    '--eta', default=50, type=float, help='actor loss parameter')

parser.add_argument(
    '--actor_lr', default=3e-4, type=float, help='learning rate for actor')

parser.add_argument(
    '--critic_lr', default=1e-3, type=float, help='learning rate for critic')

parser.add_argument(
    '--kl_targ', default=0.003, type=float, help='kl divergence target')

parser.add_argument(
    '--nums_worker', default=8, type=int, help='number of workers')

parser.add_argument(
    '--train_epochs', default=10, type=int, help='training epochs')

parser.add_argument(
    '--batch_size', default=20, type=int, help='trianing batch size')

parser.add_argument(
    '--training_steps', default=4000, type=int, help='steps number for training')

parser.add_argument(
    '--max_episodes', default=1000000000, type=int, help='max trianing episodes')

parser.add_argument(
    '--animate', default=False, type=bool, help='whether to animate environment')

parser.add_argument(
    '--save_network', default=False, type=bool, help='whether to save network')

parser.add_argument(
    '--load_network', default=False, type=bool, help='whether to load network')

parser.add_argument(
    '--test_algorithm', default=False, type=bool, help='wether to test algorithm')

parser.add_argument(
    '--eval_algorithm', default=False, type=bool, help='whether to evaluate algorithm')

parser.add_argument(
    '--env_name', default='RoboschoolAnt-v1', type=str, help='gym env name')

parser.add_argument(
    '--model_name', default='ppo', type=str, help='save or load model name')

# global value
args = parser.parse_args()

class PPO(object):
    def __init__(self, env, args):
        self.init_param(env, args)
        self.session = tf.InteractiveSession()

        self._build_ph()
        self.model = self._build_network()
        self._build_trainning()
        self._build_summary()

        self.merge_all = tf.summary.merge_all()
        self.writer = tf.summary.FileWriter('../tensorboard/ppo/{}'.format(
            time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
            ), self.session.graph)
        self.session.run(tf.global_variables_initializer())

    def _build_ph(self):
        self.obs_ph = tf.placeholder(tf.float32, [None, self.obs_dim],
                                     'obs_ph')
        self.act_ph = tf.placeholder(tf.float32, [None, self.act_dim],
                                     'act_ph')
        self.adv_ph = tf.placeholder(tf.float32, [None, ],
                                     'adv_ph')
        self.ret_ph = tf.placeholder(tf.float32, [None, ],
                                     'ret_ph')

        self.lr_ph = tf.placeholder(tf.float32, name='lr_ph')

        self.beta_ph = tf.placeholder(tf.float32, name='beta_ph')

        self.old_log_vars_ph = tf.placeholder(tf.float32, [self.act_dim, ],
                                              'old_log_vars')
        self.old_means_ph = tf.placeholder(tf.float32, [None, self.act_dim],
                                           'old_means')

    def _build_network(self):
        # build actor network
        hid1_size = self.obs_dim * 10
        hid3_size = self.act_dim * 10
        hid2_size = int(np.sqrt(hid1_size * hid3_size))

        self.actor_network = tl.layers.InputLayer(self.obs_ph, name = 'actor_network_input')
        self.actor_network = tl.layers.DenseLayer(self.actor_network, n_units = hid1_size, act = tf.nn.tanh,
            W_init = tf.random_normal_initializer(stddev=np.sqrt(1.0 / self.obs_dim)), name = 'actor_tanh1')
        self.actor_network = tl.layers.DenseLayer(self.actor_network, n_units = hid2_size, act = tf.nn.tanh,
            W_init = tf.random_normal_initializer(stddev=np.sqrt(1.0 / float(hid1_size))), name = 'actor_tanh2')
        self.actor_network = tl.layers.DenseLayer(self.actor_network, n_units = hid3_size, act = tf.nn.tanh,
            W_init = tf.random_normal_initializer(stddev=np.sqrt(1.0 / float(hid2_size))), name = 'actor_tanh3')
        self.actor_network = tl.layers.DenseLayer(self.actor_network, n_units = self.act_dim, act = tf.nn.tanh,
            W_init = tf.random_normal_initializer(stddev=np.sqrt(1.0 / float(hid3_size))), name = 'means')

        # build critic network
        hid1_size = self.obs_dim * 10
        hid3_size = 5
        hid2_size = int(np.sqrt(hid1_size * hid3_size))

        self.critic_network = tl.layers.InputLayer(self.obs_ph, name = 'critic_network_input')
        self.critic_network = tl.layers.DenseLayer(self.critic_network, n_units = hid1_size, act = tf.nn.tanh,
            W_init = tf.random_normal_initializer(stddev=np.sqrt(1.0 / self.obs_dim)), name = 'critic_tanh1')
        self.critic_network = tl.layers.DenseLayer(self.critic_network, n_units = hid2_size, act = tf.nn.tanh,
            W_init = tf.random_normal_initializer(stddev=np.sqrt(1.0 / float(hid1_size))), name = 'critic_tanh2')
        self.critic_network = tl.layers.DenseLayer(self.critic_network, n_units = hid3_size, act = tf.nn.tanh,
            W_init = tf.random_normal_initializer(stddev=np.sqrt(1.0 / float(hid2_size))), name = 'critic_tanh3')
        self.critic_network = tl.layers.DenseLayer(self.critic_network, n_units = 1, act = tf.nn.tanh,
            W_init = tf.random_normal_initializer(stddev=np.sqrt(1.0 / float(hid3_size))), name = 'value')

        # build variance network
        logvar_speed = (10 * hid3_size) // 48
        log_vars = tf.get_variable('logvars', (logvar_speed, self.act_dim), tf.float32,
                                   tf.constant_initializer(0.0))

        self.means = self.actor_network.outputs
        self.log_vars = tf.reduce_sum(log_vars, axis=0) - self.args.log_vars
        self.value = self.critic_network.outputs

        # sample action from norm distributiion
        with tf.variable_scope('sample_action'):
            self.sampled_act = (self.means +
                                tf.exp(self.log_vars / 2.0) *
                                tf.random_normal(shape=(self.act_dim,)))

        return [self.actor_network, self.critic_network]

    def _build_trainning(self):
        # logprob
        self.logp = -0.5 * tf.reduce_sum(self.log_vars) + -0.5 * tf.reduce_sum(tf.square(self.act_ph - self.means) / tf.exp(self.log_vars), axis=1)
        self.logp_old = -0.5 * tf.reduce_sum(self.old_log_vars_ph) + -0.5 * tf.reduce_sum(tf.square(self.act_ph - self.old_means_ph) / tf.exp(self.old_log_vars_ph), axis=1)

        with tf.variable_scope('kl'):
            self.kl = 0.5 * tf.reduce_mean(tf.reduce_sum(tf.exp(self.old_log_vars_ph - self.log_vars)) +
                tf.reduce_sum(tf.square(self.means - self.old_means_ph) / tf.exp(self.log_vars), axis=1) -
                self.act_dim +
                tf.reduce_sum(self.log_vars) - tf.reduce_sum(self.old_log_vars_ph))

        with tf.variable_scope('entropy'):
            self.entropy = 0.5 * (self.act_dim * (np.log(2 * np.pi) + 1) +
                                  tf.reduce_sum(self.log_vars))

        with tf.variable_scope('actor_loss'):
            loss1 = -tf.reduce_mean(self.adv_ph * tf.exp(self.logp - self.logp_old))
            loss2 = tf.reduce_mean(self.beta_ph * self.kl)
            loss3 = self.args.eta * tf.square(tf.maximum(0.0, self.kl - 2.0 * self.args.kl_targ))
            self.actor_loss = loss1 + loss2 + loss3

        self.actor_opt = tf.train.AdamOptimizer(self.lr_ph).minimize(self.actor_loss)

        with tf.variable_scope('critic_loss'):
            self.critic_loss = tf.reduce_mean(tf.square(tf.squeeze(self.value) - self.ret_ph))
        self.critic_opt = tf.train.AdamOptimizer(self.args.critic_lr).minimize(self.critic_loss)

    def _build_summary(self):
        self.score_tb = tf.placeholder(tf.float32, name='score_tb')
        self.actor_loss_tb = tf.placeholder(tf.float32, name='actor_loss_tb')
        self.critic_loss_tb = tf.placeholder(tf.float32, name='critic_loss_tb')
        self.entropy_tb = tf.placeholder(tf.float32, name='entropy_tb')
        self.kl_tb = tf.placeholder(tf.float32, name='kl_tb')
        self.lr_tb = tf.placeholder(tf.float32, name='lr_tb')
        self.beta_tb = tf.placeholder(tf.float32, name='beta_tb')

        with tf.name_scope('loss'):
            tf.summary.scalar('actor_loss', self.actor_loss_tb)
            tf.summary.scalar('critic_loss', self.critic_loss_tb)

        with tf.name_scope('param'):
            tf.summary.scalar('entropy', self.entropy_tb)
            tf.summary.scalar('kl', self.kl_tb)
            tf.summary.scalar('lr', self.lr_tb)
            tf.summary.scalar('beta', self.beta_tb)

        tf.summary.scalar('score', self.score_tb)

    def update_actor(self, obs, acts, advs, rets, score):
        feed_dict = {
        self.obs_ph: obs,
        self.act_ph: acts,
        self.adv_ph: advs,
        self.ret_ph: rets,
        self.beta_ph: self.beta,
        self.lr_ph: self.args.actor_lr * self.lr_multiplier
        }

        old_means_np, old_log_vars_np = self.session.run([self.means, self.log_vars],
                                                      feed_dict)
        feed_dict[self.old_log_vars_ph] = old_log_vars_np
        feed_dict[self.old_means_ph] = old_means_np

        for e in range(self.args.train_epochs):
            self.session.run(self.actor_opt, feed_dict)
            kl = self.session.run(self.kl, feed_dict)
            if kl > self.args.kl_targ * 4: # early stopping
                break

        # magic setting
        if kl > self.args.kl_targ * 2:
            self.beta = np.minimum(35, 1.5 * self.beta)
            if self.beta > 30 and self.lr_multiplier > 0.1:
                self.lr_multiplier /= 1.5
        elif kl < self.args.kl_targ / 2.0:
            self.beta = np.maximum(1.0 / 35.0, self.beta / 1.5)
            if self.beta < (1.0 / 30.0) and self.lr_multiplier < 10:
                self.lr_multiplier *= 1.5

        stats = self._visualize_stats(feed_dict, score)
        self._visualize_tensorboard(stats)

        if self.args.save_network and self.time_step % 10 == 0.:
            self.save_network(self.args.model_name)

        return stats

    def update_critic(self, x, y):
        num_batches = max(x.shape[0] // 256, 1)
        batch_size = x.shape[0] // num_batches
        if self.replay_buffer_x is None:
            x_train, y_train = x, y
        else:
            x_train = np.concatenate([x, self.replay_buffer_x])
            y_train = np.concatenate([y, self.replay_buffer_y])
        self.replay_buffer_x = x
        self.replay_buffer_y = y
        for e in range(self.critic_epochs):
            x_train, y_train = shuffle(x_train, y_train)
            for j in range(num_batches):
                start = j * batch_size
                end = (j + 1) * batch_size
                obs = x_train[start:end, :]
                ret = y_train[start:end]
                feed_dict = {self.obs_ph: obs,
                             self.ret_ph: ret}
                self.session.run(self.critic_opt, feed_dict=feed_dict)

    def _visualize_stats(self, feed_dict, score):
        kl, entropy, actor_loss, critic_loss = self.session.run(
            [self.kl, self.entropy, self.actor_loss, self.critic_loss],
            feed_dict)

        stats = OrderedDict()
        stats["Score"] = score
        stats["LearningRate"] = self.args.actor_lr * self.lr_multiplier
        stats["Beta"] = self.beta
        stats["KL-divergence"] = kl
        stats["Entropy"] = entropy
        stats["ActorLoss"] = actor_loss
        stats["CriticLoss"] = critic_loss

        return stats

    def _visualize_tensorboard(self, stats):
        feed_dict = {
        self.score_tb: stats["Score"],
        self.lr_tb: stats["LearningRate"],
        self.beta_tb: stats["Beta"],
        self.kl_tb: stats["KL-divergence"],
        self.entropy_tb: stats["Entropy"],
        self.actor_loss_tb: stats["ActorLoss"],
        self.critic_loss_tb: stats["CriticLoss"],
        }

        self.time_step += 1
        summary = self.session.run(self.merge_all, feed_dict)
        self.writer.add_summary(summary, self.time_step)

    def sample(self, obs):
        # obs = np.reshape(obs, (1, self.obs_dim))
        feed_dict = {self.obs_ph: obs}
        if self.args.test_algorithm:
            return self.session.run(self.means, feed_dict=feed_dict)
        else:
            return self.session.run(self.sampled_act, feed_dict=feed_dict)

    def get_value(self, obs):
        values = self.value.eval(feed_dict = {self.obs_ph: obs})
        return values

    def convert_action(self, action):
        return action * self.act_high

    def init_param(self, env, args):
        self.args = args
        # env param
        self.obs_dim = env.observation_space.shape[0]
        self.act_dim = env.action_space.shape[0]
        self.act_high = env.action_space.high

        # value init
        self.time_step = 0
        self.score = 0

        # actor param
        self.beta = 1
        self.lr_multiplier = 1.0

        # critic param
        self.replay_buffer_x = None
        self.replay_buffer_y = None
        self.critic_epochs = 10

    def save_network(self, model_name):
        for i in range(len(self.model)):
            tl.files.save_npz(self.model[i].all_params,
                              name='../model/ppo/{}_{}.npz'.format(model_name, i),
                              sess=self.session)

    def load_network(self, model_name):
        for i in range(len(self.model)):
            params = tl.files.load_npz(name='../model/ppo/{}_{}.npz'.format(model_name, i))
            tl.files.assign_params(self.session, params, self.model[i])

@ray.remote
class RayEnvironment(object):
    def __init__(self, env):
        self.env = env
        state = self.env.reset()
        self.shape = state.shape

    def step(self, action):
        if self.done:
            return [np.zeros(self.shape), 0.0, True]
        else:
            state, reward, done, info = self.env.step(action)
            self.done = done
            return [state, reward, done]

    def reset(self):
        self.done = False
        return self.env.reset()

def run_episode(envs, agent, trajectories, animate=args.animate):
    terminates = [False for _ in range(len(envs))]
    terminates_idxs = [0 for _ in range(len(envs))]

    paths_obs, paths_act, paths_rew = [], [], []

    states = [env.reset.remote() for env in envs]
    states = ray.get(states)

    while not all(terminates):
        # if animate:
        #     env.render()
        paths_obs.append(states)
        actions = agent.sample(states)
        paths_act.append(actions)
        next_step = [env.step.remote(actions[i]) for i, env in enumerate(envs)]
        next_step = ray.get(next_step)

        states = [batch[0] for batch in next_step]
        rewards = [batch[1] for batch in next_step]
        dones = [batch[2] for batch in next_step]

        paths_rew.append(rewards)

        for i, d in enumerate(dones):
            if d:
                terminates[i] = True
            else:
                terminates_idxs[i] += 1

    for i in range(len(envs)):
        obs = []
        acts = []
        rews = []

        for j in range(len(paths_rew)):
            obs.append(paths_obs[j][i])
            acts.append(paths_act[j][i])
            rews.append(paths_rew[j][i])
            if terminates_idxs[i] == j:
                break

        obs = np.asarray(obs)
        acts = np.asarray(acts)
        rews = np.asarray(rews)

        acts = np.reshape(acts, (len(rews), agent.act_dim))
        trajectory = {
        'obs': obs,
        'acts': acts,
        'rewards': rews
        }

        trajectories.append(trajectory)

    return trajectories

def run_policy(envs, agent, training_steps, batch_size):
    trajectories = []

    for e in range((batch_size // args.nums_worker) + 1):
        trajectories = run_episode(envs, agent, trajectories)

    mean_step = np.mean([len(t['rewards']) for t in trajectories])
    score = np.mean([t['rewards'].sum() for t in trajectories])
    return trajectories, score, mean_step

def discount(x, gamma):
    return scipy.signal.lfilter([1.0], [1.0, -gamma], x[::-1])[::-1]

def add_disc_sum_rew(trajectories, gamma):
    for trajectory in trajectories:
        rewards = trajectory['rewards'] * (1 - gamma)
        disc_sum_rew = discount(rewards, gamma)
        trajectory['disc_sum_rew'] = disc_sum_rew

def add_value(trajectories, agent):
    for trajectory in trajectories:
        obs = trajectory['obs']
        values = agent.get_value(obs)
        trajectory['values'] = np.squeeze(np.asarray(values))

def add_gae(trajectories, gamma, lam):
    for trajectory in trajectories:
        rewards = trajectory['rewards'] * (1 - gamma)
        values = trajectory['values']
        # temporal differences
        tds = rewards - values + np.append(values[1:] * gamma, 0)
        advantages = discount(tds, gamma * lam)
        advs = np.asarray(advantages)
        trajectory['advs'] = advs

def build_train_set(trajectories):
    observes = np.concatenate([t['obs'] for t in trajectories])
    actions = np.concatenate([t['acts'] for t in trajectories])
    disc_sum_rew = np.concatenate([t['disc_sum_rew'] for t in trajectories])
    advantages = np.concatenate([t['advs'] for t in trajectories])
    # normalize advantages
    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-6)

    return observes, actions, advantages, disc_sum_rew

def print_stats(stats):
    print("*********** Iteration {} ************".format(stats["Iteration"]))
    table = []
    for k, v in stats.items():
        table.append([k, v])

    print(tabulate(table, tablefmt="grid"))

def train():
    env = gym.make(args.env_name)
    agent = PPO(env, args)
    if args.eval_algorithm:
        env = wrappers.Monitor(env, './model/{}'.format(args.model_name), force=True)
    e = 0

    envs = [gym.make(args.env_name) for _ in range(args.nums_worker)]
    envs = [RayEnvironment.remote(envs[i]) for i in range(args.nums_worker)]

    if args.load_network:
        agent.load_network(args.model_name)
    while e < (args.max_episodes):
        trajectories, score, mean_step = run_policy(envs, agent, args.training_steps, args.batch_size)
        e += len(trajectories)
        add_value(trajectories, agent)
        add_disc_sum_rew(trajectories, args.gamma)
        add_gae(trajectories, args.gamma, args.lambda_gae)
        obs, acts, advs, rets = build_train_set(trajectories)

        stats = agent.update_actor(obs, acts, advs, rets, score)
        agent.update_critic(obs, rets)
        stats["AverageStep"] = mean_step
        stats["Iteration"] = e
        print_stats(stats)

def apply_wechat():
    threads = []
    t1 = threading.Thread(target=train)
    threads.append(t1)
    t2 = threading.Thread(target=util.wechat_display)
    threads.append(t2)

    for t in threads:
        t.start()

if __name__ == "__main__":
    ray.init()
    train()
    # apply_wechat()
    # util.wechat_display()