import tensorflow as tf
import numpy as np
from baselines.common.vec_env.dummy_vec_env import DummyVecEnv

N_TRIALS = 10000
N_EPISODES = 100

_sess_config = tf.compat.v1.ConfigProto(
    allow_soft_placement=True,
    intra_op_parallelism_threads=1,
    inter_op_parallelism_threads=1
)

def simple_test(env_fn, learn_fn, min_reward_fraction, n_trials=N_TRIALS):
    def seeded_env_fn():
        env = env_fn()
        env.seed(0)
        return env

    np.random.seed(0)
    env = DummyVecEnv([seeded_env_fn])
    with tf.Graph().as_default(), tf.compat.v1.Session(config=_sess_config).as_default():
        tf.compat.v1.set_random_seed(0)
        model = learn_fn(env)
        sum_rew = 0
        done = True
        for i in range(n_trials):
            if done:
                obs = env.reset()
                state = model.initial_state
            if state is not None:
                a, v, state, _ = model.step(obs, S=state, M=[False])
            else:
                a, v, _, _ = model.step(obs)
            obs, rew, done, _ = env.step(a)
            sum_rew += float(rew)
        print("Reward in {} trials is {}".format(n_trials, sum_rew))
        assert sum_rew > min_reward_fraction * n_trials, \
            'sum of rewards {} is less than {} of the total number of trials {}'.format(sum_rew, min_reward_fraction, n_trials)

def reward_per_episode_test(env_fn, learn_fn, min_avg_reward, n_trials=N_EPISODES):
    env = DummyVecEnv([env_fn])
    with tf.Graph().as_default(), tf.compat.v1.Session(config=_sess_config).as_default():
        model = learn_fn(env)
        N_TRIALS = 100
        observations, actions, rewards = rollout(env, model, N_TRIALS)
        rewards = [sum(r) for r in rewards]
        avg_rew = sum(rewards) / N_TRIALS
        print("Average reward in {} episodes is {}".format(n_trials, avg_rew))
        assert avg_rew > min_avg_reward, \
            'average reward in {} episodes ({}) is less than {}'.format(n_trials, avg_rew, min_avg_reward)

def rollout(env, model, n_trials):
    rewards = []
    actions = []
    observations = []
    for i in range(n_trials):
        obs = env.reset()
        state = model.initial_state if hasattr(model, 'initial_state') else None
        episode_rew = []
        episode_actions = []
        episode_obs = []
        while True:
            if state is not None:
                a, v, state, _ = model.step(obs, S=state, M=[False])
            else:
                a,v, _, _ = model.step(obs)

            obs, rew, done, _ = env.step(a)
            episode_rew.append(rew)
            episode_actions.append(a)
            episode_obs.append(obs)
            if done:
                break
        rewards.append(episode_rew)
        actions.append(episode_actions)
        observations.append(episode_obs)
    return observations, actions, rewards


def smoketest(argstr, **kwargs):
    import tempfile
    import subprocess
    import os
    argstr = 'python -m baselines.run ' + argstr
    for key, value in kwargs:
        argstr += ' --{}={}'.format(key, value)
    tempdir = tempfile.mkdtemp()
    env = os.environ.copy()
    env['OPENAI_LOGDIR'] = tempdir
    subprocess.run(argstr.split(' '), env=env)
    return tempdir