# Copyright 2017 reinforce.io. All Rights Reserved.
#
# 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.
# ==============================================================================

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import unittest

import numpy as np
from six.moves import xrange

from tensorforce import Configuration
from tensorforce.agents import TRPOAgent
from tensorforce.core.networks import layered_network_builder
from tensorforce.execution import Runner
from tensorforce.contrib.openai_gym import OpenAIGym


class TestQuickstartExample(unittest.TestCase):

    def test_example(self):
        passed = 0

        for _ in xrange(3):
            # Create an OpenAIgym environment
            env = OpenAIGym('CartPole-v0')

            # Create a Trust Region Policy Optimization agent
            agent = TRPOAgent(config=Configuration(
                log_level='info',
                batch_size=100,
                baseline=dict(
                    type='mlp',
                    size=32,
                    hidden_layers=1,
                    epochs=20,
                    update_batch_size=32
                ),
                gae_rewards=True,
                normalize_rewards=False,
                gae_lambda=0.97,
                max_kl_divergence=0.005,
                cg_iterations=20,
                cg_damping=0.01,
                ls_max_backtracks=20,
                ls_override=False,
                states=env.states,
                actions=env.actions,
                network=layered_network_builder([
                    dict(type='dense', size=32, activation='tanh'),
                    dict(type='dense', size=32, activation='tanh')
                ])
            ))
            runner = Runner(agent=agent, environment=env)

            def episode_finished(r):
                # Test if mean reward over 50 should ensure that learning took off
                avg_reward = np.mean(r.episode_rewards[-50:])
                return r.episode < 100 or avg_reward < 50.0

            runner.run(episodes=2000, max_timesteps=200, episode_finished=episode_finished)

            if runner.episode < 2000:
                passed += 1

        print('Quick start example passed = {}'.format(passed))
        self.assertTrue(passed >= 2)

@@ -918,18 +918,17 @@
  import 
-TR
+P
 POAgent%0A
@@ -1386,10 +1386,9 @@
 t = 
-TR
+P
 POAg
@@ -1478,12 +1478,14 @@
 ize=
-100,
+4096,%0A
 %0A   
@@ -1501,125 +1501,100 @@
     
-baseline=dict(%0A                    type='mlp',%0A                    size=32,%0A                    hidden_layers=1,%0A    
+gae_lambda=0.97,%0A                learning_rate=0.001,%0A                entropy_penalty=0.01,%0A
     
@@ -1612,18 +1612,17 @@
  epochs=
-20
+5
 ,%0A      
@@ -1635,18 +1635,17 @@
     
- 
-   update
+optimizer
 _bat
@@ -1656,317 +1656,46 @@
 ize=
-32%0A                ),%0A                gae_rewards=True,%0A                normalize_rewards=False,%0A                gae_lambda=0.97,%0A                max_kl_divergence=0.005,%0A                cg_iterations=20,%0A                cg_damping=0.01,%0A                ls_max_backtracks=20,%0A                ls_override=False
+512,%0A                loss_clipping=0.2
 ,%0A  

Change quickstart test to PPO