import numpy as np

class QTableFactory:
    @staticmethod
    def create_obs_qtable(num_actions):
        return ObsQTable(num_actions)


class ObsQTable:
    def __init__(self, num_actions=7):
        self.qtable = {}
        self.num_actions = num_actions


    def set_value(self, state, action, value):
        if state not in self.qtable:
            self.qtable[state] = np.zeros(self.num_actions)

        self.qtable[state][action] = value

    def get_value(self, state, action):
        if state in self.qtable:
            if action is None:
                return self.qtable[state]
            elif action == -1:
                return 0.0
            else:
                return self.qtable[state][action]
        else:
            if action is None:
                return np.zeros(self.num_actions)
            else:
                return 0.0


    def get_qtable(self):
        return self.qtable
