def Market():
    env = raw_env()
    env = wrappers.AssertOutOfBoundsWrapper(env)
    return env

class raw_env(AECEnv):
    metadata = {
        "name": "Market",
    }

    def __init__(self):
        self.num_actions = 4
        self.agents = ["A", "B"]
        self.possible_agents = self.agents[:]
        self._it = 0
        self._max_it = 1000000
        self.observation_spaces = {agent: MultiDiscrete([self.num_actions, self.num_actions]) for agent in self.agents}
        # [Box(low=0,high=14,shape=(2,), dtype=int)]
        self.action_spaces = {agent: Discrete(self.num_actions) for agent in self.agents}
        self.prices = [(num+1) / (self.num_actions + 1) for num in range(self.num_actions)]
        self.act_A = 0
        self.act_B = 0

        
    def reset(self, seed=None, options=None):
        self.agents = self.possible_agents[:]
        self._agent_selector = agent_selector(self.agents)
        self.agent_selection = self._agent_selector.next()
        self.rewards = {agent: 0 for agent in self.agents}
        self._cumulative_rewards = {agent: 0 for agent in self.agents}
        self.terminations = {agent: False for agent in self.agents}
        self.truncations = {agent: False for agent in self.agents}
        self.infos = {agent: {} for agent in self.agents}
        self._it = 0
        self.observations = {agent: np.array([random.randint(0, self.num_actions - 1), random.randint(0, self.num_actions - 1)]) for agent in self.agents}
        self.act_A = self.observations['A'][0]
        self.act_B = self.observations['A'][1]
        return self.observations, self.infos


    def step(self, action):
  
        if (
            self.terminations[self.agent_selection]
            or self.truncations[self.agent_selection]
        ):
            self._was_dead_step(action)
            return
        
        agent = self.agent_selection
        
        if agent == 'A':
            self.act_A = action
        else:
            self.act_B = action
                    
        rewards = compute_rewards(np.array([self.prices[self.act_A], self.prices[self.act_B]]), 0.25)
        
        self.rewards['A'], self.rewards['B'] = rewards[0], rewards[1]
        
        for i in self.agents:
                self.observations[i] = np.array([self.act_A, self.act_B])
                
        self.agent_selection = self._agent_selector.next()
        self._accumulate_rewards()

        if self._it >= self._max_it : 
            self.terminations = {a: True for a in self.agents}
        else : 
            self.terminations = {a: False for a in self.agents}

        return self.observations, self.rewards, self.terminations, self.truncations, self.infos
        
    def render(self):
        pass

    def observation_space(self, agent):
        return self.observation_spaces[agent]

    def action_space(self, agent):
        return self.action_spaces[agent]
    
    def action_space(self, agent):
        return self.action_spaces[agent]    
    
    def observe(self, agent):
        return self.observations[agent]   
    
    def seed(self, seed):
        np.random.seed(seed)

def _get_env():
    """This function is needed to provide callables for DummyVectorEnv."""
    return PettingZooEnv(Market())

def _creat_DQN(env, learning_rate = 0.01) :
    observation_space = env.observation_space
    net = Net(
            state_shape=observation_space.shape or observation_space.n,
            action_shape=env.action_space.shape or env.action_space.n,
            hidden_sizes=[8,8],
            device="cuda" if torch.cuda.is_available() else "cpu",
        ).to("cuda" if torch.cuda.is_available() else "cpu")
    optim = torch.optim.Adam(net.parameters(), lr=learning_rate, betas = (0.9, 0.999))
    agent_learn = DQNPolicy(
        model=net,
        optim=optim,
        discount_factor=0.99,
        estimation_step=1,
        is_double = False,
        target_update_freq=0)
    return agent_learn

def _creat_PPO(env, learning_rate = 0.01) :
    net = Net(
            state_shape=env.observation_space.shape,
            action_shape=env.action_space.shape or env.action_space.n,
            hidden_sizes=[16,16],
            device="cuda" if torch.cuda.is_available() else "cpu",
        ).to("cuda" if torch.cuda.is_available() else "cpu")
    optim = torch.optim.Adam(net.parameters(), lr=learning_rate, betas = (0.9, 0.999))
    agent_learn = PPOPolicy(
        actor=net,
        critic=net,
        optim=optim,
        dist_fn = Multinomial,
        discount_factor=0.99)
    return agent_learn