from ..value_function import ValueFunction
from ...utils import check_random_state

_MEAN = 0.0
_STD = 0.3


class TabularVF(ValueFunction):
    '''
    Tabular Value Function.

    Can be used for both state V(s) and state-action Q(s, a) values.

    # Arguments

        init: boolean. Wheter to return a randomly initialized value
            when accessed and does not exist in the table.

    # Examples

        ```
        v = TabularVF()
        v[state] = 3.2
        print(q[state]) # 3.2

        q = TabularVF()
        q[state, action] = 1.8
        print(q[state, action]) # 1.8
        ```
    '''

    def __init__(self, init=True, random_state=None):
        self.init = init
        self.random_state = check_random_state(random_state)
        self._d = {}

    def __setitem__(self, key, value):
        '''
        Sets the state or state-action value.

        # Arguments
            key: `state` or `(state, action)`
            value: a scalar
        '''
        self._d[key] = value

    #################
    # ValueFunction #
    #################

    def __getitem__(self, key):
        '''
        Returns the state or state-action value.

        # Arguments
            key: `state` or `(state, action)`
        '''
        if key not in self._d:
            if self.init:
                self._d[key] = self.random_state.normal(_MEAN, _STD)
            else:
                self._d[key] = 0
        return self._d[key]

@@ -257,17 +257,16 @@
 guments%0A
-%0A
         
@@ -399,18 +399,16 @@
  Example
-s%0A
 %0A       
@@ -411,24 +411,30 @@
      %60%60%60
+python
 %0A       
  v = Tab
@@ -425,16 +425,20 @@
 %0A       
+ %3E%3E%3E
  v = Tab
@@ -453,16 +453,20 @@
 %0A       
+ %3E%3E%3E
  v%5Bstate
@@ -472,32 +472,36 @@
 e%5D = 3.2%0A       
+ %3E%3E%3E
  print(q%5Bstate%5D)
@@ -500,18 +500,24 @@
 %5Bstate%5D)
+%0A      
  
-#
  3.2%0A%0A  
@@ -521,16 +521,20 @@
 %0A       
+ %3E%3E%3E
  q = Tab
@@ -549,16 +549,20 @@
 %0A       
+ %3E%3E%3E
  q%5Bstate
@@ -584,16 +584,20 @@
 %0A       
+ %3E%3E%3E
  print(q
@@ -612,18 +612,24 @@
 action%5D)
+%0A      
  
-#
  1.8%0A   

update docs