from typing import Optional
from mlagents_envs.communicator_objects.unity_output_pb2 import UnityOutputProto
from mlagents_envs.communicator_objects.unity_input_pb2 import UnityInputProto


class Communicator(object):
    def __init__(self, worker_id=0, base_port=5005):
        """
        Python side of the communication. Must be used in pair with the right Unity Communicator equivalent.

        :int worker_id: Offset from base_port. Used for training multiple environments simultaneously.
        :int base_port: Baseline port number to connect to Unity environment over. worker_id increments over this.
        """

    def initialize(self, inputs: UnityInputProto) -> UnityOutputProto:
        """
        Used to exchange initialization parameters between Python and the Environment
        :param inputs: The initialization input that will be sent to the environment.
        :return: UnityOutput: The initialization output sent by Unity
        """

    def exchange(self, inputs: UnityInputProto) -> Optional[UnityOutputProto]:
        """
        Used to send an input and receive an output from the Environment
        :param inputs: The UnityInput that needs to be sent the Environment
        :return: The UnityOutputs generated by the Environment
        """

    def close(self):
        """
        Sends a shutdown signal to the unity environment, and closes the connection.
        """