'''
Implements wrappers on top of the basic messenger environments
'''
import random

from messenger.envs.base import MessengerEnv
from messenger.envs.stage_one import StageOne
from messenger.envs.stage_two import StageTwo


class TwoEnvWrapper(MessengerEnv):
    '''
    Switches between two Messenger environments
    '''
    def __init__(self, stage:int, split_1:str, split_2:str, prob_env_1=0.5, **kwargs):
        if stage == 1:
            self.env_1 = StageOne(split=split_1)
            self.env_2 = StageOne(split=split_2)
        elif stage == 2:
            self.env_1 = StageTwo(split=split_1)
            self.env_2 = StageTwo(split=split_2)
        
        self.prob_env_1 = prob_env_1
        self.cur_env = None
        
    def reset(self, **kwargs):
        if random.random() < self.prob_env_1:
            self.cur_env = self.env_1
        else:
            self.cur_env = self.env_2
        return self.cur_env.reset(**kwargs)
    
    def step(self, action):
        return self.cur_env.step(action)