import gym
from gym import spaces, error, utils
from gym.utils import seeding
import numpy as np
from scipy.spatial.distance import pdist, squareform
import configparser
from os import path
import matplotlib.pyplot as plt
from matplotlib.pyplot import gca

font = {'family' : 'sans-serif',
        'weight' : 'bold',
        'size'   : 14}

class FlockingEnv(gym.Env):

    def __init__(self):

        config_file = path.join(path.dirname(__file__), "params_flock.cfg")
        config = configparser.ConfigParser()
        config.read(config_file)
        config = config['flock']

        self.fig = None
        self.line1 = None
        self.filter_len = int(config['filter_length'])
        self.nx_system = 4
        self.n_nodes = int(config['network_size']) 
        self.comm_radius = float(config['comm_radius'])
        self.dt = float(config['system_dt'])
        self.v_max = float(config['max_vel_init'])
        self.v_bias = self.v_max  # 0.5 * self.v_max
        self.r_max = float(config['max_rad_init'])
        self.std_dev = float(config['std_dev']) * self.dt

        self.pooling = []
        if config.getboolean('sum_pooling'):
            self.pooling.append(np.nansum)
        if config.getboolean('min_pooling'):
            self.pooling.append(np.nanmin)
        if config.getboolean('max_pooling'):
            self.pooling.append(np.nanmax)
        self.n_pools = len(self.pooling)

        # number of features and outputs
        self.n_features = int(config['N_features'])
        self.nx = int(self.n_features / self.n_pools / self.filter_len)
        self.nu = int(config['N_outputs'])  # outputs

        self.x_agg = np.zeros((self.n_nodes, self.nx * self.filter_len, self.n_pools))
        self.x = np.zeros((self.n_nodes, self.nx_system))
        self.u = np.zeros((self.n_nodes, self.nu))
        self.mean_vel = np.zeros((self.n_nodes, self.nu))

        # TODO
        self.max_accel = 40
        self.max_z = 200  

        # self.b = np.ones((self.n_nodes,1))

        # self.action_space = spaces.Box(low=-self.max_accel, high=self.max_accel, shape=(self.n_nodes, 2), dtype=np.float32 )
        # self.observation_space = spaces.Box(low=-self.max_z, high=self.max_z, shape=(
        # self.n_nodes, self.nx * self.filter_len * self.n_pools) , dtype=np.float32)

        self.action_space = spaces.Box(low=-self.max_accel, high=self.max_accel, shape=(2,) , dtype=np.float32 )
        self.observation_space = spaces.Box(low=-self.max_z, high=self.max_z, shape=(self.n_features, ), dtype=np.float32)

        self.seed()

    def render(self, mode='human'):

        if self.fig is None:
            plt.ion()
            fig = plt.figure()
            ax = fig.add_subplot(111)
            line1, = ax.plot(self.x[:, 0], self.x[:, 1], 'bo')  # Returns a tuple of line objects, thus the comma
            ax.plot([0], [0], 'kx')
            plt.ylim(-1.0 * self.r_max, 1.0 * self.r_max)
            plt.xlim(-1.0 * self.r_max, 1.0 * self.r_max)
            a = gca()
            a.set_xticklabels(a.get_xticks(), font)
            a.set_yticklabels(a.get_yticks(), font)
            plt.title('GNN Controller')
            self.fig = fig
            self.line1 = line1

        self.line1.set_xdata(self.x[:, 0])
        self.line1.set_ydata(self.x[:, 1])
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()


    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def step(self, u):
        x = self.x
        x_ = np.zeros((self.n_nodes, self.nx_system))

        #u = np.vstack((np.zeros((self.n_leaders, 2)), u))
        # x position
        x_[:, 0] = x[:, 0] + x[:, 2] * self.dt
        # y position
        x_[:, 1] = x[:, 1] + x[:, 3] * self.dt
        # x velocity
        x_[:, 2] = x[:, 2] +  0.1 * u[:, 0] * self.dt + np.random.normal(0, self.std_dev,(self.n_nodes,))
        # y velocity
        x_[:, 3] = x[:, 3] +  0.1 * u[:, 1] * self.dt + np.random.normal(0, self.std_dev,(self.n_nodes,))
        # TODO - check the 0.1
        self.x = x_
        self.x_agg = self.aggregate(self.x, self.x_agg)
        self.u = u


        return self._get_obs(), -self.instant_cost(), False, {}

    def instant_cost(self):  # sum of differences in velocities
        return np.sum(np.var(self.x[:, 2:4], axis=0)) #+ np.sum(np.square(self.u)) * 0.00001
        #return np.sum(np.square(self.x[:,2:4] - self.mean_vel))

    def _get_obs(self):
        reshaped = self.x_agg.reshape((self.n_nodes, self.n_features))
        clipped = np.clip(reshaped, a_min=-self.max_z, a_max=self.max_z)
        return clipped #[self.n_leaders:, :]

    def reset(self):
        x = np.zeros((self.n_nodes, self.nx_system))
        degree = 0
        min_dist = 0

        while degree < 2 or min_dist < 0.1:  # < 0.25:  # 0.25:  #0.5: #min_dist < 0.25:
            # randomly initialize the state of all agents
            length = np.sqrt(np.random.uniform(0, self.r_max, size=(self.n_nodes,)))
            angle = np.pi * np.random.uniform(0, 2, size=(self.n_nodes,))
            x[:, 0] = length * np.cos(angle)
            x[:, 1] = length * np.sin(angle)

            bias = np.random.uniform(low=-self.v_bias, high=self.v_bias, size=(2,))
            x[:, 2] = np.random.uniform(low=-self.v_max, high=self.v_max, size=(self.n_nodes,)) + bias[0]
            x[:, 3] = np.random.uniform(low=-self.v_max, high=self.v_max, size=(self.n_nodes,)) + bias[1]

            # compute distances between agents
            x_t_loc = x[:, 0:2]  # x,y location determines connectivity
            a_net = squareform(pdist(x_t_loc.reshape((self.n_nodes, 2)), 'euclidean'))

            # no self loops
            a_net = a_net + 2 * self.comm_radius * np.eye(self.n_nodes)

            # compute minimum distance between agents and degree of network
            min_dist = np.min(np.min(a_net))
            a_net = a_net < self.comm_radius
            degree = np.min(np.sum(a_net.astype(int), axis=1))

            self.mean_vel = np.mean(x[:,2:4],axis=0)

        self.x = x
        self.x_agg = np.zeros((self.n_nodes, self.nx * self.filter_len, self.n_pools))
        self.x_agg = self.aggregate(self.x, self.x_agg)

        return self._get_obs()

    # def render(self, mode='human'):
    #     pass

    def close(self):
        pass

    def aggregate(self, xt, x_agg):
        """
        Perform aggegration operation for all possible pooling operations using helper functions get_pool and get_comms
        Args:
            x_agg (): Last time step's aggregated info
            xt (): Current state of all agents

        Returns:
            Aggregated state values
        """

        x_features = self.get_x_features(xt)
        a_net = self.get_connectivity(xt)
        for k in range(0, self.n_pools):
            comm_data = self.get_comms(np.dstack((x_features, self.get_features(x_agg[:, :, k]))), a_net)
            x_agg[:, :, k] = self.get_pool(comm_data, self.pooling[k])
        return x_agg

    def get_connectivity(self, x):
        """
        Get the adjacency matrix of the network based on agent locations by computing pairwise distances using pdist
        Args:
            x (): current states of all agents

        Returns: adjacency matrix of network

        """
        x_t_loc = x[:, 0:2]  # x,y location determines connectivity
        a_net = squareform(pdist(x_t_loc.reshape((self.n_nodes, 2)), 'euclidean'))
        a_net = (a_net < self.comm_radius).astype(float)
        np.fill_diagonal(a_net, 0)
        return a_net

    def get_x_features(self, xt): # TODO
        """
        Compute the non-linear features necessary for implementing Turner 2003
        Args:
            xt (): current state of all agents

        Returns: matrix of features for each agent

        """

        diff = xt.reshape((self.n_nodes, 1, self.nx_system)) - xt.reshape((1, self.n_nodes, self.nx_system))
        r2 = np.multiply(diff[:, :, 0], diff[:, :, 0]) + np.multiply(diff[:, :, 1], diff[:, :, 1]) + np.eye(
            self.n_nodes)
        return np.dstack((diff[:, :, 2], np.divide(diff[:, :, 0], np.multiply(r2, r2)), np.divide(diff[:, :, 0], r2),
                          diff[:, :, 3], np.divide(diff[:, :, 1], np.multiply(r2, r2)), np.divide(diff[:, :, 1], r2)))

    def get_features(self, agg):
        """
        Matrix of
        Args:
            agg (): the aggregated matrix from the last time step

        Returns: matrix of aggregated features from all nodes at current time

        """
        return np.tile(agg[:, :-self.nx].reshape((self.n_nodes, 1, -1)), (1, self.n_nodes, 1))  # TODO check indexing

    def get_comms(self, mat, a_net):
        """
        Enforces that agents who are not connected in the network cannot observe each others' states
        Args:
            mat (): matrix of state information for the whole graph
            a_net (): adjacency matrix for flock network (weighted networks unsupported for now)

        Returns:
            mat (): sparse matrix with NaN values where agents can't communicate

        """
        a_net[a_net == 0] = np.nan
        return mat * a_net.reshape(self.n_nodes, self.n_nodes, 1)

    def get_pool(self, mat, func):
        """
        Perform pooling operations on the matrix of state information. The replacement of values with NaNs for agents who
        can't communicate must already be enforced.
        Args:
            mat (): matrix of state information
            func (): pooling function (np.nansum(), np.nanmin() or np.nanmax()). Must ignore NaNs.

        Returns:
            information pooled from neighbors for each agent

        """
        return func(mat, axis=1).reshape((self.n_nodes, self.n_features))  # TODO check this axis = 1

    def controller(self):
        """
        The controller for flocking from Turner 2003.
        Args:
            x (): the current state
        Returns: the optimal action
        """
        x = self.x

        s_diff = x.reshape((self.n_nodes, 1, self.nx_system)) - x.reshape((1, self.n_nodes, self.nx_system))
        r2 = np.multiply(s_diff[:, :, 0], s_diff[:, :, 0]) + np.multiply(s_diff[:, :, 1], s_diff[:, :, 1]) + np.eye(
            self.n_nodes)
        p = np.dstack((s_diff, self.potential_grad(s_diff[:, :, 0], r2), self.potential_grad(s_diff[:, :, 1], r2)))
        p_sum = np.nansum(p, axis=1).reshape((self.n_nodes, self.nx_system + 2))
        return np.hstack(((- p_sum[:, 4] - p_sum[:, 2]).reshape((-1, 1)), (- p_sum[:, 3] - p_sum[:, 5]).reshape(-1, 1)))

    def potential_grad(self, pos_diff, r2):
        """
        Computes the gradient of the potential function for flocking proposed in Turner 2003.
        Args:
            pos_diff (): difference in a component of position among all agents
            r2 (): distance squared between agents

        Returns: corresponding component of the gradient of the potential

        """
        grad = -2.0 * np.divide(pos_diff, np.multiply(r2, r2)) + 2 * np.divide(pos_diff, r2)
        grad[r2 > self.comm_radius] = 0
        return grad


[((20, 17, 20, 44), 'configparser.ConfigParser', 'configparser.ConfigParser', ({}, {}), '()', False, 'import configparser\n'), ((50, 21, 50, 86), 'numpy.zeros', 'np.zeros', ({(50, 30, 50, 85): '(self.n_nodes, self.nx * self.filter_len, self.n_pools)'}, {}), '((self.n_nodes, self.nx * self.filter_len, self.n_pools))', True, 'import numpy as np\n'), ((51, 17, 51, 57), 'numpy.zeros', 'np.zeros', ({(51, 26, 51, 56): '(self.n_nodes, self.nx_system)'}, {}), '((self.n_nodes, self.nx_system))', True, 'import numpy as np\n'), ((52, 17, 52, 50), 'numpy.zeros', 'np.zeros', ({(52, 26, 52, 49): '(self.n_nodes, self.nu)'}, {}), '((self.n_nodes, self.nu))', True, 'import numpy as np\n'), ((53, 24, 53, 57), 'numpy.zeros', 'np.zeros', ({(53, 33, 53, 56): '(self.n_nodes, self.nu)'}, {}), '((self.n_nodes, self.nu))', True, 'import numpy as np\n'), ((65, 28, 65, 112), 'gym.spaces.Box', 'spaces.Box', (), '', False, 'from gym import spaces, error, utils\n'), ((66, 33, 66, 122), 'gym.spaces.Box', 'spaces.Box', (), '', False, 'from gym import spaces, error, utils\n'), ((94, 31, 94, 54), 'gym.utils.seeding.np_random', 'seeding.np_random', ({(94, 49, 94, 53): 'seed'}, {}), '(seed)', False, 'from gym.utils import seeding\n'), ((99, 13, 99, 53), 'numpy.zeros', 'np.zeros', ({(99, 22, 99, 52): '(self.n_nodes, self.nx_system)'}, {}), '((self.n_nodes, self.nx_system))', True, 'import numpy as np\n'), ((124, 18, 124, 72), 'numpy.clip', 'np.clip', (), '', True, 'import numpy as np\n'), ((128, 12, 128, 52), 'numpy.zeros', 'np.zeros', ({(128, 21, 128, 51): '(self.n_nodes, self.nx_system)'}, {}), '((self.n_nodes, self.nx_system))', True, 'import numpy as np\n'), ((158, 21, 158, 86), 'numpy.zeros', 'np.zeros', ({(158, 30, 158, 85): '(self.n_nodes, self.nx * self.filter_len, self.n_pools)'}, {}), '((self.n_nodes, self.nx * self.filter_len, self.n_pools))', True, 'import numpy as np\n'), ((199, 8, 199, 34), 'numpy.fill_diagonal', 'np.fill_diagonal', ({(199, 25, 199, 30): 'a_net', (199, 32, 199, 33): '(0)'}, {}), '(a_net, 0)', True, 'import numpy as np\n'), ((19, 32, 19, 54), 'os.path.dirname', 'path.dirname', ({(19, 45, 19, 53): '__file__'}, {}), '(__file__)', False, 'from os import path\n'), ((73, 12, 73, 21), 'matplotlib.pyplot.ion', 'plt.ion', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((74, 18, 74, 30), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((78, 12, 78, 57), 'matplotlib.pyplot.ylim', 'plt.ylim', ({(78, 21, 78, 38): '(-1.0 * self.r_max)', (78, 40, 78, 56): '(1.0 * self.r_max)'}, {}), '(-1.0 * self.r_max, 1.0 * self.r_max)', True, 'import matplotlib.pyplot as plt\n'), ((79, 12, 79, 57), 'matplotlib.pyplot.xlim', 'plt.xlim', ({(79, 21, 79, 38): '(-1.0 * self.r_max)', (79, 40, 79, 56): '(1.0 * self.r_max)'}, {}), '(-1.0 * self.r_max, 1.0 * self.r_max)', True, 'import matplotlib.pyplot as plt\n'), ((80, 16, 80, 21), 'matplotlib.pyplot.gca', 'gca', ({}, {}), '()', False, 'from matplotlib.pyplot import gca\n'), ((83, 12, 83, 39), 'matplotlib.pyplot.title', 'plt.title', ({(83, 22, 83, 38): '"""GNN Controller"""'}, {}), "('GNN Controller')", True, 'import matplotlib.pyplot as plt\n'), ((107, 56, 107, 105), 'numpy.random.normal', 'np.random.normal', ({(107, 73, 107, 74): '(0)', (107, 76, 107, 88): 'self.std_dev', (107, 89, 107, 104): '(self.n_nodes,)'}, {}), '(0, self.std_dev, (self.n_nodes,))', True, 'import numpy as np\n'), ((109, 56, 109, 105), 'numpy.random.normal', 'np.random.normal', ({(109, 73, 109, 74): '(0)', (109, 76, 109, 88): 'self.std_dev', (109, 89, 109, 104): '(self.n_nodes,)'}, {}), '(0, self.std_dev, (self.n_nodes,))', True, 'import numpy as np\n'), ((119, 22, 119, 52), 'numpy.var', 'np.var', (), '', True, 'import numpy as np\n'), ((139, 19, 139, 83), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n'), ((155, 28, 155, 52), 'numpy.mean', 'np.mean', (), '', True, 'import numpy as np\n'), ((213, 101, 214, 25), 'numpy.eye', 'np.eye', ({(214, 12, 214, 24): 'self.n_nodes'}, {}), '(self.n_nodes)', True, 'import numpy as np\n'), ((267, 109, 268, 25), 'numpy.eye', 'np.eye', ({(268, 12, 268, 24): 'self.n_nodes'}, {}), '(self.n_nodes)', True, 'import numpy as np\n'), ((134, 29, 134, 83), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n'), ((135, 28, 135, 73), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n'), ((136, 31, 136, 44), 'numpy.cos', 'np.cos', ({(136, 38, 136, 43): 'angle'}, {}), '(angle)', True, 'import numpy as np\n'), ((137, 31, 137, 44), 'numpy.sin', 'np.sin', ({(137, 38, 137, 43): 'angle'}, {}), '(angle)', True, 'import numpy as np\n'), ((140, 22, 140, 95), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n'), ((141, 22, 141, 95), 'numpy.random.uniform', 'np.random.uniform', (), '', True, 'import numpy as np\n'), ((151, 30, 151, 43), 'numpy.min', 'np.min', ({(151, 37, 151, 42): 'a_net'}, {}), '(a_net)', True, 'import numpy as np\n'), ((213, 13, 213, 54), 'numpy.multiply', 'np.multiply', ({(213, 25, 213, 38): 'diff[:, :, (0)]', (213, 40, 213, 53): 'diff[:, :, (0)]'}, {}), '(diff[:, :, (0)], diff[:, :, (0)])', True, 'import numpy as np\n'), ((213, 57, 213, 98), 'numpy.multiply', 'np.multiply', ({(213, 69, 213, 82): 'diff[:, :, (1)]', (213, 84, 213, 97): 'diff[:, :, (1)]'}, {}), '(diff[:, :, (1)], diff[:, :, (1)])', True, 'import numpy as np\n'), ((215, 88, 215, 116), 'numpy.divide', 'np.divide', ({(215, 98, 215, 111): 'diff[:, :, (0)]', (215, 113, 215, 115): 'r2'}, {}), '(diff[:, :, (0)], r2)', True, 'import numpy as np\n'), ((216, 88, 216, 116), 'numpy.divide', 'np.divide', ({(216, 98, 216, 111): 'diff[:, :, (1)]', (216, 113, 216, 115): 'r2'}, {}), '(diff[:, :, (1)], r2)', True, 'import numpy as np\n'), ((267, 13, 267, 58), 'numpy.multiply', 'np.multiply', ({(267, 25, 267, 40): 's_diff[:, :, (0)]', (267, 42, 267, 57): 's_diff[:, :, (0)]'}, {}), '(s_diff[:, :, (0)], s_diff[:, :, (0)])', True, 'import numpy as np\n'), ((267, 61, 267, 106), 'numpy.multiply', 'np.multiply', ({(267, 73, 267, 88): 's_diff[:, :, (1)]', (267, 90, 267, 105): 's_diff[:, :, (1)]'}, {}), '(s_diff[:, :, (1)], s_diff[:, :, (1)])', True, 'import numpy as np\n'), ((270, 16, 270, 36), 'numpy.nansum', 'np.nansum', (), '', True, 'import numpy as np\n'), ((283, 69, 283, 92), 'numpy.divide', 'np.divide', ({(283, 79, 283, 87): 'pos_diff', (283, 89, 283, 91): 'r2'}, {}), '(pos_diff, r2)', True, 'import numpy as np\n'), ((148, 51, 148, 71), 'numpy.eye', 'np.eye', ({(148, 58, 148, 70): 'self.n_nodes'}, {}), '(self.n_nodes)', True, 'import numpy as np\n'), ((215, 66, 215, 85), 'numpy.multiply', 'np.multiply', ({(215, 78, 215, 80): 'r2', (215, 82, 215, 84): 'r2'}, {}), '(r2, r2)', True, 'import numpy as np\n'), ((216, 66, 216, 85), 'numpy.multiply', 'np.multiply', ({(216, 78, 216, 80): 'r2', (216, 82, 216, 84): 'r2'}, {}), '(r2, r2)', True, 'import numpy as np\n'), ((283, 42, 283, 61), 'numpy.multiply', 'np.multiply', ({(283, 54, 283, 56): 'r2', (283, 58, 283, 60): 'r2'}, {}), '(r2, r2)', True, 'import numpy as np\n')]