from typing import Tuple, NoReturn, Sequence

import numpy as np
from numpy import ndarray

from akademy.models.base_models.replay_memory import ReplayMemory
from akademy.common.typedefs import State, Action


class ExperienceReplayMemory(ReplayMemory):
    """
    Implementation of the Experience Replay model as first described in the
    following paper:
        Lin, LJ. Self-improving reactive agents based on reinforcement learning,
        planning and teaching. Mach Learn 8, 293–321 (1992).
        https://doi.org/10.1007/BF00992699
    Args:
        state_shape: Tuple of a Tensor's dimensions e.g. (251, ) for 50 window
        memsize: total number of possible experiences to hold
    """
    def __init__(self, state_shape: Sequence[int], memsize: int = 10000):
        self.mem_count = 0
        self.memsize = memsize
        self.state_shape = state_shape

        # allocate numpy arrays as buffers to hold replay experiences
        self.states = np.zeros((self.memsize, *self.state_shape), dtype=np.float32)
        self.actions = np.zeros(self.memsize, dtype=np.int64)
        self.rewards = np.zeros(self.memsize, dtype=np.float32)
        self.next_states = np.zeros((self.memsize, *self.state_shape), dtype=np.float32)
        self.dones = np.zeros(self.memsize, dtype=bool)

    def add(
            self, state: State, action: Action,
            reward: float, next_state: State, done: bool
    ) -> NoReturn:
        """
        Add a replay experience as constituent parts to the appropriate buffers
        using the mem_index to treat them like ring buffers.
        Args:
            state: Object e.g. ndarray representing the current state of an environment.
            action: Object e.g. float | int representing action taken by Agent.
            reward: Value of the reward generated by the Agent's action.
            next_state: Object e.g. ndarray representing the resulting state of an environment
                after an Agent's action is applied.
            done: boolean indicating if this was a terminal state.
        """
        # treat buffers like FIFO stacks
        mem_index = self.mem_count % self.memsize

        # adds each part of the experience into the appropriate index of
        # the corresponding ndarrays.
        self.states[mem_index] = state
        self.actions[mem_index] = action
        self.rewards[mem_index] = reward
        self.next_states[mem_index] = next_state
        self.dones[mem_index] = 1 - done  # flip the bit here

        # keeps a measure of each addition such that modulo operations
        # can be used to effectively determine the last-added reference
        # such that the memory acts as a FIFO queue.
        self.mem_count += 1

    def sample(
            self, batch_size: int = 64
    ) -> Tuple[ndarray, ndarray, ndarray, ndarray, ndarray]:
        """
        Gets a list of <batch_size> many random indices from within
        a range of [0, min(memcount, memsize)] and returns the elements from
        each ring buffer corresponding to those as a tuple as:

        Args:
            batch_size: number of total random samples to return.

        Returns:
            Tuple of ndarrays as [state, action, reward, next-states, dones]
        """
        # control bounds of sample size to ensure
        #   A.) no more than max amount sampled
        #   B.) no more tha min available sampled.
        _mem_max = min(self.mem_count, self.memsize)

        # randomly select indices to be sampled
        batch_indices = np.random.choice(_mem_max, batch_size, replace=True)

        # create sub-array from each specified index, from each component.
        states = self.states[batch_indices]
        actions = self.actions[batch_indices]
        rewards = self.rewards[batch_indices]
        next_states = self.next_states[batch_indices]
        dones = self.dones[batch_indices]

        return states, actions, rewards, next_states, dones

    def __len__(self) -> int:
        """
        Length is reported as the minimum of the count of added experiences and
        the maximum number of experiences such that the total possible range
        of length values is [0, memsize].
        Returns:
            integer value representing the current capacity of the
            ExperienceReplay object in the range of (0, memsize]
        """
        return min(self.mem_count, self.memsize)
