import os
import os.path as osp
import time
from collections import deque
import pickle

from baselines.ddpg.ddpg_learner import DDPG
from baselines.ddpg.models import Actor, Critic
from baselines.ddpg.memory import Memory
from baselines.ddpg.noise import AdaptiveParamNoiseSpec, NormalActionNoise, OrnsteinUhlenbeckActionNoise
from baselines.common import set_global_seeds

from baselines import logger
import tensorflow as tf
import numpy as np

try:
    from mpi4py import MPI
except ImportError:
    MPI = None

def learn(network, env,
          seed=None,
          total_timesteps=None,
          nb_epochs=None, # with default settings, perform 1M steps total
          nb_epoch_cycles=20,
          nb_rollout_steps=100,
          reward_scale=1.0,
          render=False,
          render_eval=False,
          noise_type='adaptive-param_0.2',
          normalize_returns=False,
          normalize_observations=True,
          critic_l2_reg=1e-2,
          actor_lr=1e-4,
          critic_lr=1e-3,
          popart=False,
          gamma=0.99,
          clip_norm=None,
          nb_train_steps=50, # per epoch cycle and MPI worker,
          nb_eval_steps=100,
          batch_size=64, # per MPI worker
          tau=0.01,
          eval_env=None,
          param_noise_adaption_interval=50,
          load_path=None,
          **network_kwargs):

    set_global_seeds(seed)

    if total_timesteps is not None:
        assert nb_epochs is None
        nb_epochs = int(total_timesteps) // (nb_epoch_cycles * nb_rollout_steps)
    else:
        nb_epochs = 500

    if MPI is not None:
        rank = MPI.COMM_WORLD.Get_rank()
    else:
        rank = 0

    nb_actions = env.action_space.shape[-1]
    assert (np.abs(env.action_space.low) == env.action_space.high).all()  # we assume symmetric actions.

    memory = Memory(limit=int(1e6), action_shape=env.action_space.shape, observation_shape=env.observation_space.shape)
    critic = Critic(nb_actions, ob_shape=env.observation_space.shape, network=network, **network_kwargs)
    actor = Actor(nb_actions, ob_shape=env.observation_space.shape, network=network, **network_kwargs)

    action_noise = None
    param_noise = None
    if noise_type is not None:
        for current_noise_type in noise_type.split(','):
            current_noise_type = current_noise_type.strip()
            if current_noise_type == 'none':
                pass
            elif 'adaptive-param' in current_noise_type:
                _, stddev = current_noise_type.split('_')
                param_noise = AdaptiveParamNoiseSpec(initial_stddev=float(stddev), desired_action_stddev=float(stddev))
            elif 'normal' in current_noise_type:
                _, stddev = current_noise_type.split('_')
                action_noise = NormalActionNoise(mu=np.zeros(nb_actions), sigma=float(stddev) * np.ones(nb_actions))
            elif 'ou' in current_noise_type:
                _, stddev = current_noise_type.split('_')
                action_noise = OrnsteinUhlenbeckActionNoise(mu=np.zeros(nb_actions), sigma=float(stddev) * np.ones(nb_actions))
            else:
                raise RuntimeError('unknown noise type "{}"'.format(current_noise_type))

    max_action = env.action_space.high
    logger.info('scaling actions by {} before executing in env'.format(max_action))

    agent = DDPG(actor, critic, memory, env.observation_space.shape, env.action_space.shape,
        gamma=gamma, tau=tau, normalize_returns=normalize_returns, normalize_observations=normalize_observations,
        batch_size=batch_size, action_noise=action_noise, param_noise=param_noise, critic_l2_reg=critic_l2_reg,
        actor_lr=actor_lr, critic_lr=critic_lr, enable_popart=popart, clip_norm=clip_norm,
        reward_scale=reward_scale)
    logger.info('Using agent with the following configuration:')
    logger.info(str(agent.__dict__.items()))

    if load_path is not None:
        load_path = osp.expanduser(load_path)
        ckpt = tf.train.Checkpoint(model=agent)
        manager = tf.train.CheckpointManager(ckpt, load_path, max_to_keep=None)
        ckpt.restore(manager.latest_checkpoint)
        print("Restoring from {}".format(manager.latest_checkpoint))

    eval_episode_rewards_history = deque(maxlen=100)
    episode_rewards_history = deque(maxlen=100)

    # Prepare everything.
    agent.initialize()

    agent.reset()

    obs = env.reset()
    if eval_env is not None:
        eval_obs = eval_env.reset()
    nenvs = obs.shape[0]

    episode_reward = np.zeros(nenvs, dtype = np.float32) #vector
    episode_step = np.zeros(nenvs, dtype = int) # vector
    episodes = 0 #scalar
    t = 0 # scalar

    epoch = 0



    start_time = time.time()

    epoch_episode_rewards = []
    epoch_episode_steps = []
    epoch_actions = []
    epoch_qs = []
    epoch_episodes = 0
    for epoch in range(nb_epochs):
        for cycle in range(nb_epoch_cycles):
            # Perform rollouts.
            if nenvs > 1:
                # if simulating multiple envs in parallel, impossible to reset agent at the end of the episode in each
                # of the environments, so resetting here instead
                agent.reset()
            for t_rollout in range(nb_rollout_steps):
                # Predict next action.
                action, q, _, _ = agent.step(tf.constant(obs), apply_noise=True, compute_Q=True)
                action, q = action.numpy(), q.numpy()

                # Execute next action.
                if rank == 0 and render:
                    env.render()

                # max_action is of dimension A, whereas action is dimension (nenvs, A) - the multiplication gets broadcasted to the batch
                new_obs, r, done, info = env.step(max_action * action)  # scale for execution in env (as far as DDPG is concerned, every action is in [-1, 1])
                # note these outputs are batched from vecenv

                t += 1
                if rank == 0 and render:
                    env.render()
                episode_reward += r
                episode_step += 1

                # Book-keeping.
                epoch_actions.append(action)
                epoch_qs.append(q)
                agent.store_transition(obs, action, r, new_obs, done) #the batched data will be unrolled in memory.py's append.

                obs = new_obs

                for d in range(len(done)):
                    if done[d]:
                        # Episode done.
                        epoch_episode_rewards.append(episode_reward[d])
                        episode_rewards_history.append(episode_reward[d])
                        epoch_episode_steps.append(episode_step[d])
                        episode_reward[d] = 0.
                        episode_step[d] = 0
                        epoch_episodes += 1
                        episodes += 1
                        if nenvs == 1:
                            agent.reset()


            # Train.
            epoch_actor_losses = []
            epoch_critic_losses = []
            epoch_adaptive_distances = []
            for t_train in range(nb_train_steps):
                # Adapt param noise, if necessary.
                if memory.nb_entries >= batch_size and t_train % param_noise_adaption_interval == 0:
                    batch = agent.memory.sample(batch_size=batch_size)
                    obs0 = tf.constant(batch['obs0'])
                    distance = agent.adapt_param_noise(obs0)
                    epoch_adaptive_distances.append(distance)

                cl, al = agent.train()
                epoch_critic_losses.append(cl)
                epoch_actor_losses.append(al)
                agent.update_target_net()

            # Evaluate.
            eval_episode_rewards = []
            eval_qs = []
            if eval_env is not None:
                nenvs_eval = eval_obs.shape[0]
                eval_episode_reward = np.zeros(nenvs_eval, dtype = np.float32)
                for t_rollout in range(nb_eval_steps):
                    eval_action, eval_q, _, _ = agent.step(eval_obs, apply_noise=False, compute_Q=True)
                    eval_obs, eval_r, eval_done, eval_info = eval_env.step(max_action * eval_action)  # scale for execution in env (as far as DDPG is concerned, every action is in [-1, 1])
                    if render_eval:
                        eval_env.render()
                    eval_episode_reward += eval_r

                    eval_qs.append(eval_q)
                    for d in range(len(eval_done)):
                        if eval_done[d]:
                            eval_episode_rewards.append(eval_episode_reward[d])
                            eval_episode_rewards_history.append(eval_episode_reward[d])
                            eval_episode_reward[d] = 0.0

        if MPI is not None:
            mpi_size = MPI.COMM_WORLD.Get_size()
        else:
            mpi_size = 1

        # Log stats.
        # XXX shouldn't call np.mean on variable length lists
        duration = time.time() - start_time
        stats = agent.get_stats()
        combined_stats = stats.copy()
        combined_stats['rollout/return'] = np.mean(epoch_episode_rewards)
        combined_stats['rollout/return_std'] = np.std(epoch_episode_rewards)
        combined_stats['rollout/return_history'] = np.mean(episode_rewards_history)
        combined_stats['rollout/return_history_std'] = np.std(episode_rewards_history)
        combined_stats['rollout/episode_steps'] = np.mean(epoch_episode_steps)
        combined_stats['rollout/actions_mean'] = np.mean(epoch_actions)
        combined_stats['rollout/Q_mean'] = np.mean(epoch_qs)
        combined_stats['train/loss_actor'] = np.mean(epoch_actor_losses)
        combined_stats['train/loss_critic'] = np.mean(epoch_critic_losses)
        combined_stats['train/param_noise_distance'] = np.mean(epoch_adaptive_distances)
        combined_stats['total/duration'] = duration
        combined_stats['total/steps_per_second'] = float(t) / float(duration)
        combined_stats['total/episodes'] = episodes
        combined_stats['rollout/episodes'] = epoch_episodes
        combined_stats['rollout/actions_std'] = np.std(epoch_actions)
        # Evaluation statistics.
        if eval_env is not None:
            combined_stats['eval/return'] = eval_episode_rewards
            combined_stats['eval/return_history'] = np.mean(eval_episode_rewards_history)
            combined_stats['eval/Q'] = eval_qs
            combined_stats['eval/episodes'] = len(eval_episode_rewards)
        def as_scalar(x):
            if isinstance(x, np.ndarray):
                assert x.size == 1
                return x[0]
            elif np.isscalar(x):
                return x
            else:
                raise ValueError('expected scalar, got %s'%x)

        combined_stats_sums = np.array([ np.array(x).flatten()[0] for x in combined_stats.values()])
        if MPI is not None:
            combined_stats_sums = MPI.COMM_WORLD.allreduce(combined_stats_sums)

        combined_stats = {k : v / mpi_size for (k,v) in zip(combined_stats.keys(), combined_stats_sums)}

        # Total statistics.
        combined_stats['total/epochs'] = epoch + 1
        combined_stats['total/steps'] = t

        for key in sorted(combined_stats.keys()):
            logger.record_tabular(key, combined_stats[key])

        if rank == 0:
            logger.dump_tabular()
        logger.info('')
        logdir = logger.get_dir()
        if rank == 0 and logdir:
            if hasattr(env, 'get_state'):
                with open(os.path.join(logdir, 'env_state.pkl'), 'wb') as f:
                    pickle.dump(env.get_state(), f)
            if eval_env and hasattr(eval_env, 'get_state'):
                with open(os.path.join(logdir, 'eval_env_state.pkl'), 'wb') as f:
                    pickle.dump(eval_env.get_state(), f)


    return agent

[((49, 4, 49, 26), 'baselines.common.set_global_seeds', 'set_global_seeds', ({(49, 21, 49, 25): 'seed'}, {}), '(seed)', False, 'from baselines.common import set_global_seeds\n'), ((66, 13, 66, 104), 'baselines.ddpg.models.Critic', 'Critic', (), '', False, 'from baselines.ddpg.models import Actor, Critic\n'), ((67, 12, 67, 102), 'baselines.ddpg.models.Actor', 'Actor', (), '', False, 'from baselines.ddpg.models import Actor, Critic\n'), ((91, 12, 95, 34), 'baselines.ddpg.ddpg_learner.DDPG', 'DDPG', (), '', False, 'from baselines.ddpg.ddpg_learner import DDPG\n'), ((96, 4, 96, 64), 'baselines.logger.info', 'logger.info', ({(96, 16, 96, 63): '"""Using agent with the following configuration:"""'}, {}), "('Using agent with the following configuration:')", False, 'from baselines import logger\n'), ((106, 35, 106, 52), 'collections.deque', 'deque', (), '', False, 'from collections import deque\n'), ((107, 30, 107, 47), 'collections.deque', 'deque', (), '', False, 'from collections import deque\n'), ((119, 21, 119, 56), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((120, 19, 120, 47), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((128, 17, 128, 28), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((58, 15, 58, 40), 'mpi4py.MPI.COMM_WORLD.Get_rank', 'MPI.COMM_WORLD.Get_rank', ({}, {}), '()', False, 'from mpi4py import MPI\n'), ((100, 20, 100, 45), 'os.path.expanduser', 'osp.expanduser', ({(100, 35, 100, 44): 'load_path'}, {}), '(load_path)', True, 'import os.path as osp\n'), ((101, 15, 101, 47), 'tensorflow.train.Checkpoint', 'tf.train.Checkpoint', (), '', True, 'import tensorflow as tf\n'), ((102, 18, 102, 79), 'tensorflow.train.CheckpointManager', 'tf.train.CheckpointManager', (), '', True, 'import tensorflow as tf\n'), ((229, 43, 229, 73), 'numpy.mean', 'np.mean', ({(229, 51, 229, 72): 'epoch_episode_rewards'}, {}), '(epoch_episode_rewards)', True, 'import numpy as np\n'), ((230, 47, 230, 76), 'numpy.std', 'np.std', ({(230, 54, 230, 75): 'epoch_episode_rewards'}, {}), '(epoch_episode_rewards)', True, 'import numpy as np\n'), ((231, 51, 231, 83), 'numpy.mean', 'np.mean', ({(231, 59, 231, 82): 'episode_rewards_history'}, {}), '(episode_rewards_history)', True, 'import numpy as np\n'), ((232, 55, 232, 86), 'numpy.std', 'np.std', ({(232, 62, 232, 85): 'episode_rewards_history'}, {}), '(episode_rewards_history)', True, 'import numpy as np\n'), ((233, 50, 233, 78), 'numpy.mean', 'np.mean', ({(233, 58, 233, 77): 'epoch_episode_steps'}, {}), '(epoch_episode_steps)', True, 'import numpy as np\n'), ((234, 49, 234, 71), 'numpy.mean', 'np.mean', ({(234, 57, 234, 70): 'epoch_actions'}, {}), '(epoch_actions)', True, 'import numpy as np\n'), ((235, 43, 235, 60), 'numpy.mean', 'np.mean', ({(235, 51, 235, 59): 'epoch_qs'}, {}), '(epoch_qs)', True, 'import numpy as np\n'), ((236, 45, 236, 72), 'numpy.mean', 'np.mean', ({(236, 53, 236, 71): 'epoch_actor_losses'}, {}), '(epoch_actor_losses)', True, 'import numpy as np\n'), ((237, 46, 237, 74), 'numpy.mean', 'np.mean', ({(237, 54, 237, 73): 'epoch_critic_losses'}, {}), '(epoch_critic_losses)', True, 'import numpy as np\n'), ((238, 55, 238, 88), 'numpy.mean', 'np.mean', ({(238, 63, 238, 87): 'epoch_adaptive_distances'}, {}), '(epoch_adaptive_distances)', True, 'import numpy as np\n'), ((243, 48, 243, 69), 'numpy.std', 'np.std', ({(243, 55, 243, 68): 'epoch_actions'}, {}), '(epoch_actions)', True, 'import numpy as np\n'), ((274, 8, 274, 23), 'baselines.logger.info', 'logger.info', ({(274, 20, 274, 22): '""""""'}, {}), "('')", False, 'from baselines import logger\n'), ((275, 17, 275, 33), 'baselines.logger.get_dir', 'logger.get_dir', ({}, {}), '()', False, 'from baselines import logger\n'), ((220, 23, 220, 48), 'mpi4py.MPI.COMM_WORLD.Get_size', 'MPI.COMM_WORLD.Get_size', ({}, {}), '()', False, 'from mpi4py import MPI\n'), ((226, 19, 226, 30), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((247, 52, 247, 89), 'numpy.mean', 'np.mean', ({(247, 60, 247, 88): 'eval_episode_rewards_history'}, {}), '(eval_episode_rewards_history)', True, 'import numpy as np\n'), ((261, 34, 261, 79), 'mpi4py.MPI.COMM_WORLD.allreduce', 'MPI.COMM_WORLD.allreduce', ({(261, 59, 261, 78): 'combined_stats_sums'}, {}), '(combined_stats_sums)', False, 'from mpi4py import MPI\n'), ((270, 12, 270, 59), 'baselines.logger.record_tabular', 'logger.record_tabular', ({(270, 34, 270, 37): 'key', (270, 39, 270, 58): 'combined_stats[key]'}, {}), '(key, combined_stats[key])', False, 'from baselines import logger\n'), ((273, 12, 273, 33), 'baselines.logger.dump_tabular', 'logger.dump_tabular', ({}, {}), '()', False, 'from baselines import logger\n'), ((63, 12, 63, 40), 'numpy.abs', 'np.abs', ({(63, 19, 63, 39): 'env.action_space.low'}, {}), '(env.action_space.low)', True, 'import numpy as np\n'), ((204, 38, 204, 78), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((254, 17, 254, 31), 'numpy.isscalar', 'np.isscalar', ({(254, 29, 254, 30): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((144, 45, 144, 61), 'tensorflow.constant', 'tf.constant', ({(144, 57, 144, 60): 'obs'}, {}), '(obs)', True, 'import tensorflow as tf\n'), ((190, 27, 190, 53), 'tensorflow.constant', 'tf.constant', ({(190, 39, 190, 52): "batch['obs0']"}, {}), "(batch['obs0'])", True, 'import tensorflow as tf\n'), ((278, 26, 278, 63), 'os.path.join', 'os.path.join', ({(278, 39, 278, 45): 'logdir', (278, 47, 278, 62): '"""env_state.pkl"""'}, {}), "(logdir, 'env_state.pkl')", False, 'import os\n'), ((281, 26, 281, 68), 'os.path.join', 'os.path.join', ({(281, 39, 281, 45): 'logdir', (281, 47, 281, 67): '"""eval_env_state.pkl"""'}, {}), "(logdir, 'eval_env_state.pkl')", False, 'import os\n'), ((259, 41, 259, 52), 'numpy.array', 'np.array', ({(259, 50, 259, 51): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((81, 52, 81, 72), 'numpy.zeros', 'np.zeros', ({(81, 61, 81, 71): 'nb_actions'}, {}), '(nb_actions)', True, 'import numpy as np\n'), ((81, 96, 81, 115), 'numpy.ones', 'np.ones', ({(81, 104, 81, 114): 'nb_actions'}, {}), '(nb_actions)', True, 'import numpy as np\n'), ((84, 63, 84, 83), 'numpy.zeros', 'np.zeros', ({(84, 72, 84, 82): 'nb_actions'}, {}), '(nb_actions)', True, 'import numpy as np\n'), ((84, 107, 84, 126), 'numpy.ones', 'np.ones', ({(84, 115, 84, 125): 'nb_actions'}, {}), '(nb_actions)', True, 'import numpy as np\n')]