"""
One-dimensional visual servo task.

This task gives Becca a chance to build a comparatively large number
of sensors into a few informative features. However, due to the
construction of the task, it's not strictly necessary to build
complex features to do well on it.
"""
import os

import matplotlib.pyplot as plt
import numpy as np

import becca.brain as becca_brain
from becca.base_world import World as BaseWorld
import becca_test.world_tools as wtools


class World(BaseWorld):
    """
    One-dimensional visual servo world

    In this world, Becca can direct its gaze left and right
    along a mural. It is rewarded for directing it near the center.
    Optimal performance is a reward of somewhere around .9 per time step.

    Some of this world's attributes are defined in base_world.py.
    The rest are defined below.
    """
    def __init__(self, lifespan=None):
        """
        Set up the world

        Parameters
        ----------
        lifespan : int
            The number of time steps to continue the world.
        """
        BaseWorld.__init__(self, lifespan)
        self.name = 'image_1D'
        print("Entering", self.name)

        # fov_span : int
        #     The world pixelizes its field of view into a superpixel array
        #     that is fov_span X fov_span.
        self.fov_span = 5
        self.n_sensors = self.fov_span ** 2
        # self.n_sensors = 2 * self.fov_span ** 2
        self.n_actions = 8
        # jump_fraction : float
        #     The fraction of time steps on which the agent jumps to
        #     a random position.
        self.jump_fraction = .1
        # step_cost : float
        #     The punishment per position step taken.
        self.step_cost = .1

        self.visualize_interval = 1e6
        # print_features : bool
        #     If True, plot and save visualizations of each of the features
        #     each time the world is visualized,
        #     rendered so that they represent what they mean in this world.
        self.print_features = False

        # Initialize the image to be used as the environment
        module_path = os.path.dirname(os.path.abspath(__file__))
        # image_filename : str
        #     The file name of the image including the relative path.
        self.image_filename = os.path.join(module_path,
                                           'images',
                                           'bar_test.png')
        # data : array of floats
        #     The image, read in and stored as a 2D numpy array.
        self.data = plt.imread(self.image_filename)
        # Convert it to grayscale if it's in color
        if self.data.shape[2] == 3:
            # Collapse the three RGB matrices into one b/w value matrix
            self.data = np.sum(self.data, axis=2) / 3.0
        # Define the size of the field of view, its range of
        # allowable positions, and its initial position
        image_width = self.data.shape[1]
        # max_step_size : int
        #     The largest step size allowed, in pixels in the original image.
        self.max_step_size = image_width / 2
        # target_column : int
        #     The column index that marks the center of the rewarded region.
        self.target_column = image_width / 2
        # reward_region_width : int
        #     The width of the region, in number of columns, within which
        #     the center of the field of view gets rewarded.
        self.reward_region_width = image_width / 8
        # noise_magnitude : float
        #     A scaling factor that drives how much inaccurate each movement
        #     will be.
        self.noise_magnitude = 0.1

        # column_history : list if ints
        #     A time series of the location (measured in column pixels) of the
        #     center of the brain's field of view.
        self.column_history = []
        # fov_height, fov_width : float
        #     The height and width (number of rows) of the field of view,
        #     in pixels.
        self.fov_height = np.min(self.data.shape)
        self.fov_width = self.fov_height
        # column_min, column_max : int
        #     The low and high bounds on where the field of view
        #     can be centered.
        self.column_min = np.ceil(self.fov_width / 2)
        self.column_max = np.floor(self.data.shape[1] - self.column_min)
        # column_position : int
        #     The current location of the center of the field of view.
        self.column_position = np.random.random_integers(self.column_min,
                                                         self.column_max)
        # block_width : int
        #     The width of each superpixel, in number of columns.
        self.block_width = self.fov_width / (self.fov_span + 2)
        self.sensors = np.zeros(self.n_sensors)
        self.action = np.zeros(self.n_actions)
        self.reward = 0.

    def step(self, action):
        """
        Advance the world by one time step

        Parameters
        ----------
        action : array of floats
            The set of action commands to execute.

        Returns
        -------
        self.reward : float
            The amount of reward or punishment given by the world.
        self.sensors : array of floats
            The values of each of the sensors.
        """
        self.timestep += 1
        self.action = action.ravel()
        self.action[np.nonzero(self.action)] = 1.

        # Actions 0-3 move the field of view to a higher-numbered row
        # (downward in the image) with varying magnitudes, and
        # actions 4-7 do the opposite.
        raw_col_step = (self.action[0] * self.max_step_size / 2 +
                        self.action[1] * self.max_step_size / 4 +
                        self.action[2] * self.max_step_size / 8 +
                        self.action[3] * self.max_step_size / 16 -
                        self.action[4] * self.max_step_size / 2 -
                        self.action[5] * self.max_step_size / 4 -
                        self.action[6] * self.max_step_size / 8 -
                        self.action[7] * self.max_step_size / 16)
        noise_factor = (
            self.noise_magnitude * np.random.random_sample() * 2.0 -
            self.noise_magnitude * np.random.random_sample() * 2.0 + 1.)
        column_step = int(raw_col_step * noise_factor)
        self.column_position = self.column_position + column_step
        self.column_position = max(self.column_position, self.column_min)
        self.column_position = min(self.column_position, self.column_max)
        self.column_history.append(self.column_position)

        # At random intervals, jump to a random position in the world.
        if np.random.random_sample() < self.jump_fraction:
            self.column_position = np.random.random_integers(self.column_min,
                                                             self.column_max)
        # Create the sensory input vector.
        fov = self.data[:, int(self.column_position - self.fov_width / 2):
                        int(self.column_position + self.fov_width / 2)]
        # Calculate center surround features for the image.
        center_surround_pixels = wtools.center_surround(fov,
                                                        self.fov_span,
                                                        self.fov_span)
        # unsplit_sensors = center_surround_pixels.ravel()
        self.sensors = center_surround_pixels.ravel()
        # These can be positive or negative, so split them into two
        # sets of sensors--one for the positive values and one for the
        # negative ones. Then stack them together for one big sensor array.
        # self.sensors = np.concatenate((np.maximum(unsplit_sensors, 0),
        #                           np.abs(np.minimum(unsplit_sensors, 0))))

        # Calculate the reward.
        self.reward = 0
        if (np.abs(self.column_position - self.target_column) <
                self.reward_region_width / 2.0):
            self.reward += 1.
        self.reward -= (np.abs(column_step) /
                        self.max_step_size * self.step_cost)
        return self.sensors, self.reward

    def visualize(self, brain):
        """
        Show what's going on in the world.
        """
        '''
        if self.print_features:
            feature_set = ft.get_feature_set(brain)
        for i_level, level_features in enumerate(feature_set):
            print('Features in level', i_level)
            for i_feature, feature in enumerate(level_features):
                print('Feature', i_feature, 'level', i_level)
                fig = plt.figure(num=99)
                fig.clf()
                fig, axarr = plt.subplots(1, 2 ** (i_level + 1),
                                          num=99, figsize=(3 * world.width,
                                                           3 * world.depth))
                for i_snap, snap in enumerate(feature):
                    print('    ', i_snap, ':', np.where(snap > 0)[0])
                    axis = axarr[i_snap]
                    # Convert projection to sensor activities
                    world.convert_sensors_to_detectors(snap)
                    plot_robot(world, axis, 0., 0., np.pi/2)
                    plot_sensors(world, axis, 0., 0., np.pi/2)

                for axis in axarr:
                    plt.sca(axis)
                    plt.axis('equal')
                    plt.axis('off')
                    plt.ylim((-world.depth, world.depth))
                    plt.xlim((-world.width, world.width))

                fig.canvas.draw()

                filename = '_'.join(('level', str(i_level).zfill(2),
                                     'sequence', str(i_feature).zfill(4),
                                     world.name, 'world.png'))
                full_filename = os.path.join(world.features_directory,
                                             filename)
                plt.title(filename)
                plt.savefig(full_filename, format='png')


            projections = brain.get_index_projections()[0]
            wtools.print_pixel_array_features(
                projections,
                self.fov_span ** 2 * 2,
                0,
                self.fov_span, self.fov_span,
                world_name=self.name)

        # Periodically show the state history and inputs as perceived by Becca.
        print(''.join(["world is ", str(self.timestep), " timesteps old"]))
        fig = plt.figure(11)
        plt.clf()
        plt.plot(self.column_history, 'k.', alpha=.2)
        plt.title(''.join(['Column history for ', self.name]))
        plt.xlabel('time step')
        plt.ylabel('position (pixels)')
        fig.show()
        fig.canvas.draw()

        fig = plt.figure(12)
        sensed_image = np.reshape(
            0.5 * (self.sensors[:len(self.sensors)/2] -
                   self.sensors[len(self.sensors)/2:] + 1),
            (self.fov_span, self.fov_span))
        plt.gray()
        plt.imshow(sensed_image, interpolation='nearest')
        plt.title("Image sensed")
        fig.show()
        fig.canvas.draw()
        '''


if __name__ == "__main__":
    becca_brain.run(World())
