import torch
import numpy as np


class Context:

    def __init__(self,logger,coord_logger,save_dir,path_models,args):
        self.logger = logger
        self.coord_logger=coord_logger
        self.load_model_path=args.load_model_path
        self.save_model=args.save_model
        self.path_models=path_models
        self.save_dir=save_dir

        if args.env_name == "sawyer_door":self.object = np.load(open("resources/door_goals.npy", "rb"), allow_pickle=True).item()
        if args.env_name == "sawyer_pickup": self.object = np.load(open("resources/pickup_goals.npy", "rb"), allow_pickle=True).item()
        if args.env_type == "maze" and args.env_name == "PointMaze1-v1" or args.env_name == "PointMaze1Pos-v1":
            self.loaded_obs = torch.load("resources/obs.pt")
            self.loaded_pos = torch.load("resources/pos.pt")
        if args.env_type == "maze" and args.env_name == "AntMaze1-v1" or args.env_name == "AntMaze1Pos-v1":
            self.loaded_obs = torch.load("resources/obs_antmaze.pt")
            self.loaded_pos = torch.load("resources/pos_antmaze.pt")

        if args.env_type == "maze" and args.env_name == "PointPush-v1" or args.env_name == "PointPushPos-v1":
            self.loaded_obs = torch.load("resources/obs_push.pt")
            self.loaded_pos = torch.load("resources/pos_push.pt")
        if args.env_type == "maze" and args.env_name == "AntPush-v1" or args.env_name == "AntPushPos-v1":
            self.loaded_obs = torch.load("resources/obs_antpush.pt")
            self.loaded_pos = torch.load("resources/pos_antpush.pt")
        if args.env_type == "maze" and args.env_name == "AntMaze-v1" or args.env_name == "AntMazePos-v1":
            self.loaded_obs = torch.load("resources/obs_ant4room.pt")
            self.loaded_pos = torch.load("resources/pos_ant4room.pt")
