from multiprocessing import Process, Pipe
import gym

def worker(conn, env):
    '''
    conn = <multiprocessing.connection.Connection object at 0x7f9aacbb5d68>
    env = <LTLEnv<ZonesEnv5<Zones-5-v0>>>
    '''
    while True:
        cmd, data = conn.recv()
        if cmd == "step":
            obs, reward, done, info = env.step(data)
            if done:
                obs = env.reset()
            conn.send((obs, reward, done, info))
        elif cmd == "reset":
            obs = env.reset()
            conn.send(obs)
        elif cmd == "kill":
            return
        else:
            raise NotImplementedError

class ParallelEnv(gym.Env):
    """A concurrent execution of environments in multiple processes."""
    def __init__(self, envs):
        assert len(envs) >= 1, "No environment given."

        self.envs = envs  ## {list:16}
        self.observation_space = self.envs[0].observation_space  # Dict(features:Box(76,))
        self.action_space = self.envs[0].action_space  # Box(2,)

        self.locals = []
        for env in self.envs[1:]:
            local, remote = Pipe()
            self.locals.append(local)
            p = Process(target=worker, args=(remote, env))
            p.daemon = True
            p.start()  # turn to def worker(conn, env):
            remote.close()

    def __del__(self):
        for local in self.locals:
            local.send(("kill", None))

    def reset(self):
        for local in self.locals:
            local.send(("reset", None))
        results = [self.envs[0].reset()] + [local.recv() for local in self.locals]
        return results

    def step(self, actions):
        for local, action in zip(self.locals, actions[1:]):
            local.send(("step", action))
        obs, reward, done, info = self.envs[0].step(actions[0])
        if done:
            obs = self.envs[0].reset()
        results = zip(*[(obs, reward, done, info)] + [local.recv() for local in self.locals])
        return results

    def render(self):
        raise NotImplementedError