import multiprocessing as mp
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.nn.parallel import DistributedDataParallelCPU as DDPC

from rlpyt.utils.quick_args import save__init__args
from rlpyt.utils.collections import namedarraytuple
from rlpyt.utils.synchronize import RWLock
from rlpyt.utils.logging import logger
from rlpyt.models.utils import strip_ddp_state_dict

AgentInputs = namedarraytuple("AgentInputs",
    ["observation", "prev_action", "prev_reward"])
AgentStep = namedarraytuple("AgentStep", ["action", "agent_info"])


class BaseAgent:
    """
    The agent performs many functions, including: action-selection during
    sampling, returning policy-related values to use in training (e.g. action
    probabilities), storing recurrent state during sampling, managing model
    device, and performing model parameter communication between processes.
    The agent is both interfaces: sampler<-->neural network<-->algorithm.
    Typically, each algorithm and environment combination will require at
    least some of its own agent functionality.

    The base agent automatically carries out some of these roles.  It assumes
    there is one neural network model.  Agents using multiple models might
    need to extend certain funcionality to include those models, depending on
    how they are used.
    """
    
    recurrent = False
    alternating = False

    def __init__(self, ModelCls=None, model_kwargs=None, initial_model_state_dict=None):
        """
        Arguments are saved but no model initialization occurs.

        Args:
            ModelCls: The model class to be used.
            model_kwargs (optional): Any keyword arguments to pass when instantiating the model.
            initial_model_state_dict (optional): Initial model parameter values.
        """

        save__init__args(locals())
        self.model = None  # type: torch.nn.Module
        self.shared_model = None
        self.distribution = None
        self.device = torch.device("cpu")
        self._mode = None
        if self.model_kwargs is None:
            self.model_kwargs = dict()
        # The rest only for async operations:
        self._rw_lock = RWLock()
        self._send_count = mp.RawValue("l", 0)
        self._recv_count = 0

    def __call__(self, observation, prev_action, prev_reward):
        """Returns values from model forward pass on training data (i.e. used
        in algorithm)."""
        raise NotImplementedError

    def initialize(self, env_spaces, share_memory=False, **kwargs):
        """
        Instantiates the neural net model(s) according to the environment
        interfaces.  

        Uses shared memory as needed--e.g. in CpuSampler, workers have a copy
        of the agent for action-selection.  The workers automatically hold
        up-to-date parameters in ``model``, because they exist in shared
        memory, constructed here before worker processes fork. Agents with
        additional model components (beyond ``self.model``) for
        action-selection should extend this method to share those, as well.

        Typically called in the sampler during startup.

        Args:
            env_spaces: passed to ``make_env_to_model_kwargs()``, typically namedtuple of 'observation' and 'action'.
            share_memory (bool): whether to use shared memory for model parameters.
        """
        self.env_model_kwargs = self.make_env_to_model_kwargs(env_spaces)
        self.model = self.ModelCls(**self.env_model_kwargs,
            **self.model_kwargs)
        if share_memory:
            self.model.share_memory()
            # Store the shared_model (CPU) under a separate name, in case the
            # model gets moved to GPU later:
            self.shared_model = self.model
        if self.initial_model_state_dict is not None:
            self.model.load_state_dict(self.initial_model_state_dict)
        self.env_spaces = env_spaces
        self.share_memory = share_memory

    def make_env_to_model_kwargs(self, env_spaces):
        """Generate any keyword args to the model which depend on environment interfaces."""
        return {}

    def to_device(self, cuda_idx=None):
        """Moves the model to the specified cuda device, if not ``None``.  If
        sharing memory, instantiates a new model to preserve the shared (CPU)
        model.  Agents with additional model components (beyond
        ``self.model``) for action-selection or for use during training should
        extend this method to move those to the device, as well.

        Typically called in the runner during startup.
        """
        if cuda_idx is None:
            return
        if self.shared_model is not None:
            self.model = self.ModelCls(**self.env_model_kwargs,
                **self.model_kwargs)
            self.model.load_state_dict(self.shared_model.state_dict())
        self.device = torch.device("cuda", index=cuda_idx)
        self.model.to(self.device)
        logger.log(f"Initialized agent model on device: {self.device}.")

    def data_parallel(self):
        """Wraps the model with PyTorch's DistributedDataParallel.  The
        intention is for rlpyt to create a separate Python process to drive
        each GPU (or CPU-group for CPU-only, MPI-like configuration). Agents
        with additional model components (beyond ``self.model``) which will
        have gradients computed through them should extend this method to wrap
        those, as well.

        Typically called in the runner during startup.
        """
        if self.device.type == "cpu":
            self.model = DDPC(self.model)
            logger.log("Initialized DistributedDataParallelCPU agent model.")
        else:
            self.model = DDP(self.model,
                device_ids=[self.device.index], output_device=self.device.index)
            logger.log("Initialized DistributedDataParallel agent model on "
                f"device {self.device}.")

    def async_cpu(self, share_memory=True):
        """Used in async runner only; creates a new model instance to be used
        in the sampler, separate from the model shared with the optimizer
        process.  The sampler can operate asynchronously, and choose when to
        copy the optimizer's (shared) model parameters into its model (under
        read-write lock). The sampler model may be stored in shared memory,
        as well, to instantly share values with sampler workers.  Agents with
        additional model components (beyond ``self.model``) should extend this
        method to do the same with those, if using in asynchronous mode.

        Typically called in the runner during startup.

        TODO: double-check wording if this happens in sampler and optimizer."""
        if self.device.type != "cpu":
            return
        assert self.shared_model is not None
        self.model = self.ModelCls(**self.env_model_kwargs,
            **self.model_kwargs)
        # TODO: might need strip_ddp_state_dict.
        self.model.load_state_dict(self.shared_model.state_dict())
        if share_memory:  # Not needed in async_serial.
            self.model.share_memory()  # For CPU workers in async_cpu.
        logger.log("Initialized async CPU agent model.")

    def collector_initialize(self, global_B=1, env_ranks=None):
        """If needed to initialize within CPU sampler (e.g. vector epsilon-greedy,
        see EpsilonGreedyAgent for details)."""
        pass

    @torch.no_grad()  # Hint: apply this decorator on overriding method.
    def step(self, observation, prev_action, prev_reward):
        """Returns selected actions for environment instances in sampler."""
        raise NotImplementedError  # return type: AgentStep

    def reset(self):
        pass

    def reset_one(self, idx):
        pass

    def parameters(self):
        """Parameters to be optimized (overwrite in subclass if multiple models)."""
        return self.model.parameters()

    def state_dict(self):
        """Returns model parameters for saving."""
        return self.model.state_dict()

    def load_state_dict(self, state_dict):
        """Load model parameters, should expect format returned from ``state_dict()``."""
        self.model.load_state_dict(state_dict)

    def train_mode(self, itr):
        """Go into training mode (e.g. see PyTorch's ``Module.train()``)."""
        self.model.train()
        self._mode = "train"

    def sample_mode(self, itr):
        """Go into sampling mode."""
        self.model.eval()
        self._mode = "sample"

    def eval_mode(self, itr):
        """Go into evaluation mode.  Example use could be to adjust epsilon-greedy."""
        self.model.eval()
        self._mode = "eval"

    def sync_shared_memory(self):
        """Copies model parameters into shared_model, e.g. to make new values
        available to sampler workers.  If running CPU-only, these will be the
        same object--no copy necessary.  If model is on GPU, copy to CPU is
        performed. (Requires ``initialize(share_memory=True)`` called
        previously.  Not used in async mode.

        Typically called in the XXX during YY.
        """
        if self.shared_model is not self.model:  # (self.model gets trained)
            self.shared_model.load_state_dict(strip_ddp_state_dict(
                self.model.state_dict()))

    def send_shared_memory(self):
        """Used in async mode only, in optimizer process; copies parameters
        from trained model (maybe GPU) to shared model, which the sampler can
        access. Does so under write-lock, and increments send-count which sampler
        can check.

        Typically called in the XXX during YY."""
        if self.shared_model is not self.model:
            with self._rw_lock.write_lock:
                self.shared_model.load_state_dict(
                    strip_ddp_state_dict(self.model.state_dict()))
                self._send_count.value += 1

    def recv_shared_memory(self):
        """Used in async mode, in sampler process; copies parameters from
        model shared with optimizer into local model, if shared model has been
        updated.  Does so under read-lock.  (Local model might also be shared
        with sampler workers).

        Typically called in the XXX during YY."""
        if self.shared_model is not self.model:
            with self._rw_lock:
                if self._recv_count < self._send_count.value:
                    self.model.load_state_dict(self.shared_model.state_dict())
                    self._recv_count = self._send_count.value

    def toggle_alt(self):
        pass  # Only needed for recurrent alternating agent, but might get called.


AgentInputsRnn = namedarraytuple("AgentInputsRnn",  # Training only.
    ["observation", "prev_action", "prev_reward", "init_rnn_state"])


class RecurrentAgentMixin:
    """
    Mixin class to manage recurrent state during sampling (so the sampler
    remains agnostic).  To be used like ``class
    MyRecurrentAgent(RecurrentAgentMixin, MyAgent):``.
    """
    recurrent = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._prev_rnn_state = None
        self._sample_rnn_state = None  # Store during eval.

    def reset(self):
        """Sets the recurrent state to ``None``, which built-in PyTorch
        modules conver to zeros."""
        self._prev_rnn_state = None

    def reset_one(self, idx):
        """Sets the recurrent state corresponding to one environment instance
        to zero.  Assumes rnn state is in cudnn-compatible shape: [N,B,H],
        where B corresponds to environment index."""
        if self._prev_rnn_state is not None:
            self._prev_rnn_state[:, idx] = 0  # Automatic recursion in namedarraytuple.

    def advance_rnn_state(self, new_rnn_state):
        """Sets the recurrent state to the newly computed one (i.e. recurrent agents should
        call this at the end of their ``step()``). """
        self._prev_rnn_state = new_rnn_state

    @property
    def prev_rnn_state(self):
        return self._prev_rnn_state

    def train_mode(self, itr):
        """If coming from sample mode, store the rnn state elsewhere and clear it."""
        if self._mode == "sample":
            self._sample_rnn_state = self._prev_rnn_state
        self._prev_rnn_state = None
        super().train_mode(itr)

    def sample_mode(self, itr):
        """If coming from non-sample modes, restore the last sample-mode rnn state."""
        if self._mode != "sample":
            self._prev_rnn_state = self._sample_rnn_state
        super().sample_mode(itr)

    def eval_mode(self, itr):
        """If coming from sample mode, store the rnn state elsewhere and clear it."""
        if self._mode == "sample":
            self._sample_rnn_state = self._prev_rnn_state
        self._prev_rnn_state = None
        super().eval_mode(itr)


class AlternatingRecurrentAgentMixin:
    """
    Maintain an alternating pair of recurrent states to use when stepping in
    the sampler. Automatically swap them out when ``advance_rnn_state()`` is
    called, so it otherwise behaves like regular recurrent agent.  Should use
    only in alternating samplers, where two sets of environment instances take
    turns stepping (no special class needed for feedforward agents).  Use in
    place of ``RecurrentAgentMixin``.
    """

    recurrent = True
    alternating = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._alt = 0
        self._prev_rnn_state = None
        self._prev_rnn_state_pair = [None, None]
        self._sample_rnn_state_pair = [None, None]

    def reset(self):
        self._prev_rnn_state_pair = [None, None]
        self._prev_rnn_state = None
        self._alt = 0
        # Leave _sample_rnn_state_pair alone.

    def advance_rnn_state(self, new_rnn_state):
        """To be called inside agent.step()."""
        self._prev_rnn_state_pair[self._alt] = new_rnn_state
        self._alt ^= 1
        self._prev_rnn_state = self._prev_rnn_state_pair[self._alt]

    @property
    def prev_rnn_state(self):
        return self._prev_rnn_state

    def train_mode(self, itr):
        if self._mode == "sample":
            self._sample_rnn_state_pair = self._prev_rnn_state_pair
        self._prev_rnn_state_pair = [None, None]
        self._prev_rnn_state = None
        self._alt = 0
        super().train_mode(itr)

    def sample_mode(self, itr):
        if self._mode != "sample":
            self._prev_rnn_state_pair = self._sample_rnn_state_pair
            self._alt = 0
            self._prev_rnn_state = self._prev_rnn_state_pair[0]
        super().sample_mode(itr)

    def eval_mode(self, itr):
        if self._mode == "sample":
            self._sample_rnn_state_pair = self._prev_rnn_state_pair
        self._prev_rnn_state_pair = [None, None]
        self._prev_rnn_state = None
        self._alt = 0
        super().eval_mode(itr)

    def get_alt(self):
        return self._alt

    def toggle_alt(self):
        self._alt ^= 1
        self._prev_rnn_state = self._prev_rnn_state_pair[self._alt]
