<reponame>idoyehe/PCC-RL
# Copyright 2019 <NAME> and <NAME>
#
# 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.

import gym
import network_sim
import tensorflow as tf

from stable_baselines.common.policies import FeedForwardPolicy
from stable_baselines import PPO1
import os
from shutil import rmtree
import sys
import inspect

currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
parentdir = os.path.dirname(currentdir)
sys.path.insert(0, parentdir)
from common.simple_arg_parse import arg_or_default

# arch_str = arg_or_default("--arch", default="32,16")
arch_str = arg_or_default("--arch", default="64,32")
if arch_str == "":
    arch = []
else:
    arch = [int(layer_width) for layer_width in arch_str.split(",")]
print("Architecture is: %s" % str(arch))

training_sess = None

env = gym.make('PccNs-v0')
# env = gym.make('CartPole-v0')

gamma = arg_or_default("--gamma", default=0.99)
print("gamma = %f" % gamma)


class MyMlpPolicy(FeedForwardPolicy):

    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False, **_kwargs):
        super(MyMlpPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse, net_arch=[{"pi": arch, "vf": arch}],
                                          feature_extraction="mlp", **_kwargs)
        global training_sess
        training_sess = sess


model = PPO1(MyMlpPolicy, env, verbose=1, schedule='constant', timesteps_per_actorbatch=8192, optim_batchsize=2048, gamma=gamma)

# for i in range(0, 6):
# with model.graph.as_default():
#     saver = tf.compat.v1.train.Saver()
#     saver.save(training_sess, "./pcc_model_%d.ckpt" % i)
model.learn(total_timesteps=(100 * 8192))

env.testing(True)
obs = env.reset()
for _ in range(10 * 8192):
    action, _states = model.predict(obs)
    obs, reward, done, info = env.step(action)

env.reset()
env.testing(False)

#
# # Save the model to the location specified below.
#
# default_export_dir = parentdir + "/tmp/pcc_saved_models/model_A/"
#
# rmtree(parentdir + '/tmp', ignore_errors=True)
#
# export_dir = arg_or_default("--model-dir", default=default_export_dir)
# with model.graph.as_default():
#     pol = model.policy_pi  # act_model
#
#     obs_ph = pol.obs_ph
#     act = pol.deterministic_action
#     sampled_act = pol.action
#
#     obs_input = tf.compat.v1.saved_model.utils.build_tensor_info(obs_ph)
#     outputs_tensor_info = tf.saved_model.utils.build_tensor_info(act)
#     stochastic_act_tensor_info = tf.saved_model.utils.build_tensor_info(sampled_act)
#     signature = tf.compat.v1.saved_model.signature_def_utils.build_signature_def(
#         inputs={"ob": obs_input},
#         outputs={"act": outputs_tensor_info, "stochastic_act": stochastic_act_tensor_info},
#         method_name=tf.saved_model.PREDICT_METHOD_NAME)
#
#     signature_map = {tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
#                          signature}
#
#     model_builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(export_dir)
#     model_builder.add_meta_graph_and_variables(model.sess,
#                                                tags=[tf.saved_model.SERVING],
#                                                signature_def_map=signature_map,
#                                                clear_devices=True)
#     model_builder.save(as_text=True)
