from ddpg import *

import opensim as osim

from osim.http.client import Client

from osim.env import *

ENV_NAME = 'learning_to_run'

PATH = 'models/'

EPISODES = 100000

TEST = 1

def main():

    env = ei(True,seed=0,diff=0)

    env.reset()

    agent = DDPG(env)

    returns = []

    rewards = []

    rs = RunningStats()

    for episode in xrange(EPISODES):

        state = env.reset()

        reward_episode = []

        print "episode:",episode

        # Train

        demo = 50

        n_step = 3

        s,s1 = [],[]

        ea = engineered_action(np.random.rand())

        if np.random.rand() < 0.5:

            for i in range(demo):

                ob = env.step(ea)[0]

        ob = env.step(ea)[0]

        s = ob

        ob = env.step(ea)[0]

        s1 = ob

        s = process_state(s,s1,center=True) #s, which stands for state, is the new ob

        rs.normalize(s)

        for step in xrange(1000):

            ac = agent.action(s)

            print(ac)

            ac = np.clip(ac + agent.exploration_noise.noise(),0.05,0.95)

            temp = 0

            for i in range(n_step):

                ob, rew, new, _ = env.step(ac+agent.exploration_noise.noise()*0.2,0.05,0.95))

                rew = (rew/0.01 + int(new) * 0.1 + int((ob[2]/0.70)<1.0) * -1.)

                temp += rew

                if new: 

                    break

                s1 = ob

            rew = temp

            print(rew)

            s1 = process_state(s1,ob,center=True)

            rs.normalize(s1)

            agent.perceive(s,ac,rew,s1,new)

            s = s1

            s1 = ob

            reward_episode.append(rew)

            if new:

                break

        if episode % 5 == 0:

            print("episode reward = %.2f" % sum(reward_episode))

        # Testing:

        #if episode % 1 == 0:

        if episode % 100 == 0 and episode > 50:

            agent.save_model(PATH, episode)

            total_return = 0

            ave_reward = 0

            for i in xrange(TEST):

                state = env.reset()

                reward_per_step = 0

                for i in range(demo):

                    ob = env.step(ea)[0]

                ob = env.step(ea)[0]

                s = ob

                ob = env.step(ea)[0]

                s1 = ob

                s = process_state(s,s1,center=True) #s, which stands for state, is the new ob

                rs.normalize(s)

                for j in xrange(1000):

                    ac = agent.action(s)

                    temp = 0

                    for i in range(n_step):

                        ob, rew, new, _ = env.step(ac)

                        rew = (rew/0.01 + int(new) * 0.1 + int((ob[2]/0.70)<1.0) * -1.)

                        temp += rew

                        if new: 

                            break

                        s1 = ob

                    rew = temp

                    s1 = process_state(s1,ob,center=True)

                    rs.normalize(s1)

                    s = s1

                    s1 = ob

                    total_return += rew

                    if new:

                        break

                    reward_per_step += (rew - reward_per_step)/(j+1)

                ave_reward += reward_per_step

            ave_return = total_return/TEST

            ave_reward = ave_reward/TEST

            returns.append(ave_return)

            rewards.append(ave_reward)

            print 'episode: ',episode,'Evaluation Average Return:',ave_return, '  Evaluation Average Reward: ', ave_reward

if __name__ == '__main__':

    main()