#!/usr/bin/env python3
import os
from baselines.common.cmd_util import make_mujoco_env, mujoco_arg_parser
from baselines.common import tf_util as U
from baselines import logger
from imitation_env_2_vel_symmetry_8ms import robotisImitationEnv
import gym
import numpy as np

def train(num_timesteps, seed, env, model_path=None):
    env_id = 'walkman-v1'
    from baselines.ppo1 import mlp_policy, pposgd_simple
    U.make_session(num_cpu=1).__enter__()

    def policy_fn(name, ob_space, ac_space):
        return mlp_policy.MlpPolicy(name=name, ob_space=ob_space, ac_space=ac_space,
                                    hid_size=64, num_hid_layers=2)

    logger.log("NOTE: reward will be scaled by a factor of 10  in logged stats. Check the monitor for unscaled reward.")
    pi, oldpi = pposgd_simple.learn(env, policy_fn,
                             max_timesteps=num_timesteps,
                             timesteps_per_actorbatch=2048,
                             clip_param=0.1, entcoeff=0.0,
                             optim_epochs=10,
                             optim_stepsize=1e-4,
                             optim_batchsize=128,
                             gamma=0.99,
                             lam=0.95,
                             schedule='constant',
                             )
    # env.close()
    if model_path:
        U.save_state(model_path)

    return pi

def retrain(num_timesteps, seed, env, model_path=None):
    env_id = 'walkman-v1'
    from baselines.ppo1 import mlp_policy, pposgd_simple
    U.make_session(num_cpu=1).__enter__()

    def policy_fn(name, ob_space, ac_space):
        return mlp_policy.MlpPolicy(name=name, ob_space=ob_space, ac_space=ac_space,
                                    hid_size=64, num_hid_layers=2)

    # env = make_mujoco_env(env_id, seed)

    # parameters below were the best found in a simple random search
    # these are good enough to make humanoid walk, but whether those are
    # an absolute best or not is not certain
    # env = RewScale(env, 0.1)
    # env = gym.make(env_id)
    logger.log("NOTE: reward will be scaled by a factor of 10  in logged stats. Check the monitor for unscaled reward.")
    pi, old_pi = pposgd_simple.learn(env, policy_fn,
                             max_timesteps=num_timesteps,
                             timesteps_per_actorbatch=8,
                             clip_param=0.1, entcoeff=0.0,
                             optim_epochs=10,
                             optim_stepsize=1e-4,
                             optim_batchsize=8,
                             gamma=0.99,
                             lam=0.95,
                             schedule='constant',
                             )
    # env.close()
    if model_path:
        U.save_state(model_path)

    return pi, old_pi

def contrain(num_timesteps, pi_in, oldpi_in, seed, env, model_path=None):
    from baselines.ppo1 import mlp_policy, pposgd_simple
    # U.make_session(num_cpu=1).__enter__()

    def policy_fn(name, ob_space, ac_space):
        return mlp_policy.MlpPolicy(name=name, ob_space=ob_space, ac_space=ac_space,
                                    hid_size=64, num_hid_layers=2)

    logger.log("NOTE: reward will be scaled by a factor of 10  in logged stats. Check the monitor for unscaled reward.")
    pi = pposgd_simple.learn2(env, policy_fn,
                             max_timesteps=num_timesteps,
                             timesteps_per_actorbatch=2048,
                             clip_param=0.1, entcoeff=0.0,
                             optim_epochs=10,
                             optim_stepsize=1e-4,
                             optim_batchsize=128,
                             gamma=0.99,
                             lam=0.95,
                             schedule='constant',pi_in=pi_in, oldpi_in=oldpi_in
                             )
    # env.close()
    if model_path:
        U.save_state(model_path)

    return pi

class RewScale(gym.RewardWrapper):
    def __init__(self, env, scale):
        gym.RewardWrapper.__init__(self, env)
        self.scale = scale

    def reward(self, r):
        return r * self.scale

def main():
    logger.configure()
    parser = mujoco_arg_parser()
    parser.add_argument('--model-path', default=os.path.join(logger.get_dir(), 'humanoid_policy'))
    parser.set_defaults(num_timesteps=int(4e7))
    env_id = 'walkman-v1'

    args = parser.parse_args()

    con = False
    args.play = False
    if not args.play:
        # train the model
        # env = gym.make(env_id)
        env = robotisImitationEnv()
        train(num_timesteps=args.num_timesteps, seed=args.seed, env=env, model_path=args.model_path)
        env.close()
    else:
        if con:
            env = robotisImitationEnv()
            pi, oldpi = retrain(num_timesteps=1, env=env, seed=args.seed)
            args.model_path = "models/8500"
            U.load_state(args.model_path)
            contrain(num_timesteps=args.num_timesteps, seed=args.seed, env=env, model_path=args.model_path, pi_in=pi, oldpi_in=oldpi)
            env.close()

        else:
            # construct the model object, load pre-trained model and render
            print('Play')
            data = open('13000.txt','a')
            env = robotisImitationEnv()
            pi, oldpi = retrain(num_timesteps=1, env=env, seed=args.seed)
            
            args.model_path = "models/13000"
            U.load_state(args.model_path)
            # env = make_mujoco_env('Humanoid-v2', seed=0)
            print('reset',file = data)
            ob = env.reset()
            for i in range(34):
                print(ob[i],file = data,end=' ')
            print(ob[34],file = data)
            count = 0
            while True:
                
                #print('go to cycle',file = data)
                action = pi.act(stochastic=False, ob=ob)[0]
                #noise = np.random.randn(action.shape[0])
                #action = action+ 0.05*noise
                ob, _, done, _ = env.step(action)
                for i in range(34):
                    print(ob[i],file = data,end=' ')
                print(ob[34],file = data)

                # env.render()
                count+=1
                print(count)
                if done:
                    data.close()
                    ob = env.reset()
            
            print('next2')
            count=0
            args.model_path = "models/2900_1"
            U.load_state(args.model_path)

            while ob[2]<8.5 or env.phase!=0:
                #print('go to cycle',file = data)
                action = pi.act(stochastic=False, ob=ob)[0]
                #noise = np.random.randn(action.shape[0])
                #action = action+ 0.05*noise
                ob, _, done, _ = env.step(action)
                for i in range(34):
                    print(ob[i],file = data,end=' ')
                print(ob[34],file = data)
                # env.render()
                count+=1
                if done:
                    data.close()
                    ob = env.reset()
            print("next3")
            count=0
            args.model_path = "models/6500_4"
            U.load_state(args.model_path)
            while ob[2]<8.0 or env.phase==0:
                #print('go to cycle',file = data)
                action = pi.act(stochastic=False, ob=ob)[0]
                #noise = np.random.randn(action.shape[0])
                #action = action+ 0.05*noise
                ob, _, done, _ = env.step(action)
                for i in range(34):
                    print(ob[i],file = data,end=' ')
                print(ob[34],file = data)
                # env.render()
                count+=1
                if done:
                    data.close()
                    ob = env.reset()
             
            print('next4')
            count=0
            args.model_path = "models/6500_4"
            U.load_state(args.model_path)

            while ob[2]<12:
                #print('go to cycle',file = data)
                action = pi.act(stochastic=False, ob=ob)[0]
                #noise = np.random.randn(action.shape[0])
                #action = action+ 0.05*noise
                ob, _, done, _ = env.step(action)
                for i in range(34):
                    print(ob[i],file = data,end=' ')
                print(ob[34],file = data)
                # env.render()
                count+=1
                if done:
                    data.close()
                    ob = env.reset()
                    
             

if __name__ == '__main__':
    main()