# Path: HDPy/rl.py
# class Plant(object):
#     """A template for Actor-Critic *plants*. The *Plant* describes the
#     interaction of the Actor-Critic with the environment. Given a robot
#     which follows a certain *Policy*, the environment generates rewards
#     and robot states.
#     
#     An additional instance to :py:class:`PuPy.Normalization` may be 
#     supplied in ``norm`` for normalizing sensor values.
#     """
#     def __init__(self, state_space_dim=None, norm=None):
#         self._state_space_dim = state_space_dim
#         self.normalization = None
#         self.set_normalization(norm)
#     
#     def state_input(self, state):
#         """Return the state-part of the critic input
#         (i.e. the reservoir input).
#         
#         The state-part is derived from the current robot ``state`` and
#         possibly also its ``action``. As return format, a Nx1 numpy
#         vector is expected, where 2 dimensions should exist (e.g.
#         :py:meth:`numpy.atleast_2d`).
#         
#         Although the reservoir input will consist of both, the *state*
#         and *action*, this method must only return the *state* part of
#         it.
#         
#         """
#         raise NotImplementedError()
#     
#     def reward(self, epoch):
#         """A reward generated by the *Plant* based on the current
#         sensor readings in ``epoch``. The reward is single-dimensional.
#         
#         The reward is evaluated in every step. It builds the foundation
#         of the approximated return.
#         """
#         raise NotImplementedError()
#     
#     def state_space_dim(self):
#         """Return the dimension of the state space.
#         This value is equal to the size of the vector returned by
#         :py:meth:`state_input`.
#         """
#         if self._state_space_dim is None:
#             raise NotImplementedError()
#         return self._state_space_dim
#     
#     def set_normalization(self, norm):
#         """Set the normalization instance to ``norm``."""
#         if norm is None:
#             norm = PuPy.Normalization()
#         self.normalization = norm
#     
#     def reset(self):
#         """Reset plant to initial state."""
#         pass


from ...rl import Plant
import numpy as np
import warnings

class SpeedReward(Plant):
    """A :py:class:`Plant` with focus on the speed of the robot.
    """
    def __init__(self):
        super(SpeedReward, self).__init__(state_space_dim=2)
    
    def state_input(self, state):
        """Return the location, sampled from the *GPS* (x,y) values.
        The sample is an average over the last 10 GPS coordinates.
        """
        sio =  np.atleast_2d([
            self.normalization.normalize_value('puppyGPS_x', state['puppyGPS_x'][-10:]).mean(),
            self.normalization.normalize_value('puppyGPS_y', state['puppyGPS_y'][-10:]).mean()
        ]).T
        return sio
    
    def reward(self, epoch):
        """Return the covered distance and -1.0 if the robot tumbled.
        The speed measurement is taken from the 100th to the last sample.
        """
        if (epoch['accelerometer_z'][-100:] < 1.0).sum() > 80:
            return -1.0
        
        x = epoch['puppyGPS_x']
        y = epoch['puppyGPS_y']
        return np.linalg.norm(np.array([x[-1] - x[-100], y[-1] - y[-100]]))


# Path: HDPy/rl.py
# class Plant(object):
#     """A template for Actor-Critic *plants*. The *Plant* describes the
#     interaction of the Actor-Critic with the environment. Given a robot
#     which follows a certain *Policy*, the environment generates rewards
#     and robot states.
#     
#     An additional instance to :py:class:`PuPy.Normalization` may be 
#     supplied in ``norm`` for normalizing sensor values.
#     """
#     def __init__(self, state_space_dim=None, norm=None):
#         self._state_space_dim = state_space_dim
#         self.normalization = None
#         self.set_normalization(norm)
#     
#     def state_input(self, state):
#         """Return the state-part of the critic input
#         (i.e. the reservoir input).
#         
#         The state-part is derived from the current robot ``state`` and
#         possibly also its ``action``. As return format, a Nx1 numpy
#         vector is expected, where 2 dimensions should exist (e.g.
#         :py:meth:`numpy.atleast_2d`).
#         
#         Although the reservoir input will consist of both, the *state*
#         and *action*, this method must only return the *state* part of
#         it.
#         
#         """
#         raise NotImplementedError()
#     
#     def reward(self, epoch):
#         """A reward generated by the *Plant* based on the current
#         sensor readings in ``epoch``. The reward is single-dimensional.
#         
#         The reward is evaluated in every step. It builds the foundation
#         of the approximated return.
#         """
#         raise NotImplementedError()
#     
#     def state_space_dim(self):
#         """Return the dimension of the state space.
#         This value is equal to the size of the vector returned by
#         :py:meth:`state_input`.
#         """
#         if self._state_space_dim is None:
#             raise NotImplementedError()
#         return self._state_space_dim
#     
#     def set_normalization(self, norm):
#         """Set the normalization instance to ``norm``."""
#         if norm is None:
#             norm = PuPy.Normalization()
#         self.normalization = norm
#     
#     def reset(self):
#         """Reset plant to initial state."""
#         pass


# Path: HDPy/puppy/plant/plants.py
from ...rl import Plant
import numpy as np
import warnings


class SpeedReward(Plant):
    """A :py:class:`Plant` with focus on the speed of the robot.
    """
    def __init__(self):
        super(SpeedReward, self).__init__(state_space_dim=2)
    
    def state_input(self, state):
        """Return the location, sampled from the *GPS* (x,y) values.
        The sample is an average over the last 10 GPS coordinates.
        """
        sio =  np.atleast_2d([
            self.normalization.normalize_value('puppyGPS_x', state['puppyGPS_x'][-10:]).mean(),
            self.normalization.normalize_value('puppyGPS_y', state['puppyGPS_y'][-10:]).mean()
        ]).T
        return sio
    
    def reward(self, epoch):
        """Return the covered distance and -1.0 if the robot tumbled.
        The speed measurement is taken from the 100th to the last sample.
        """
        if (epoch['accelerometer_z'][-100:] < 1.0).sum() > 80:
            return -1.0
        
        x = epoch['puppyGPS_x']
        y = epoch['puppyGPS_y']
        return np.linalg.norm(np.array([x[-1] - x[-100], y[-1] - y[-100]]))


class LineFollower(Plant):