# Path: babyai/evaluate.py
# def evaluate_demo_agent(agent, episodes):
#     logs = {"num_frames_per_episode": [], "return_per_episode": []}
# 
#     number_of_demos = len(agent.demos)
# 
#     for demo_id in range(min(number_of_demos, episodes)):
#         logs["num_frames_per_episode"].append(len(agent.demos[demo_id]))
# 
#     return logs
# 
# def batch_evaluate(agent, env_name, seed, episodes, return_obss_actions=False, pixel=False):
#     num_envs = min(256, episodes)
# 
#     envs = []
#     for i in range(num_envs):
#         env = gym.make(env_name)
#         if pixel:
#             env = RGBImgPartialObsWrapper(env)
#         envs.append(env)
#     env = ManyEnvs(envs)
# 
#     logs = {
#         "num_frames_per_episode": [],
#         "return_per_episode": [],
#         "observations_per_episode": [],
#         "actions_per_episode": [],
#         "seed_per_episode": []
#     }
# 
#     for i in range((episodes + num_envs - 1) // num_envs):
#         seeds = range(seed + i * num_envs, seed + (i + 1) * num_envs)
#         env.seed(seeds)
# 
#         many_obs = env.reset()
# 
#         cur_num_frames = 0
#         num_frames = np.zeros((num_envs,), dtype='int64')
#         returns = np.zeros((num_envs,))
#         already_done = np.zeros((num_envs,), dtype='bool')
#         if return_obss_actions:
#             obss = [[] for _ in range(num_envs)]
#             actions = [[] for _ in range(num_envs)]
#         while (num_frames == 0).any():
#             action = agent.act_batch(many_obs)['action']
#             if return_obss_actions:
#                 for i in range(num_envs):
#                     if not already_done[i]:
#                         obss[i].append(many_obs[i])
#                         actions[i].append(action[i].item())
#             many_obs, reward, done, _ = env.step(action)
#             agent.analyze_feedback(reward, done)
#             done = np.array(done)
#             just_done = done & (~already_done)
#             returns += reward * just_done
#             cur_num_frames += 1
#             num_frames[just_done] = cur_num_frames
#             already_done[done] = True
# 
#         logs["num_frames_per_episode"].extend(list(num_frames))
#         logs["return_per_episode"].extend(list(returns))
#         logs["seed_per_episode"].extend(list(seeds))
#         if return_obss_actions:
#             logs["observations_per_episode"].extend(obss)
#             logs["actions_per_episode"].extend(actions)
# 
#     return logs
# 
# def evaluate(agent, env, episodes, model_agent=True, offsets=None):
#     # Initialize logs
#     if model_agent:
#         agent.model.eval()
#     logs = {"num_frames_per_episode": [], "return_per_episode": [], "observations_per_episode": []}
# 
#     if offsets:
#         count = 0
# 
#     for i in range(episodes):
#         if offsets:
#             # Ensuring test on seed offsets that generated successful demonstrations
#             while count != offsets[i]:
#                 obs = env.reset()
#                 count += 1
# 
#         obs = env.reset()
#         agent.on_reset()
#         done = False
# 
#         num_frames = 0
#         returnn = 0
#         obss = []
#         while not done:
#             action = agent.act(obs)['action']
#             obss.append(obs)
#             obs, reward, done, _ = env.step(action)
#             agent.analyze_feedback(reward, done)
#             num_frames += 1
#             returnn += reward
# 
# 
#         logs["observations_per_episode"].append(obss)
#         logs["num_frames_per_episode"].append(num_frames)
#         logs["return_per_episode"].append(returnn)
#     if model_agent:
#         agent.model.train()
#     return logs


import argparse
import gym
import time
import datetime
import babyai.utils as utils
from babyai.evaluate import evaluate_demo_agent, batch_evaluate, evaluate
                    help="origin of the demonstrations: human | agent (REQUIRED or --model or --demos REQUIRED)")
parser.add_argument("--demos", default=None,
                    help="name of the demos file (REQUIRED or --demos-origin or --model REQUIRED)")
parser.add_argument("--episodes", type=int, default=1000,
                    help="number of episodes of evaluation (default: 1000)")
parser.add_argument("--seed", type=int, default=int(1e9),
                    help="random seed")
parser.add_argument("--argmax", action="store_true", default=False,
                    help="action with highest probability is selected for model agent")
parser.add_argument("--contiguous-episodes", action="store_true", default=False,
                    help="Make sure episodes on which evaluation is done are contiguous")
parser.add_argument("--worst-episodes-to-show", type=int, default=10,
                    help="The number of worse episodes to show")


def main(args, seed, episodes):
    # Set seed for all randomness sources
    utils.seed(seed)

    # Define agent

    env = gym.make(args.env)
    env.seed(seed)
    agent = utils.load_agent(env, args.model, args.demos, args.demos_origin, args.argmax, args.env)
    if args.model is None and args.episodes > len(agent.demos):
        # Set the number of episodes to be the number of demos
        episodes = len(agent.demos)

    # Evaluate
    if isinstance(agent, utils.DemoAgent):

# Path: babyai/evaluate.py
# def evaluate_demo_agent(agent, episodes):
#     logs = {"num_frames_per_episode": [], "return_per_episode": []}
# 
#     number_of_demos = len(agent.demos)
# 
#     for demo_id in range(min(number_of_demos, episodes)):
#         logs["num_frames_per_episode"].append(len(agent.demos[demo_id]))
# 
#     return logs
# 
# def batch_evaluate(agent, env_name, seed, episodes, return_obss_actions=False, pixel=False):
#     num_envs = min(256, episodes)
# 
#     envs = []
#     for i in range(num_envs):
#         env = gym.make(env_name)
#         if pixel:
#             env = RGBImgPartialObsWrapper(env)
#         envs.append(env)
#     env = ManyEnvs(envs)
# 
#     logs = {
#         "num_frames_per_episode": [],
#         "return_per_episode": [],
#         "observations_per_episode": [],
#         "actions_per_episode": [],
#         "seed_per_episode": []
#     }
# 
#     for i in range((episodes + num_envs - 1) // num_envs):
#         seeds = range(seed + i * num_envs, seed + (i + 1) * num_envs)
#         env.seed(seeds)
# 
#         many_obs = env.reset()
# 
#         cur_num_frames = 0
#         num_frames = np.zeros((num_envs,), dtype='int64')
#         returns = np.zeros((num_envs,))
#         already_done = np.zeros((num_envs,), dtype='bool')
#         if return_obss_actions:
#             obss = [[] for _ in range(num_envs)]
#             actions = [[] for _ in range(num_envs)]
#         while (num_frames == 0).any():
#             action = agent.act_batch(many_obs)['action']
#             if return_obss_actions:
#                 for i in range(num_envs):
#                     if not already_done[i]:
#                         obss[i].append(many_obs[i])
#                         actions[i].append(action[i].item())
#             many_obs, reward, done, _ = env.step(action)
#             agent.analyze_feedback(reward, done)
#             done = np.array(done)
#             just_done = done & (~already_done)
#             returns += reward * just_done
#             cur_num_frames += 1
#             num_frames[just_done] = cur_num_frames
#             already_done[done] = True
# 
#         logs["num_frames_per_episode"].extend(list(num_frames))
#         logs["return_per_episode"].extend(list(returns))
#         logs["seed_per_episode"].extend(list(seeds))
#         if return_obss_actions:
#             logs["observations_per_episode"].extend(obss)
#             logs["actions_per_episode"].extend(actions)
# 
#     return logs
# 
# def evaluate(agent, env, episodes, model_agent=True, offsets=None):
#     # Initialize logs
#     if model_agent:
#         agent.model.eval()
#     logs = {"num_frames_per_episode": [], "return_per_episode": [], "observations_per_episode": []}
# 
#     if offsets:
#         count = 0
# 
#     for i in range(episodes):
#         if offsets:
#             # Ensuring test on seed offsets that generated successful demonstrations
#             while count != offsets[i]:
#                 obs = env.reset()
#                 count += 1
# 
#         obs = env.reset()
#         agent.on_reset()
#         done = False
# 
#         num_frames = 0
#         returnn = 0
#         obss = []
#         while not done:
#             action = agent.act(obs)['action']
#             obss.append(obs)
#             obs, reward, done, _ = env.step(action)
#             agent.analyze_feedback(reward, done)
#             num_frames += 1
#             returnn += reward
# 
# 
#         logs["observations_per_episode"].append(obss)
#         logs["num_frames_per_episode"].append(num_frames)
#         logs["return_per_episode"].append(returnn)
#     if model_agent:
#         agent.model.train()
#     return logs


# Path: scripts/evaluate.py
import argparse
import gym
import time
import datetime
import babyai.utils as utils
from babyai.evaluate import evaluate_demo_agent, batch_evaluate, evaluate

                    help="origin of the demonstrations: human | agent (REQUIRED or --model or --demos REQUIRED)")
parser.add_argument("--demos", default=None,
                    help="name of the demos file (REQUIRED or --demos-origin or --model REQUIRED)")
parser.add_argument("--episodes", type=int, default=1000,
                    help="number of episodes of evaluation (default: 1000)")
parser.add_argument("--seed", type=int, default=int(1e9),
                    help="random seed")
parser.add_argument("--argmax", action="store_true", default=False,
                    help="action with highest probability is selected for model agent")
parser.add_argument("--contiguous-episodes", action="store_true", default=False,
                    help="Make sure episodes on which evaluation is done are contiguous")
parser.add_argument("--worst-episodes-to-show", type=int, default=10,
                    help="The number of worse episodes to show")


def main(args, seed, episodes):
    # Set seed for all randomness sources
    utils.seed(seed)

    # Define agent

    env = gym.make(args.env)
    env.seed(seed)
    agent = utils.load_agent(env, args.model, args.demos, args.demos_origin, args.argmax, args.env)
    if args.model is None and args.episodes > len(agent.demos):
        # Set the number of episodes to be the number of demos
        episodes = len(agent.demos)

    # Evaluate
    if isinstance(agent, utils.DemoAgent):

        logs = evaluate_demo_agent(agent, episodes)