"""Markov Decision Process Implementations."""

from SafeRLBench import EnvironmentBase
from SafeRLBench.spaces import DiscreteSpace


class MDP(EnvironmentBase):
    """Discrete Markov Decision Process Environment.

    Attributes
    ----------

    """

    def __init__(self, transitions, rewards):
        """MDP initialization.

        Parameters
        ----------
        transitions : array
            Array holding transition matrix for each action. The dimension of
            the state and action spaces will be deduced from this array.
        rewards : array
            Array holding the reward matrix for each action. It needs to comply
            with the dimensions deduced from the transitions array.
        """
        self.transitions = transitions
        self.rewards = rewards

        # determine state and action space
        self.action_space = DiscreteSpace(len(transitions))
        self.state_space = DiscreteSpace(transitions[0].shape[0])

    def _update(self, action):
        pass

    def _reset(self):
        pass

@@ -41,16 +41,36 @@
 ns.%22%22%22%0A%0A
+import numpy as np%0A%0A
 from Saf
@@ -148,16 +148,45 @@
 Space%0A%0A%0A
+# TODO: Implementation, Docs%0A
 class MD
@@ -288,16 +288,119 @@
 --------
+%0A    transitions :%0A%0A    rewards :%0A%0A    action_space :%0A%0A    state_space :%0A%0A    init_state :%0A%0A    state :
 %0A%0A    %22%22
@@ -445,16 +445,44 @@
  rewards
+, init_state=None, seed=None
 ):%0A     
@@ -1158,47 +1158,520 @@
     
-def _update(self, action):%0A        pass
+    # setup current state and store the initial state for reset%0A        self.init_state = init_state%0A        self.state = init_state%0A%0A    def _update(self, action):%0A        prev_state = self.state%0A        # choose next state%0A        self.state = np.random.choice(np.arange(self.action_space.dimension),%0A                                      p=self.transitions%5Baction%5D%5Bself.state%5D)%0A        # determine reward%0A        reward = self.rewards%5Baction%5D%5Bprev_state, self.state%5D%0A%0A        return action, self.state, reward
 %0A%0A  
@@ -1702,9 +1702,33 @@
     
-pass
+self.state = self.init_state
 %0A

MDP. WIP