from collections import deque, defaultdict
import os
import sys
import logging
import time
import json
import gym
import torch.nn as nn
import torch
import numpy as np
import matplotlib.pyplot as plt

from model import RL_Policy, Semantic_Mapping
from utils.storage import GlobalRolloutStorage
from envs import make_vec_envs
from arguments import get_args
import algo

os.environ["OMP_NUM_THREADS"] = "1"


def main():
    args = get_args()

    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    if args.cuda:
        torch.cuda.manual_seed(args.seed)

    # Setup Logging
    log_dir = "{}/models/{}/".format(args.dump_location, args.exp_name)
    dump_dir = "{}/dump/{}/".format(args.dump_location, args.exp_name)

    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    if not os.path.exists(dump_dir):
        os.makedirs(dump_dir)

    logging.basicConfig(
        filename=log_dir + 'train.log',
        level=logging.INFO)
    print("Dumping at {}".format(log_dir))
    print(args)
    logging.info(args)

    # Logging and loss variables
    num_scenes = args.num_processes
    num_episodes = int(args.num_eval_episodes)
    device = args.device = torch.device("cuda:0" if args.cuda else "cpu")

    g_masks = torch.ones(num_scenes).float().to(device)

    best_g_reward = -np.inf

    # one episode per process for both train and eval
    # for eval, one scene per process
    if args.eval:
        episode_success = []
        episode_spl = []
        episode_dist = []
        for _ in range(args.num_processes):
            episode_success.append(deque(maxlen=num_episodes))
            episode_spl.append(deque(maxlen=num_episodes))
            episode_dist.append(deque(maxlen=num_episodes))
    # for train, different episodes of same scene per process
    else:
        episode_success = deque(maxlen=1000)
        episode_spl = deque(maxlen=1000)
        episode_dist = deque(maxlen=1000)

    finished = np.zeros((args.num_processes))
    wait_env = np.zeros((args.num_processes))

    g_episode_rewards = deque(maxlen=1000)

    g_value_losses = deque(maxlen=1000)
    g_action_losses = deque(maxlen=1000)
    g_dist_entropies = deque(maxlen=1000)

    per_step_g_rewards = deque(maxlen=1000)

    g_process_rewards = np.zeros((num_scenes))

    # Starting environments
    torch.set_num_threads(1)
    envs = make_vec_envs(args)
    obs, infos = envs.reset()

    full_episode_data = []
    episode_data = [None] * num_scenes
    for e, info in enumerate(infos):
        cInfo = info.copy()
        cInfo["episode_data"]["positions"] = []
        cInfo["episode_data"]["gt_positions"] = []
        cInfo["episode_data"]["goal_rewards"] = []
        cInfo["episode_data"]["explore_rewards"] = []
        cInfo["episode_data"]["policy_goals"] = []
        cInfo["episode_data"]["used_policy"] = []
        episode_data[e] = cInfo["episode_data"]

    torch.set_grad_enabled(False)

    # Initialize map variables:
    # Full map consists of multiple channels containing the following:
    # 1. Obstacle Map
    # 2. Exploread Area (places that are known to be free or occupied)
    # 3. Current Agent Location
    # 4. Past Agent Locations
    # 5,6,7,.. : Semantic Categories
    nc = args.num_sem_categories + 4  # num channels

    # Calculating full and local map sizes
    map_size = args.map_size_cm // args.map_resolution
    full_w, full_h = map_size, map_size
    local_w = int(full_w / args.global_downscaling)
    local_h = int(full_h / args.global_downscaling)

    # Initializing full and local map
    full_map = torch.zeros(num_scenes, nc, full_w, full_h).float().to(device)
    local_map = torch.zeros(num_scenes, nc, local_w,
                            local_h).float().to(device)

    # Initial full and local pose
    full_pose = torch.zeros(num_scenes, 3).float().to(device)
    local_pose = torch.zeros(num_scenes, 3).float().to(device)

    # Origin of local map
    origins = np.zeros((num_scenes, 3))

    # Local Map Boundaries
    lmb = np.zeros((num_scenes, 4)).astype(int)

    # Planner pose inputs has 7 dimensions
    # 1-3 store continuous global agent location
    # 4-7 store local map boundaries
    planner_pose_inputs = np.zeros((num_scenes, 7))

    # get local boundary (x1, x2, y1, y2) given local agent position (x, y) and map size
    def get_local_map_boundaries(agent_loc, local_sizes, full_sizes):
        loc_r, loc_c = agent_loc
        local_w, local_h = local_sizes
        full_w, full_h = full_sizes

        if args.global_downscaling > 1:
            gx1, gy1 = loc_r - local_w // 2, loc_c - local_h // 2
            gx2, gy2 = gx1 + local_w, gy1 + local_h
            if gx1 < 0:
                gx1, gx2 = 0, local_w
            if gx2 > full_w:
                gx1, gx2 = full_w - local_w, full_w

            if gy1 < 0:
                gy1, gy2 = 0, local_h
            if gy2 > full_h:
                gy1, gy2 = full_h - local_h, full_h
        else:
            gx1, gx2, gy1, gy2 = 0, full_w, 0, full_h

        return [gx1, gx2, gy1, gy2]

    # initialize global and local maps and poses given that initial position 
    # is at map center with 0 orientation
    def init_map_and_pose():
        full_map.fill_(0.)
        full_pose.fill_(0.)
        full_pose[:, :2] = args.map_size_cm / 100.0 / 2.0 

        locs = full_pose.cpu().numpy()
        planner_pose_inputs[:, :3] = locs
        for e in range(num_scenes):
            r, c = locs[e, 1], locs[e, 0]
            loc_r, loc_c = [int(r * 100.0 / args.map_resolution),
                            int(c * 100.0 / args.map_resolution)]

            # 3x3 grid around agent location is considered explored
            full_map[e, 2:4, loc_r - 1:loc_r + 2, loc_c - 1:loc_c + 2] = 1.0

            lmb[e] = get_local_map_boundaries((loc_r, loc_c),
                                              (local_w, local_h),
                                              (full_w, full_h))

            planner_pose_inputs[e, 3:] = lmb[e]
            origins[e] = [lmb[e][2] * args.map_resolution / 100.0,
                          lmb[e][0] * args.map_resolution / 100.0, 0.]

        for e in range(num_scenes):
            local_map[e] = full_map[e, :,
                                    lmb[e, 0]:lmb[e, 1],
                                    lmb[e, 2]:lmb[e, 3]]
            local_pose[e] = full_pose[e] - \
                torch.from_numpy(origins[e]).to(device).float()

    # identical to above, except for specific environment
    def init_map_and_pose_for_env(e):
        full_map[e].fill_(0.)
        full_pose[e].fill_(0.)
        full_pose[e, :2] = args.map_size_cm / 100.0 / 2.0

        locs = full_pose[e].cpu().numpy()
        planner_pose_inputs[e, :3] = locs
        r, c = locs[1], locs[0]
        loc_r, loc_c = [int(r * 100.0 / args.map_resolution),
                        int(c * 100.0 / args.map_resolution)]

        full_map[e, 2:4, loc_r - 1:loc_r + 2, loc_c - 1:loc_c + 2] = 1.0

        lmb[e] = get_local_map_boundaries((loc_r, loc_c),
                                          (local_w, local_h),
                                          (full_w, full_h))

        planner_pose_inputs[e, 3:] = lmb[e]
        origins[e] = [lmb[e][2] * args.map_resolution / 100.0,
                      lmb[e][0] * args.map_resolution / 100.0, 0.]

        local_map[e] = full_map[e, :, lmb[e, 0]:lmb[e, 1], lmb[e, 2]:lmb[e, 3]]
        local_pose[e] = full_pose[e] - \
            torch.from_numpy(origins[e]).to(device).float()

    # reward is the newly explored area in a given step (in m^2)
    def update_intrinsic_rew(e):
        prev_explored_area = full_map[e, 1].sum(1).sum(0)
        full_map[e, :, lmb[e, 0]:lmb[e, 1], lmb[e, 2]:lmb[e, 3]] = \
            local_map[e]
        curr_explored_area = full_map[e, 1].sum(1).sum(0)
        intrinsic_rews[e] = curr_explored_area - prev_explored_area
        intrinsic_rews[e] *= (args.map_resolution / 100.)**2  # to m^2

    def get_random_goal(e):
        for _ in range(20):
            goal = np.random.rand(2)
            goal = [int(goal[0] * local_w), int(goal[1] * local_w)]
            goal = [min(goal[0], int(local_w-1)), min(goal[1], int(local_w-1))]
            if not local_map[e, 1, goal[0], goal[1]]: break
        return goal

    init_map_and_pose()

    # Global policy observation space
    ngc = 8 + args.num_sem_categories
    es = 2
    g_observation_space = gym.spaces.Box(0, 1, # binary local map
                                         (ngc,
                                          local_w,
                                          local_h), dtype='uint8')

    # Semantic Mapping
    sem_map_module = Semantic_Mapping(args).to(device)
    sem_map_module.eval()

    intrinsic_rews = torch.zeros(num_scenes).to(device)

    # Predict semantic map from frame 1
    poses = torch.from_numpy(np.asarray(
        [infos[env_idx]['sensor_pose'] for env_idx in range(num_scenes)])
    ).float().to(device)

    # args (obs, pose_obs, maps_last, poses_last)
    _, local_map, _, local_pose = \
        sem_map_module(obs, poses, local_map, local_pose)

    locs = local_pose.cpu().numpy()
    
    for e in range(num_scenes):
        r, c = locs[e, 1], locs[e, 0]
        loc_r, loc_c = [int(r * 100.0 / args.map_resolution),
                        int(c * 100.0 / args.map_resolution)]

        local_map[e, 2:4, loc_r - 1:loc_r + 2, loc_c - 1:loc_c + 2] = 1.
        episode_data[e]["positions"].append([int(loc_r + lmb[e, 0]), int(loc_c + lmb[e, 2]), int(locs[e, 2])])
        episode_data[e]["gt_positions"].append(list(infos[e]["gt_pos"]))

    global_goals = [get_random_goal(e) for e in range(num_scenes)]
     
    goal_maps = [np.zeros((local_w, local_h)) for _ in range(num_scenes)]

    for e in range(num_scenes):
        goal_maps[e][global_goals[e][0], global_goals[e][1]] = 1
        episode_data[e]["policy_goals"].append(((global_goals[e] + lmb[e, [0,2]]).tolist(), 0))
        episode_data[e]["used_policy"].append(True)

    planner_inputs = [{} for e in range(num_scenes)]
    for e, p_input in enumerate(planner_inputs):
        p_input['map_pred'] = local_map[e, 0, :, :].cpu().numpy() # obstacles
        p_input['exp_pred'] = local_map[e, 1, :, :].cpu().numpy() # explored
        p_input['pose_pred'] = planner_pose_inputs[e] # global location+local map bounds
        p_input['goal'] = goal_maps[e]  # global_goals[e]
        p_input['new_goal'] = 1
        p_input['found_goal'] = 0
        p_input['wait'] = wait_env[e] or finished[e]
        if args.visualize or args.print_images:
            local_map[e, -1, :, :] = 1e-5 # TODO: what is this?
            # single channel where each grid loc is cat ID
            p_input['sem_map_pred'] = local_map[e, 4:, :, :
                                                ].argmax(0).cpu().numpy()

    obs, _, done, infos = envs.plan_act_and_preprocess(planner_inputs)

    start = time.time()
    g_reward = 0

    torch.set_grad_enabled(False)
    spl_per_category = defaultdict(list)
    success_per_category = defaultdict(list)

    for step in range(args.num_training_frames // args.num_processes + 1):
        if finished.sum() == args.num_processes:
            break

        g_step = (step // args.num_local_steps) % args.num_global_steps # global step num in PPO
        l_step = step % args.num_local_steps # local step num in global step

        # ------------------------------------------------------------------
        # Reinitialize variables when episode ends
        l_masks = torch.FloatTensor([0 if x else 1
                                     for x in done]).to(device)
        g_masks *= l_masks

        for e, x in enumerate(done):
            if x:
                spl = infos[e]['spl']
                success = infos[e]['success']
                dist = infos[e]['distance_to_goal']
                spl_per_category[infos[e]['goal_name']].append(spl)
                success_per_category[infos[e]['goal_name']].append(success)
                if args.eval:
                    episode_success[e].append(success)
                    episode_spl[e].append(spl)
                    episode_dist[e].append(dist)
                    if len(episode_success[e]) == num_episodes:
                        finished[e] = 1
                    episode_data[e]["success"] = success
                    episode_data[e]["spl"] = spl
                    episode_data[e]["distance_to_goal"] = dist
                    full_map[e, :, lmb[e, 0]:lmb[e, 1], lmb[e, 2]:lmb[e, 3]] = local_map[e] 
                    episode_data[e]["explored_area"] = full_map[e, 1].sum(1).sum(0).item()
                    scene = episode_data[e]["scene_id"][16:-4]
                    if args.save_maps:
                        np.save('{}/maparr_{}_{}'.format(dump_dir, scene, episode_data[e]['episode_id']), full_map[e].cpu().numpy())
                    full_episode_data.append(episode_data[e])
                    
                    cInfo = infos[e].copy()
                    cInfo["episode_data"]["positions"] = []
                    cInfo["episode_data"]["gt_positions"] = []
                    cInfo["episode_data"]["goal_rewards"] = []
                    cInfo["episode_data"]["explore_rewards"] = []
                    cInfo["episode_data"]["policy_goals"] = []
                    cInfo["episode_data"]["used_policy"] = []
                    episode_data[e] = cInfo["episode_data"]
                else:
                    episode_success.append(success)
                    episode_spl.append(spl)
                    episode_dist.append(dist)
                wait_env[e] = 1.
                update_intrinsic_rew(e)
                init_map_and_pose_for_env(e)
        # ------------------------------------------------------------------

        # ------------------------------------------------------------------
        # Semantic Mapping Module
        poses = torch.from_numpy(np.asarray(
            [infos[env_idx]['sensor_pose'] for env_idx
             in range(num_scenes)])
        ).float().to(device)

        _, local_map, _, local_pose = \
            sem_map_module(obs, poses, local_map, local_pose)

        locs = local_pose.cpu().numpy()
        planner_pose_inputs[:, :3] = locs + origins
        local_map[:, 2, :, :].fill_(0.)  # Resetting current location channel
        # update current location
        for e in range(num_scenes):
            r, c = locs[e, 1], locs[e, 0]
            loc_r, loc_c = [int(r * 100.0 / args.map_resolution),
                            int(c * 100.0 / args.map_resolution)]
            local_map[e, 2:4, loc_r - 2:loc_r + 3, loc_c - 2:loc_c + 3] = 1.
            if args.eval and not wait_env[e]:
                episode_data[e]["positions"].append([int(loc_r + lmb[e, 0]), int(loc_c + lmb[e, 2]), int(locs[e, 2])])
                episode_data[e]["gt_positions"].append(list(infos[e]["gt_pos"]))

        # ------------------------------------------------------------------

        # ------------------------------------------------------------------
        # Global Policy
        if l_step == args.num_local_steps - 1:
            # For every global step, update the full and local maps
            for e in range(num_scenes):
                if wait_env[e] == 1:  # New episode
                    wait_env[e] = 0.
                else:
                    update_intrinsic_rew(e)
                # update global map and pose based on new position in old local frame
                full_map[e, :, lmb[e, 0]:lmb[e, 1], lmb[e, 2]:lmb[e, 3]] = \
                    local_map[e]
                full_pose[e] = local_pose[e] + \
                    torch.from_numpy(origins[e]).to(device).float()

                # center the local frame based on new position
                locs = full_pose[e].cpu().numpy()
                r, c = locs[1], locs[0]
                loc_r, loc_c = [int(r * 100.0 / args.map_resolution),
                                int(c * 100.0 / args.map_resolution)]

                lmb[e] = get_local_map_boundaries((loc_r, loc_c),
                                                  (local_w, local_h),
                                                  (full_w, full_h))

                # compute new local map and pose based on new local frame
                planner_pose_inputs[e, 3:] = lmb[e]
                origins[e] = [lmb[e][2] * args.map_resolution / 100.0,
                              lmb[e][0] * args.map_resolution / 100.0, 0.]

                local_map[e] = full_map[e, :,
                                        lmb[e, 0]:lmb[e, 1],
                                        lmb[e, 2]:lmb[e, 3]]
                local_pose[e] = full_pose[e] - \
                    torch.from_numpy(origins[e]).to(device).float()

            locs = local_pose.cpu().numpy()

            # Get exploration reward and metrics
            g_reward = torch.from_numpy(np.asarray(
                [infos[env_idx]['g_reward'] for env_idx in range(num_scenes)])
            ).float().to(device)
            g_reward += args.intrinsic_rew_coeff * intrinsic_rews.detach()
            for e in range(num_scenes):
                if args.eval and not wait_env[e]:
                    episode_data[e]["goal_rewards"].append(infos[e]["g_reward"])
                    episode_data[e]["explore_rewards"].append(intrinsic_rews[e].item())

            g_process_rewards += g_reward.cpu().numpy()
            g_total_rewards = g_process_rewards * \
                (1 - g_masks.cpu().numpy())
            g_process_rewards *= g_masks.cpu().numpy()
            per_step_g_rewards.append(np.mean(g_reward.cpu().numpy()))

            if np.sum(g_total_rewards) != 0:
                for total_rew in g_total_rewards:
                    if total_rew != 0:
                        g_episode_rewards.append(total_rew)

            global_goals = [get_random_goal(e) for e in range(num_scenes)]

            for e in range(num_scenes):
                if args.eval and not wait_env[e]:
                    episode_data[e]["policy_goals"].append(((global_goals[e] + lmb[e, [0,2]]).tolist(), 0))

            g_reward = 0
            g_masks = torch.ones(num_scenes).float().to(device)

        # ------------------------------------------------------------------

        # ------------------------------------------------------------------
        # Update long-term goal if target object is found
        found_goal = [0 for _ in range(num_scenes)]
        goal_maps = [np.zeros((local_w, local_h)) for _ in range(num_scenes)]

        # If goal category not found in map, goal is the location sampled by
        # policy
        for e in range(num_scenes):
            goal_maps[e][global_goals[e][0], global_goals[e][1]] = 1
            if args.eval and not wait_env[e]:
                episode_data[e]["used_policy"].append(True)

        # Else if goal category found in map, use all locations where prob of goal
        # obj existing is > 0 as the goal map for planner
        for e in range(num_scenes):
            cn = infos[e]['goal_cat_id'] + 4
            if local_map[e, cn, :, :].sum() != 0.:
                cat_semantic_map = local_map[e, cn, :, :].cpu().numpy()
                cat_semantic_scores = cat_semantic_map
                cat_semantic_scores[cat_semantic_scores > 0] = 1.
                goal_maps[e] = cat_semantic_scores
                found_goal[e] = 1
                if args.eval and not wait_env[e]:
                    episode_data[e]["used_policy"][-1] = False
        # ------------------------------------------------------------------

        # ------------------------------------------------------------------
        # Take action and get next observation
        planner_inputs = [{} for e in range(num_scenes)]
        for e, p_input in enumerate(planner_inputs):
            p_input['map_pred'] = local_map[e, 0, :, :].cpu().numpy()
            p_input['exp_pred'] = local_map[e, 1, :, :].cpu().numpy()
            p_input['pose_pred'] = planner_pose_inputs[e]
            p_input['goal'] = goal_maps[e]  # global_goals[e]
            p_input['new_goal'] = l_step == args.num_local_steps - 1
            p_input['found_goal'] = found_goal[e]
            p_input['wait'] = wait_env[e] or finished[e]
            if args.visualize or args.print_images:
                local_map[e, -1, :, :] = 1e-5
                p_input['sem_map_pred'] = local_map[e, 4:, :,
                                                    :].argmax(0).cpu().numpy()

        obs, _, done, infos = envs.plan_act_and_preprocess(planner_inputs)
        # ------------------------------------------------------------------
        # Logging

        if len(full_episode_data) % args.episode_save_interval == 0:
            with open('{}/{}_episode_data.json'.format(
                dump_dir, args.split), 'w') as f:
                    json.dump(full_episode_data, f)

        if step % args.log_interval == 0:
            end = time.time()
            time_elapsed = time.gmtime(end - start)
            log = " ".join([
                "Time: {0:0=2d}d".format(time_elapsed.tm_mday - 1),
                "{},".format(time.strftime("%Hh %Mm %Ss", time_elapsed)),
                "num timesteps {},".format(step * num_scenes),
                "FPS {},".format(int(step * num_scenes / (end - start)))
            ])

            log += "\n\tRewards:"

            if len(g_episode_rewards) > 0:
                log += " ".join([
                    " Global step mean/med rew:",
                    "{:.4f}/{:.4f},".format(
                        np.mean(per_step_g_rewards),
                        np.median(per_step_g_rewards)),
                    " Global eps mean/med/min/max eps rew:",
                    "{:.3f}/{:.3f}/{:.3f}/{:.3f},".format(
                        np.mean(g_episode_rewards),
                        np.median(g_episode_rewards),
                        np.min(g_episode_rewards),
                        np.max(g_episode_rewards))
                ])

            if args.eval:
                total_success = []
                total_spl = []
                total_dist = []
                for e in range(args.num_processes):
                    for acc in episode_success[e]:
                        total_success.append(acc)
                    for dist in episode_dist[e]:
                        total_dist.append(dist)
                    for spl in episode_spl[e]:
                        total_spl.append(spl)

                if len(total_spl) > 0:
                    log += " ObjectNav succ/spl/dtg:"
                    log += " {:.3f}/{:.3f}/{:.3f}({:.0f}),".format(
                        np.mean(total_success),
                        np.mean(total_spl),
                        np.mean(total_dist),
                        len(total_spl))
            else:
                if len(episode_success) > 100:
                    log += " ObjectNav succ/spl/dtg:"
                    log += " {:.3f}/{:.3f}/{:.3f}({:.0f}),".format(
                        np.mean(episode_success),
                        np.mean(episode_spl),
                        np.mean(episode_dist),
                        len(episode_spl))

            log += "\n\tLosses:"
            if len(g_value_losses) > 0 and not args.eval:
                log += " ".join([
                    " Policy Loss value/action/dist:",
                    "{:.3f}/{:.3f}/{:.3f},".format(
                        np.mean(g_value_losses),
                        np.mean(g_action_losses),
                        np.mean(g_dist_entropies))
                ])

            print(log)
            logging.info(log)
        # ------------------------------------------------------------------

        # ------------------------------------------------------------------
        # Save best models
        if (step * num_scenes) % args.save_interval < \
                num_scenes:
            if len(g_episode_rewards) >= 1000 and \
                    (np.mean(g_episode_rewards) >= best_g_reward) \
                    and not args.eval:
                torch.save(g_policy.state_dict(),
                           os.path.join(log_dir, "model_best.pth"))
                best_g_reward = np.mean(g_episode_rewards)

        # Save periodic models
        if (step * num_scenes) % args.save_periodic < \
                num_scenes:
            total_steps = step * num_scenes
            if not args.eval:
                torch.save(g_policy.state_dict(),
                           os.path.join(dump_dir,
                                        "periodic_{}.pth".format(total_steps)))
        # ------------------------------------------------------------------

    # Print and save model performance numbers during evaluation
    if args.eval:
        print("Dumping eval details...")
        
        total_success = []
        total_spl = []
        total_dist = []
        for e in range(args.num_processes):
            for acc in episode_success[e]:
                total_success.append(acc)
            for dist in episode_dist[e]:
                total_dist.append(dist)
            for spl in episode_spl[e]:
                total_spl.append(spl)

        if len(total_spl) > 0:
            log = "Final ObjectNav succ/spl/dtg:"
            log += " {:.3f}/{:.3f}/{:.3f}({:.0f}),".format(
                np.mean(total_success),
                np.mean(total_spl),
                np.mean(total_dist),
                len(total_spl))

        print(log)
        logging.info(log)
            
        # Save the spl per category
        log = "Success | SPL per category\n"
        for key in success_per_category:
            log += "{}: {} | {}\n".format(key,
                                          sum(success_per_category[key]) /
                                          len(success_per_category[key]),
                                          sum(spl_per_category[key]) /
                                          len(spl_per_category[key]))

        print(log)
        logging.info(log)

        with open('{}/{}_spl_per_cat_pred_thr.json'.format(
                dump_dir, args.split), 'w') as f:
            json.dump(spl_per_category, f)

        with open('{}/{}_success_per_cat_pred_thr.json'.format(
                dump_dir, args.split), 'w') as f:
            json.dump(success_per_category, f)
        
        with open('{}/{}_episode_data.json'.format(
                dump_dir, args.split), 'w') as f:
            json.dump(full_episode_data, f)


if __name__ == "__main__":
    main()

[(23, 'arguments.get_args', 'get_args', False, 'from arguments import get_args\n'), (25, 'numpy.random.seed', 'np.random.seed', True, 'import numpy as np\n'), (26, 'torch.manual_seed', 'torch.manual_seed', False, 'import torch\n'), (40, 'logging.basicConfig', 'logging.basicConfig', False, 'import logging\n'), (45, 'logging.info', 'logging.info', False, 'import logging\n'), (50, 'torch.device', 'torch.device', False, 'import torch\n'), (72, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (73, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (75, 'collections.deque', 'deque', False, 'from collections import deque, defaultdict\n'), (77, 'collections.deque', 'deque', False, 'from collections import deque, defaultdict\n'), (78, 'collections.deque', 'deque', False, 'from collections import deque, defaultdict\n'), (79, 'collections.deque', 'deque', False, 'from collections import deque, defaultdict\n'), (81, 'collections.deque', 'deque', False, 'from collections import deque, defaultdict\n'), (83, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (86, 'torch.set_num_threads', 'torch.set_num_threads', False, 'import torch\n'), (87, 'envs.make_vec_envs', 'make_vec_envs', False, 'from envs import make_vec_envs\n'), (102, 'torch.set_grad_enabled', 'torch.set_grad_enabled', False, 'import torch\n'), (129, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (137, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (242, 'gym.spaces.Box', 'gym.spaces.Box', False, 'import gym\n'), (299, 'time.time', 'time.time', False, 'import time\n'), (302, 'torch.set_grad_enabled', 'torch.set_grad_enabled', False, 'import torch\n'), (303, 'collections.defaultdict', 'defaultdict', False, 'from collections import deque, defaultdict\n'), (304, 'collections.defaultdict', 'defaultdict', False, 'from collections import deque, defaultdict\n'), (29, 'torch.cuda.manual_seed', 'torch.cuda.manual_seed', False, 'import torch\n'), (35, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (36, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (37, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (38, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (68, 'collections.deque', 'deque', False, 'from collections import deque, defaultdict\n'), (69, 'collections.deque', 'deque', False, 'from collections import deque, defaultdict\n'), (70, 'collections.deque', 'deque', False, 'from collections import deque, defaultdict\n'), (275, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (618, 'logging.info', 'logging.info', False, 'import logging\n'), (630, 'logging.info', 'logging.info', False, 'import logging\n'), (132, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (231, 'numpy.random.rand', 'np.random.rand', True, 'import numpy as np\n'), (248, 'model.Semantic_Mapping', 'Semantic_Mapping', False, 'from model import RL_Policy, Semantic_Mapping\n'), (251, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (457, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (506, 'time.time', 'time.time', False, 'import time\n'), (507, 'time.gmtime', 'time.gmtime', False, 'import time\n'), (570, 'logging.info', 'logging.info', False, 'import logging\n'), (634, 'json.dump', 'json.dump', False, 'import json\n'), (638, 'json.dump', 'json.dump', False, 'import json\n'), (642, 'json.dump', 'json.dump', False, 'import json\n'), (63, 'collections.deque', 'deque', False, 'from collections import deque, defaultdict\n'), (64, 'collections.deque', 'deque', False, 'from collections import deque, defaultdict\n'), (65, 'collections.deque', 'deque', False, 'from collections import deque, defaultdict\n'), (315, 'torch.FloatTensor', 'torch.FloatTensor', False, 'import torch\n'), (438, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (503, 'json.dump', 'json.dump', False, 'import json\n'), (582, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (612, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (613, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (614, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (52, 'torch.ones', 'torch.ones', False, 'import torch\n'), (120, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (121, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (125, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (126, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (578, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (581, 'os.path.join', 'os.path.join', False, 'import os\n'), (510, 'time.strftime', 'time.strftime', False, 'import time\n'), (546, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (547, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (548, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (554, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (555, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (556, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (218, 'torch.from_numpy', 'torch.from_numpy', False, 'import torch\n'), (450, 'torch.ones', 'torch.ones', False, 'import torch\n'), (521, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (522, 'numpy.median', 'np.median', True, 'import numpy as np\n'), (525, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (526, 'numpy.median', 'np.median', True, 'import numpy as np\n'), (527, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (528, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (564, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (565, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (566, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (192, 'torch.from_numpy', 'torch.from_numpy', False, 'import torch\n'), (397, 'torch.from_numpy', 'torch.from_numpy', False, 'import torch\n'), (418, 'torch.from_numpy', 'torch.from_numpy', False, 'import torch\n')]