def console(verbosity_level, module, msg):
    if verbosity_level > VERBOSITY:
        return
    else:
        print(module + ' - ' + msg) CFG_DIR = 'config'
ROCK_CFG_FILE = 'rock_sample_config.json'
LOG_DIR = 'log'
LOG_FILE = 'POMDPy.log'
def parse_map(m): class GridPosition(object):
    """
    Supported Constructors - GripPosition(i,j) or GridPosition() - defaults to 0,0
    """
    def __init__(self, i=0, j=0):
        if i is None and j is None:
            self.i = 0
            self.j = 0
        else:
            self.i = i
            self.j = j

    def __eq__(self, other):
        return self.i == other.i and self.j == other.j

    def print_position(self):
        print('(', end=' ')
        print(self.i, end=' ')
        print(',', end=' ')
        print(self.j, end=' ')
        print(')')

    def to_string(self):
        return '(' + str(self.i) + ',' + str(self.j) + ')'

    def copy(self):
        return GridPosition(self.i, self.j)

    def as_list(self):
        return [self.i, self.j]

    def manhattan_distance(self, other_position):
        return np.linalg.norm(np.subtract(self.as_list(), other_position.as_list()), 1)

    def euclidean_distance(self, other_position):
        return np.sqrt(np.power(other_position.j - self.j, 2.0) + np.power(other_position.i - self.i, 2.0)) class RockState(DiscreteState):
    """
    The state contains the position of the robot, as well as a boolean value for each rock
    representing whether it is good (true => good, false => bad).

    This class also implements DiscretizedState in order to allow the state to be easily
    converted to a List
    """

    def __init__(self, grid_position, rock_states):
        if rock_states is not None:
            assert rock_states.__len__() is not 0
        self.position = grid_position
        self.rock_states = rock_states  # list

    def distance_to(self, other_rock_state):
        """
        Distance is measured between beliefs by the sum of the num of different rocks
        """
        assert isinstance(other_rock_state, RockState)
        distance = 0
        # distance = self.position.manhattan_distance(other_rock_state.position)
        for i, j in zip(self.rock_states, other_rock_state.rock_states):
            if i != j:
                distance += 1
        return distance

    def __eq__(self, other_rock_state):
        return self.position == other_rock_state.position and self.rock_states is other_rock_state.rock_states

    def copy(self):
        return RockState(self.position, self.rock_states)

    def __hash__(self):
        """
        Returns a decimal value representing the binary state string
        :return:
        """
        return int(self.to_string(), 2)

    def to_string(self):
        state_string = self.position.to_string()
        state_string += " - "

        for i in self.rock_states:
            if i:
                state_string += "1 "
            else:
                state_string += "0 "
        return state_string

    def print_state(self):
        """
        Pretty printing
        :return:
        """
        self.position.print_position()

        print('Good: {', end=' ')
        good_rocks = []
        bad_rocks = []
        for i in range(0, self.rock_states.__len__()):
            if self.rock_states[i]:
                good_rocks.append(i)
            else:
                bad_rocks.append(i)
        for j in good_rocks:
            print(j, end=' ')
        print('}; Bad: {', end=' ')
        for k in bad_rocks:
            print(k, end=' ')
        print('}')

    def as_list(self):
        """
        Returns a list containing the (i,j) grid position boolean values
        representing the boolean rock states (good, bad)
        :return:
        """
        state_list = [self.position.i, self.position.j]
        for i in range(0, self.rock_states.__len__()):
            if self.rock_states[i]:
                state_list.append(True)
            else:
                state_list.append(False)
        return state_list

    def separate_rocks(self):
        """
        Used for the PyGame sim
        :return:
        """
        good_rocks = []
        bad_rocks = []
        for i in range(0, self.rock_states.__len__()):
            if self.rock_states[i]:
                good_rocks.append(i)
            else:
                bad_rocks.append(i)
        return good_rocks, bad_rocks class RockAction(DiscreteAction):
    """
    -The Rock sample problem Action class
    -Wrapper for storing the bin number. Also stores the rock number for checking actions
    -Handles pretty printing
    """

    def __init__(self, bin_number):
        super(RockAction, self).__init__(bin_number)
        if self.bin_number >= ActionType.CHECK:
            self.rock_no = self.bin_number - ActionType.CHECK
        else:
            self.rock_no = 0

    def copy(self):
        return RockAction(self.bin_number)

    def print_action(self):
        if self.bin_number >= ActionType.CHECK:
            print("CHECK")
        elif self.bin_number is ActionType.NORTH:
            print("NORTH")
        elif self.bin_number is ActionType.EAST:
            print("EAST")
        elif self.bin_number is ActionType.SOUTH:
            print("SOUTH")
        elif self.bin_number is ActionType.WEST:
            print("WEST")
        elif self.bin_number is ActionType.SAMPLE:
            print("SAMPLE")
        else:
            print("UNDEFINED ACTION")

    def to_string(self):
        if self.bin_number >= ActionType.CHECK:
            action = "CHECK"
        elif self.bin_number is ActionType.NORTH:
            action = "NORTH"
        elif self.bin_number is ActionType.EAST:
            action = "EAST"
        elif self.bin_number is ActionType.SOUTH:
            action = "SOUTH"
        elif self.bin_number is ActionType.WEST:
            action = "WEST"
        elif self.bin_number is ActionType.SAMPLE:
            action = "SAMPLE"
        else:
            action = "UNDEFINED ACTION"
        return action

    def distance_to(self, other_point):
        pass class ActionType(object):
    """
    Lists the possible actions and attributes an integer code to each for the Rock sample problem
    """
    NORTH = 0
    EAST = 1
    SOUTH = 2
    WEST = 3
    SAMPLE = 4
    CHECK = 5 class RockObservation(DiscreteObservation):
    """
    Default behavior is for the rock observation to say that the rock is empty
    """
    def __init__(self, is_good=False, is_empty=None):
        super(RockObservation, self).__init__(0 if is_empty else (2, 1)[is_good])
        self.is_empty = (True, is_empty)[is_empty is not None]
        self.is_good = is_good

    def distance_to(self, other_rock_observation):
        return abs(self.is_good - other_rock_observation.is_good)

    def copy(self):
        return RockObservation(self.is_good, self.is_empty)

    def __eq__(self, other_rock_observation):
        return self.is_good == other_rock_observation.is_good

    def __hash__(self):
        return (False, True)[self.is_good]

    def print_observation(self):
        if self.is_empty:
            print("EMPTY")
        elif self.is_good == 1:
            print("Good")
        elif self.is_good == 2:
            print("Bad")
        else:
            print(self.is_good)

    def to_string(self):
        if self.is_empty:
            obs = "EMPTY"
        elif self.is_good == 1:
            obs = "Good"
        elif self.is_good == 2:
            obs = "Bad"
        else:
            obs = str(self.is_good)
        return obs class DiscreteActionPool(ActionPool):
    """
     * An abstract implementation of the ActionPool interface that considers actions in terms of
     * discrete bins.
     *
    """
    def __init__(self, model):
        """
        :param model:
        """
        self.all_actions = model.get_all_actions()

    def create_action_mapping(self, belief_node):
        return DiscreteActionMapping(belief_node, self, self.create_bin_sequence(belief_node))

    def sample_an_action(self, bin_number):
        return self.all_actions[bin_number]

    def sample_random_action(self):
        return np.random.choice(self.all_actions)

    @staticmethod
    def create_bin_sequence(belief_node):
        """
        Default behavior is to make available only the legal actions for each action node
        :param belief_node:
        :return:
        """
        return belief_node.data.legal_actions() class DiscreteObservationPool(ObservationPool):
    """
    An implementation of the ObservationPool interface that is based on a discrete observation
    space.

    All of the information is stored inside the individual mapping classes, so this class serves as
    a simple factory for instances of DiscreteObservationMap.
    """

    def __init__(self, agent):
        self.agent = agent

    def create_observation_mapping(self, action_node):
        return DiscreteObservationMap(action_node, self.agent) class Model(with_metaclass(abc.ABCMeta, object)):
    """
    * Defines the Model class, which is an abstract class representing a POMDP model to be solved by
    * the solver. The core of the Model class is represented by a "black box" generative model which
    * is used to generate new steps in a simulation.

    * The key methods are:
    * - sampleAnInitState() - samples a new state from the initial belief.
    * - isTerminal() - returns true iff the given state is terminal.
    * - generateStep() - the key method representing the generative model; basically, it does
    *          (s, a) => (o, r, s')
    * - createActionPool() - this defines the way in which actions are mapped out inside the policy
    *      tree
    * - generate_particles() - this uses previous state particles to create a new set of derived particles
    *           by utilizing the step generator
    """

    def __init__(self, args):
        for k in args:
            setattr(self, k, args[k])
        pp(args)

        my_dir = os.path.dirname(__file__)
        self.weight_dir = os.path.join(my_dir, '..', '..', 'experiments', 'pickle_jar')
        self.ckpt_dir = os.path.join(my_dir, '..', '..', 'experiments', 'checkpoints')
        self.logs = os.path.join(my_dir, '..', '..', 'experiments', 'tensorboard')

        # TODO: More elegant naming convention for experiments
        if not os.path.exists(self.logs):
            os.makedirs(self.logs)

        count = len(os.listdir(self.logs))
        self.logs = os.path.join(self.logs, str(count))

        if not os.path.exists(self.weight_dir):
            os.makedirs(self.weight_dir)
        if not os.path.exists(self.ckpt_dir):
            os.makedirs(self.ckpt_dir)
        if not os.path.exists(self.logs):
            os.makedirs(self.logs)

    @abc.abstractmethod
    def reset_for_simulation(self):
        """
        The Simulator (Model) should be reset before each simulation
        :return:
        """

    @abc.abstractmethod
    def reset_for_epoch(self):
        """
        Defines behavior for resetting the simulator before each epoch
        :return:
        """

    @abc.abstractmethod
    def update(self, sim_data):
        """
        Update the state of the simulator with sim_data
        :param sim_data:
        :return:
        """

    @abc.abstractmethod
    def generate_step(self, state, action):
        """
        Generates a full StepResult, including the next state, an observation, and the reward
        *
        * For convenience, the action taken is also included in the result, as well as a flag for
        * whether or not the resulting next state is terminal.
        :param state:
        :param action:
        :return: StepResult
        """

    @abc.abstractmethod
    def sample_an_init_state(self):
        """
        Samples an initial state from the initial belief.
        :return: State
        """

    @abc.abstractmethod
    def sample_state_uninformed(self):
        """
        Samples a state from a poorly-informed prior. This is used by the provided default
        implementation of the second generateParticles() method.
        :return:
        """

    @abc.abstractmethod
    def sample_state_informed(self, belief):
        """
        :param belief:
        :return:
        """

    @staticmethod
    def get_initial_belief_state():
        """
        Return an np.array of initial belief probabilities for each state
        :return:
        """
        pass

    @abc.abstractmethod
    def belief_update(self, old_belief, action, observation):
        """
        Use bayes filter to update belief distribution
        :param old_belief:
        :param action
        :param observation
        :return:
        """

    @abc.abstractmethod
    def get_all_states(self):
        """
        :return: list of enumerated states (discrete) or range of states (continuous)
        """

    @staticmethod
    def get_transition_matrix():
        """
        Transition probability matrix, for value iteration
        :return:
        """
        pass

    @abc.abstractmethod
    def get_all_actions(self):
        """
        :return: list of enumerated actions (discrete) or range of actions (continuous)
        """

    @abc.abstractmethod
    def get_all_observations(self):
        """
        :return: list of enumerated observations (discrete) or range of observations (continuous)
        """

    @staticmethod
    def get_observation_matrix():
        """
        Observation probability matrix
        :return:
        """
        pass

    @staticmethod
    def get_reward_matrix():
        """
        Return reward matrix
        :return:
        """
        pass

    @abc.abstractmethod
    def get_legal_actions(self, state):
        """
        Given the current state of the system, return all legal actions
        :return: list of legal actions
        """

    @abc.abstractmethod
    def is_terminal(self, state):
        """
        Returns true iff the given state is terminal.
        :param state:
        :return:
        """

    @abc.abstractmethod
    def is_valid(self, state):
        """
        Returns true iff the given state is valid
        :param state:
        :return:
        """

    @abc.abstractmethod
    def create_action_pool(self):
        """
        :param solver:
        :return:
        """

    @abc.abstractmethod
    def create_root_historical_data(self, solver):
        """
        reset smart data for the root of the belief tree, if smart data is being used
        :return:
        """

    @abc.abstractmethod
    def get_max_undiscounted_return(self):
        """
        Calculate and return the highest possible undiscounted return
        :return:
        """

    @abc.abstractmethod
    def create_observation_pool(self, solver):
        """
        Return a concrete observation pool (discrete or continuous)
        :param solver:
        :return:
        """

    def generate_particles(self, previous_belief, action, obs, n_particles, prev_particles):
        """
        Generates new state particles based on the state particles of the previous node,
        * as well as on the action and observation.
        *
        * The default implementation uses rejection sampling, but this can be overridden to provide
        * a more efficient implementation.
        :param previous_belief:
        :param action:
        :param obs:
        :param n_particles:
        :return: list of particles
        """
        particles = []
        action_node = previous_belief.action_map.get_action_node(action)
        if action_node is None:
            return particles
        else:
            obs_map = action_node.observation_map
        child_node = obs_map.get_belief(obs)

        while particles.__len__() < n_particles:
            # sample a random particle
            state = random.choice(prev_particles)

            # Now generate a step in the model, and compare the observation to the actual observation.
            # Note that this comparison is done implicitly via the observation mapping, to ensure
            # that approximate observations are treated cleanly.
            result, is_legal = self.generate_step(state, action)
            if obs_map.get_belief(result.observation) is child_node:
                particles.append(result.next_state)
        return particles

    def generate_particles_uninformed(self, previous_belief, action, obs, n_particles):
        """
        Generates new state particles based on uninformed prior
        :param previous_belief:
        :param action:
        :param obs:
        :param n_particles:
        :return:
        """
        particles = []
        obs_map = previous_belief.action_map.get_action_node(action).observation_map
        child_node = obs_map.get_belief(obs)

        while particles.__len__() < n_particles:
            # sample a random particle
            state = self.sample_state_uninformed()

            # Now generate a step in the model, and compare the observation to the actual observation.
            # Note that this comparison is done implicitly via the observation mapping, to ensure
            # that approximate observations are treated cleanly.
            result, is_legal = self.generate_step(state, action)
            if obs_map.get_belief(result.observation) is child_node:
                particles.append(result.next_state)
        return particles class StepResult(object):
    """
     Represents the results of a complete step in the model, including the next state,
     * observation, and reward
     *
     * For convenience, this also includes the action taken, and a boolean flag representing
     * whether or not the resulting next state is a terminal state.
    """
    def __init__(self):
        self.action = None
        self.observation = None
        self.reward = 0
        self.next_state = None
        self.is_terminal = 0

    def print_step_result(self):
        print("------- Step Result --------")
        print("Action: ", end=' ')
        self.action.print_action()
        print("Observation: ", end=' ')
        self.observation.print_observation()
        print("Reward: ", end=' ')
        print(self.reward)
        print("Next state: ", end=' ')
        self.next_state.print_state()
        print("Is terminal: ", end=' ')
        print(self.is_terminal) class RockData(object):
    """
    Stores data about each rock
    """

    def __init__(self):
        # The number of times this rock has been checked.
        self.check_count = 0
        # The "goodness number"; +1 for each good observation of this rock, and -1 for each bad
        # observation of this rock.
        self.goodness_number = 0
        # The calculated probability that this rock is good.
        self.chance_good = 0.5

    def to_string(self):
        """
        Pretty printing
        """
        data_as_string = " Check count: " + str(self.check_count) + " Goodness number: " + \
                         str(self.goodness_number) + " Probability that rock is good: " + str(self.chance_good)
        return data_as_string class PositionAndRockData(HistoricalData):
    """
    A class to store the robot position associated with a given belief node, as well as
    explicitly calculated probabilities of goodness for each rock.
    """

    def __init__(self, model, grid_position, all_rock_data, solver):
        self.model = model
        self.solver = solver
        self.grid_position = grid_position

        # List of RockData indexed by the rock number
        self.all_rock_data = all_rock_data

        # Holds reference to the function for generating legal actions
        if self.model.preferred_actions:
            self.legal_actions = self.generate_smart_actions
        else:
            self.legal_actions = self.generate_legal_actions

    @staticmethod
    def copy_rock_data(other_data):
        new_rock_data = []
        [new_rock_data.append(RockData()) for _ in other_data]
        for i, j in zip(other_data, new_rock_data):
            j.check_count = i.check_count
            j.chance_good = i.chance_good
            j.goodness_number = i.goodness_number
        return new_rock_data

    def copy(self):
        """
        Default behavior is to return a shallow copy
        """
        return self.shallow_copy()

    def deep_copy(self):
        """
        Passes along a reference to the rock data to the new copy of RockPositionHistory
        """
        return PositionAndRockData(self.model, self.grid_position.copy(), self.all_rock_data, self.solver)

    def shallow_copy(self):
        """
        Creates a copy of this object's rock data to pass along to the new copy
        """
        new_rock_data = self.copy_rock_data(self.all_rock_data)
        return PositionAndRockData(self.model, self.grid_position.copy(), new_rock_data, self.solver)

    def update(self, other_belief):
        self.all_rock_data = other_belief.data.all_rock_data

    def any_good_rocks(self):
        any_good_rocks = False
        for rock_data in self.all_rock_data:
            if rock_data.goodness_number > 0:
                any_good_rocks = True
        return any_good_rocks

    def create_child(self, rock_action, rock_observation):
        next_data = self.deep_copy()
        next_position, is_legal = self.model.make_next_position(self.grid_position.copy(), rock_action.bin_number)
        next_data.grid_position = next_position

        if rock_action.bin_number is ActionType.SAMPLE:
            rock_no = self.model.get_cell_type(self.grid_position)
            next_data.all_rock_data[rock_no].chance_good = 0.0
            next_data.all_rock_data[rock_no].check_count = 10
            next_data.all_rock_data[rock_no].goodness_number = -10

        elif rock_action.bin_number >= ActionType.CHECK:
            rock_no = rock_action.rock_no
            rock_pos = self.model.rock_positions[rock_no]

            dist = self.grid_position.euclidean_distance(rock_pos)
            probability_correct = self.model.get_sensor_correctness_probability(dist)
            probability_incorrect = 1 - probability_correct

            rock_data = next_data.all_rock_data[rock_no]
            rock_data.check_count += 1

            likelihood_good = rock_data.chance_good
            likelihood_bad = 1 - likelihood_good

            if rock_observation.is_good:
                rock_data.goodness_number += 1
                likelihood_good *= probability_correct
                likelihood_bad *= probability_incorrect
            else:
                rock_data.goodness_number -= 1
                likelihood_good *= probability_incorrect
                likelihood_bad *= probability_correct

            if np.abs(likelihood_good) < 0.01 and np.abs(likelihood_bad) < 0.01:
                # No idea whether good or bad. reset data
                # print "Had to reset RockData"
                rock_data = RockData()
            else:
                rock_data.chance_good = old_div(likelihood_good, (likelihood_good + likelihood_bad))

        return next_data

    def generate_legal_actions(self):
        legal_actions = []
        all_actions = range(0, 5 + self.model.n_rocks)
        new_pos = self.grid_position.copy()
        i = new_pos.i
        j = new_pos.j

        for action in all_actions:
            if action is ActionType.NORTH:
                new_pos.i -= 1
            elif action is ActionType.EAST:
                new_pos.j += 1
            elif action is ActionType.SOUTH:
                new_pos.i += 1
            elif action is ActionType.WEST:
                new_pos.j -= 1

            if not self.model.is_valid_pos(new_pos):
                new_pos.i = i
                new_pos.j = j
                continue
            else:
                if action is ActionType.SAMPLE:
                    rock_no = self.model.get_cell_type(new_pos)
                    if 0 > rock_no or rock_no >= self.model.n_rocks:
                        continue
                new_pos.i = i
                new_pos.j = j
                legal_actions.append(action)
        return legal_actions

    def generate_smart_actions(self):

        smart_actions = []

        n_rocks = self.model.n_rocks

        # check if we are currently on top of a rock
        rock_no = self.model.get_cell_type(self.grid_position)

        # if we are on top of a rock, and it has been checked, sample it
        if 0 <= rock_no < n_rocks:
            rock_data = self.all_rock_data[rock_no]
            if rock_data.chance_good == 1.0 or rock_data.goodness_number > 0:
                smart_actions.append(ActionType.SAMPLE)
                return smart_actions

        worth_while_rock_found = False
        north_worth_while = False
        south_worth_while = False
        east_worth_while = False
        west_worth_while = False

        # Check to see which rocks are worthwhile

        # Only pursue one worthwhile rock at a time to prevent the agent from getting confused and
        # doing nothing
        for i in range(0, n_rocks):
            # Once an interesting rock is found, break out of the for loop

            if worth_while_rock_found:
                break
            rock_data = self.all_rock_data[i]
            if rock_data.chance_good != 0.0 and rock_data.goodness_number >= 0:
                worth_while_rock_found = True
                pos = self.model.rock_positions[i]
                if pos.i > self.grid_position.i:
                    south_worth_while = True
                elif pos.i < self.grid_position.i:
                    north_worth_while = True
                if pos.j > self.grid_position.j:
                    east_worth_while = True
                elif pos.j < self.grid_position.j:
                    west_worth_while = True

        # If no worth while rocks were found, just head east
        if not worth_while_rock_found:
            smart_actions.append(ActionType.EAST)
            return smart_actions

        if north_worth_while:
            smart_actions.append(ActionType.NORTH)
        if south_worth_while:
            smart_actions.append(ActionType.SOUTH)
        if east_worth_while:
            smart_actions.append(ActionType.EAST)
        if west_worth_while:
            smart_actions.append(ActionType.WEST)

        # see which rocks we might want to check
        for i in range(0, n_rocks):
            rock_data = self.all_rock_data[i]
            if rock_data.chance_good != 0.0 and rock_data.chance_good != 1.0 and np.abs(rock_data.goodness_number) < 2:
                smart_actions.append(ActionType.CHECK + i)

        return smart_actions
from builtins import str
from builtins import map
from builtins import hex
from builtins import range
from past.utils import old_div
from pomdpy.util import console, config_parser
from .grid_position import GridPosition
from .rock_state import RockState
from .rock_action import RockAction, ActionType
from .rock_observation import RockObservation
from pomdpy.discrete_pomdp import DiscreteActionPool, DiscreteObservationPool
from pomdpy.pomdp import Model, StepResult
from .rock_position_history import RockData, PositionAndRockData
import logging
import json
import numpy as np
from __future__ import print_function
from __future__ import absolute_import
from __future__ import division

module = "RockModel"


class RSCellType:
    """
    Rocks are enumerated 0, 1, 2, ...
    other cell types should be negative.
    """
    ROCK = 0
    EMPTY = -1
    GOAL = -2
    OBSTACLE = -3



class RockModel(Model):