#!/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()

[((37, 9, 37, 36), 'logging.getLogger', 'logging.getLogger', ({(37, 27, 37, 35): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((71, 11, 71, 36), 'os.path.dirname', 'os.path.dirname', ({(71, 27, 71, 35): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((76, 5, 76, 36), 'parameterized.parameterized.expand', 'parameterized.expand', ({(76, 26, 76, 35): 'GYM_TESTS'}, {}), '(GYM_TESTS)', False, 'from parameterized import parameterized\n'), ((87, 5, 87, 36), 'parameterized.parameterized.expand', 'parameterized.expand', ({(87, 26, 87, 35): 'GYM_TESTS'}, {}), '(GYM_TESTS)', False, 'from parameterized import parameterized\n'), ((116, 20, 116, 41), 'reagent.gym.utils.build_normalizer', 'build_normalizer', ({(116, 37, 116, 40): 'env'}, {}), '(env)', False, 'from reagent.gym.utils import build_normalizer, fill_replay_buffer\n'), ((126, 20, 128, 5), 'reagent.replay_memory.circular_replay_buffer.ReplayBuffer', 'ReplayBuffer', (), '', False, 'from reagent.replay_memory.circular_replay_buffer import ReplayBuffer\n'), ((133, 4, 135, 5), 'reagent.gym.utils.fill_replay_buffer', 'fill_replay_buffer', (), '', False, 'from reagent.gym.utils import build_normalizer, fill_replay_buffer\n'), ((137, 16, 144, 5), 'reagent.gym.agents.post_step.train_with_replay_buffer_post_step', 'train_with_replay_buffer_post_step', (), '', False, 'from reagent.gym.agents.post_step import train_with_replay_buffer_post_step\n'), ((146, 12, 148, 5), 'reagent.gym.agents.agent.Agent.create_for_env', 'Agent.create_for_env', (), '', False, 'from reagent.gym.agents.agent import Agent\n'), ((150, 13, 150, 28), 'torch.utils.tensorboard.SummaryWriter', 'SummaryWriter', ({}, {}), '()', False, 'from torch.utils.tensorboard import SummaryWriter\n'), ((176, 12, 176, 73), 'reagent.gym.agents.agent.Agent.create_for_env_with_serving_policy', 'Agent.create_for_env_with_serving_policy', ({(176, 53, 176, 56): 'env', (176, 58, 176, 72): 'serving_policy'}, {}), '(env, serving_policy)', False, 'from reagent.gym.agents.agent import Agent\n'), ((184, 16, 184, 37), 'numpy.mean', 'np.mean', ({(184, 24, 184, 36): 'eval_rewards'}, {}), '(eval_rewards)', True, 'import numpy as np\n'), ((192, 4, 192, 19), 'unittest.main', 'unittest.main', ({}, {}), '()', False, 'import unittest\n'), ((120, 43, 120, 58), 'reagent.core.types.RewardOptions', 'RewardOptions', ({}, {}), '()', False, 'from reagent.core.types import RewardOptions\n'), ((130, 13, 130, 33), 'torch.device', 'torch.device', ({(130, 26, 130, 32): '"""cuda"""'}, {}), "('cuda')", False, 'import torch\n'), ((130, 50, 130, 69), 'torch.device', 'torch.device', ({(130, 63, 130, 68): '"""cpu"""'}, {}), "('cpu')", False, 'import torch\n'), ((151, 9, 151, 39), 'reagent.tensorboardX.summary_writer_context', 'summary_writer_context', ({(151, 32, 151, 38): 'writer'}, {}), '(writer)', False, 'from reagent.tensorboardX import summary_writer_context\n'), ((170, 11, 170, 32), 'numpy.max', 'np.max', ({(170, 18, 170, 31): 'train_rewards'}, {}), '(train_rewards)', True, 'import numpy as np\n'), ((91, 25, 91, 50), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((154, 25, 156, 13), 'reagent.gym.runners.gymrunner.run_episode', 'run_episode', (), '', False, 'from reagent.gym.runners.gymrunner import evaluate_for_n_episodes, run_episode\n'), ((171, 23, 171, 44), 'numpy.max', 'np.max', ({(171, 30, 171, 43): 'train_rewards'}, {}), '(train_rewards)', True, 'import numpy as np\n'), ((178, 19, 180, 5), 'reagent.gym.runners.gymrunner.evaluate_for_n_episodes', 'evaluate_for_n_episodes', (), '', False, 'from reagent.gym.runners.gymrunner import evaluate_for_n_episodes, run_episode\n'), ((81, 24, 81, 59), 'os.path.join', 'os.path.join', ({(81, 37, 81, 45): 'curr_dir', (81, 47, 81, 58): 'config_path'}, {}), '(curr_dir, config_path)', False, 'import os\n'), ((96, 24, 96, 59), 'os.path.join', 'os.path.join', ({(96, 37, 96, 45): 'curr_dir', (96, 47, 96, 58): 'config_path'}, {}), '(curr_dir, config_path)', False, 'import os\n'), ((117, 39, 117, 68), 'pprint.pformat', 'pprint.pformat', ({(117, 54, 117, 67): 'normalization'}, {}), '(normalization)', False, 'import pprint\n'), ((166, 28, 166, 50), 'numpy.mean', 'np.mean', ({(166, 36, 166, 49): 'train_rewards'}, {}), '(train_rewards)', True, 'import numpy as np\n'), ((166, 60, 166, 81), 'numpy.max', 'np.max', ({(166, 67, 166, 80): 'train_rewards'}, {}), '(train_rewards)', True, 'import numpy as np\n'), ((185, 47, 185, 67), 'numpy.max', 'np.max', ({(185, 54, 185, 66): 'eval_rewards'}, {}), '(eval_rewards)', True, 'import numpy as np\n')]