#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import logging
import os
import pprint
import unittest

import numpy as np

# pyre-fixme[21]: Could not find module `pytest`.
import pytest
import torch
from parameterized import parameterized
from reagent.core.types import RewardOptions
from reagent.gym.agents.agent import Agent
from reagent.gym.agents.post_step import train_with_replay_buffer_post_step
from reagent.gym.envs.union import Env__Union
from reagent.gym.runners.gymrunner import evaluate_for_n_episodes, run_episode
from reagent.gym.utils import build_normalizer, fill_replay_buffer
from reagent.model_managers.model_manager import ModelManager
from reagent.model_managers.union import ModelManager__Union
from reagent.replay_memory.circular_replay_buffer import ReplayBuffer
from reagent.tensorboardX import summary_writer_context
from reagent.test.base.horizon_test_base import HorizonTestBase
from torch.utils.tensorboard import SummaryWriter


try:
    # Use internal runner or OSS otherwise
    from reagent.runners.fb.fb_batch_runner import FbBatchRunner as BatchRunner
except ImportError:
    from reagent.runners.oss_batch_runner import OssBatchRunner as BatchRunner


# for seeding the environment
SEED = 0
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


"""
Put on-policy gym tests here in the format (test name, path to yaml config).
Format path to be: "configs/<env_name>/<model_name>_<env_name>_online.yaml."
NOTE: These tests should ideally finish quickly (within 10 minutes) since they are
unit tests which are run many times.
"""
GYM_TESTS = [
    ("Discrete DQN Cartpole", "configs/cartpole/discrete_dqn_cartpole_online.yaml"),
    ("Discrete C51 Cartpole", "configs/cartpole/discrete_c51_cartpole_online.yaml"),
    ("Discrete QR Cartpole", "configs/cartpole/discrete_qr_cartpole_online.yaml"),
    (
        "Discrete DQN Open Gridworld",
        "configs/open_gridworld/discrete_dqn_open_gridworld.yaml",
    ),
    ("SAC Pendulum", "configs/pendulum/sac_pendulum_online.yaml"),
    ("TD3 Pendulum", "configs/pendulum/td3_pendulum_online.yaml"),
    ("Parametric DQN Cartpole", "configs/cartpole/parametric_dqn_cartpole_online.yaml"),
    (
        "Parametric SARSA Cartpole",
        "configs/cartpole/parametric_sarsa_cartpole_online.yaml",
    ),
    (
        "Sparse DQN Changing Arms",
        "configs/sparse/discrete_dqn_changing_arms_online.yaml",
    ),
    ("SlateQ RecSim", "configs/recsim/slate_q_recsim_online.yaml"),
    ("PossibleActionsMask DQN", "configs/functionality/dqn_possible_actions_mask.yaml"),
]


curr_dir = os.path.dirname(__file__)


class TestGym(HorizonTestBase):
    # pyre-fixme[16]: Module `parameterized` has no attribute `expand`.
    @parameterized.expand(GYM_TESTS)
    def test_gym_cpu(self, name: str, config_path: str):
        logger.info(f"Starting {name} on CPU")
        self.run_from_config(
            run_test=run_test,
            config_path=os.path.join(curr_dir, config_path),
            use_gpu=False,
        )
        logger.info(f"{name} passes!")

    # pyre-fixme[16]: Module `parameterized` has no attribute `expand`.
    @parameterized.expand(GYM_TESTS)
    @pytest.mark.serial
    # pyre-fixme[56]: Argument `not torch.cuda.is_available()` to decorator factory
    #  `unittest.skipIf` could not be resolved in a global scope.
    @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
    def test_gym_gpu(self, name: str, config_path: str):
        logger.info(f"Starting {name} on GPU")
        self.run_from_config(
            run_test=run_test,
            config_path=os.path.join(curr_dir, config_path),
            use_gpu=True,
        )
        logger.info(f"{name} passes!")


def run_test(
    env: Env__Union,
    model: ModelManager__Union,
    replay_memory_size: int,
    train_every_ts: int,
    train_after_ts: int,
    num_train_episodes: int,
    passing_score_bar: float,
    num_eval_episodes: int,
    use_gpu: bool,
):
    env = env.value
    env.seed(SEED)
    env.action_space.seed(SEED)
    normalization = build_normalizer(env)
    logger.info(f"Normalization is: \n{pprint.pformat(normalization)}")

    manager: ModelManager = model.value
    runner = BatchRunner(use_gpu, manager, RewardOptions(), normalization)
    trainer = runner.initialize_trainer()
    reporter = manager.get_reporter()
    trainer.reporter = reporter
    training_policy = manager.create_policy(trainer)

    replay_buffer = ReplayBuffer(
        replay_capacity=replay_memory_size, batch_size=trainer.minibatch_size
    )

    device = torch.device("cuda") if use_gpu else torch.device("cpu")
    # first fill the replay buffer to burn_in
    train_after_ts = max(train_after_ts, trainer.minibatch_size)
    fill_replay_buffer(
        env=env, replay_buffer=replay_buffer, desired_size=train_after_ts
    )

    post_step = train_with_replay_buffer_post_step(
        replay_buffer=replay_buffer,
        env=env,
        trainer=trainer,
        training_freq=train_every_ts,
        batch_size=trainer.minibatch_size,
        device=device,
    )

    agent = Agent.create_for_env(
        env, policy=training_policy, post_transition_callback=post_step, device=device
    )

    writer = SummaryWriter()
    with summary_writer_context(writer):
        train_rewards = []
        for i in range(num_train_episodes):
            trajectory = run_episode(
                env=env, agent=agent, mdp_id=i, max_steps=env.max_steps
            )
            ep_reward = trajectory.calculate_cumulative_reward()
            train_rewards.append(ep_reward)
            logger.info(
                f"Finished training episode {i} (len {len(trajectory)})"
                f" with reward {ep_reward}."
            )

    logger.info("============Train rewards=============")
    logger.info(train_rewards)
    logger.info(f"average: {np.mean(train_rewards)};\tmax: {np.max(train_rewards)}")

    # Check whether the max score passed the score bar; we explore during training
    # the return could be bad (leading to flakiness in C51 and QRDQN).
    assert np.max(train_rewards) >= passing_score_bar, (
        f"max reward ({np.max(train_rewards)})after training for "
        f"{len(train_rewards)} episodes is less than < {passing_score_bar}.\n"
    )

    serving_policy = manager.create_serving_policy(normalization, trainer)
    agent = Agent.create_for_env_with_serving_policy(env, serving_policy)

    eval_rewards = evaluate_for_n_episodes(
        n=num_eval_episodes, env=env, agent=agent, max_steps=env.max_steps
    ).squeeze(1)

    logger.info("============Eval rewards==============")
    logger.info(eval_rewards)
    mean_eval = np.mean(eval_rewards)
    logger.info(f"average: {mean_eval};\tmax: {np.max(eval_rewards)}")
    assert (
        mean_eval >= passing_score_bar
    ), f"Eval reward is {mean_eval}, less than < {passing_score_bar}.\n"


if __name__ == "__main__":
    unittest.main()

[((1286, 1313), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (1303, 1313), False, 'import logging\n'), ((2639, 2664), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (2654, 2664), False, 'import os\n'), ((2776, 2807), 'parameterized.parameterized.expand', 'parameterized.expand', (['GYM_TESTS'], {}), '(GYM_TESTS)\n', (2796, 2807), False, 'from parameterized import parameterized\n'), ((3188, 3219), 'parameterized.parameterized.expand', 'parameterized.expand', (['GYM_TESTS'], {}), '(GYM_TESTS)\n', (3208, 3219), False, 'from parameterized import parameterized\n'), ((4117, 4138), 'reagent.gym.utils.build_normalizer', 'build_normalizer', (['env'], {}), '(env)\n', (4133, 4138), False, 'from reagent.gym.utils import build_normalizer, fill_replay_buffer\n'), ((4513, 4601), 'reagent.replay_memory.circular_replay_buffer.ReplayBuffer', 'ReplayBuffer', ([], {'replay_capacity': 'replay_memory_size', 'batch_size': 'trainer.minibatch_size'}), '(replay_capacity=replay_memory_size, batch_size=trainer.\n    minibatch_size)\n', (4525, 4601), False, 'from reagent.replay_memory.circular_replay_buffer import ReplayBuffer\n'), ((4797, 4887), 'reagent.gym.utils.fill_replay_buffer', 'fill_replay_buffer', ([], {'env': 'env', 'replay_buffer': 'replay_buffer', 'desired_size': 'train_after_ts'}), '(env=env, replay_buffer=replay_buffer, desired_size=\n    train_after_ts)\n', (4815, 4887), False, 'from reagent.gym.utils import build_normalizer, fill_replay_buffer\n'), ((4914, 5092), 'reagent.gym.agents.post_step.train_with_replay_buffer_post_step', 'train_with_replay_buffer_post_step', ([], {'replay_buffer': 'replay_buffer', 'env': 'env', 'trainer': 'trainer', 'training_freq': 'train_every_ts', 'batch_size': 'trainer.minibatch_size', 'device': 'device'}), '(replay_buffer=replay_buffer, env=env,\n    trainer=trainer, training_freq=train_every_ts, batch_size=trainer.\n    minibatch_size, device=device)\n', (4948, 5092), False, 'from reagent.gym.agents.post_step import train_with_replay_buffer_post_step\n'), ((5152, 5257), 'reagent.gym.agents.agent.Agent.create_for_env', 'Agent.create_for_env', (['env'], {'policy': 'training_policy', 'post_transition_callback': 'post_step', 'device': 'device'}), '(env, policy=training_policy, post_transition_callback=\n    post_step, device=device)\n', (5172, 5257), False, 'from reagent.gym.agents.agent import Agent\n'), ((5281, 5296), 'torch.utils.tensorboard.SummaryWriter', 'SummaryWriter', ([], {}), '()\n', (5294, 5296), False, 'from torch.utils.tensorboard import SummaryWriter\n'), ((6426, 6487), 'reagent.gym.agents.agent.Agent.create_for_env_with_serving_policy', 'Agent.create_for_env_with_serving_policy', (['env', 'serving_policy'], {}), '(env, serving_policy)\n', (6466, 6487), False, 'from reagent.gym.agents.agent import Agent\n'), ((6730, 6751), 'numpy.mean', 'np.mean', (['eval_rewards'], {}), '(eval_rewards)\n', (6737, 6751), True, 'import numpy as np\n'), ((6981, 6996), 'unittest.main', 'unittest.main', ([], {}), '()\n', (6994, 6996), False, 'import unittest\n'), ((4295, 4310), 'reagent.core.types.RewardOptions', 'RewardOptions', ([], {}), '()\n', (4308, 4310), False, 'from reagent.core.types import RewardOptions\n'), ((4625, 4645), 'torch.device', 'torch.device', (['"""cuda"""'], {}), "('cuda')\n", (4637, 4645), False, 'import torch\n'), ((4662, 4681), 'torch.device', 'torch.device', (['"""cpu"""'], {}), "('cpu')\n", (4674, 4681), False, 'import torch\n'), ((5306, 5336), 'reagent.tensorboardX.summary_writer_context', 'summary_writer_context', (['writer'], {}), '(writer)\n', (5328, 5336), False, 'from reagent.tensorboardX import summary_writer_context\n'), ((6140, 6161), 'numpy.max', 'np.max', (['train_rewards'], {}), '(train_rewards)\n', (6146, 6161), True, 'import numpy as np\n'), ((3419, 3444), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (3442, 3444), False, 'import torch\n'), ((5434, 5502), 'reagent.gym.runners.gymrunner.run_episode', 'run_episode', ([], {'env': 'env', 'agent': 'agent', 'mdp_id': 'i', 'max_steps': 'env.max_steps'}), '(env=env, agent=agent, mdp_id=i, max_steps=env.max_steps)\n', (5445, 5502), False, 'from reagent.gym.runners.gymrunner import evaluate_for_n_episodes, run_episode\n'), ((6209, 6230), 'numpy.max', 'np.max', (['train_rewards'], {}), '(train_rewards)\n', (6215, 6230), True, 'import numpy as np\n'), ((6508, 6603), 'reagent.gym.runners.gymrunner.evaluate_for_n_episodes', 'evaluate_for_n_episodes', ([], {'n': 'num_eval_episodes', 'env': 'env', 'agent': 'agent', 'max_steps': 'env.max_steps'}), '(n=num_eval_episodes, env=env, agent=agent,\n    max_steps=env.max_steps)\n', (6531, 6603), False, 'from reagent.gym.runners.gymrunner import evaluate_for_n_episodes, run_episode\n'), ((2997, 3032), 'os.path.join', 'os.path.join', (['curr_dir', 'config_path'], {}), '(curr_dir, config_path)\n', (3009, 3032), False, 'import os\n'), ((3657, 3692), 'os.path.join', 'os.path.join', (['curr_dir', 'config_path'], {}), '(curr_dir, config_path)\n', (3669, 3692), False, 'import os\n'), ((4178, 4207), 'pprint.pformat', 'pprint.pformat', (['normalization'], {}), '(normalization)\n', (4192, 4207), False, 'import pprint\n'), ((5917, 5939), 'numpy.mean', 'np.mean', (['train_rewards'], {}), '(train_rewards)\n', (5924, 5939), True, 'import numpy as np\n'), ((5949, 5970), 'numpy.max', 'np.max', (['train_rewards'], {}), '(train_rewards)\n', (5955, 5970), True, 'import numpy as np\n'), ((6799, 6819), 'numpy.max', 'np.max', (['eval_rewards'], {}), '(eval_rewards)\n', (6805, 6819), True, 'import numpy as np\n')]