"""
Python Navigation Abstract Class
@Author: Kristian Gibson
TODO: Comments + Sphynx Docs Structured Text
TODO: Bug-fix, testing

"""

from abc import ABC
from scipy import integrate
from scipy.ndimage import interpolation
from spatialmath.base.transforms2d import *
from spatialmath.base.vectors import *
from spatialmath import SE2, SE3
from matplotlib import cm
from abc import ABC, abstractmethod
import matplotlib as mpl
import matplotlib.pyplot as plt
import copy
from roboticstoolbox.mobile.OccGrid import BaseOccupancyGrid, BinaryOccupancyGrid
from roboticstoolbox.mobile.Animations import VehiclePolygon
from colored import fg, attr

try:
    from progress.bar import FillingCirclesBar
    _progress = True
except:
    _progress = False

class PlannerBase(ABC):
    r"""
    Mobile robot motion planner (superclass)

    :param occgrid: occupancy grid, defaults to None
    :type occgrid: :class:`OccGrid` instance of ndarray(N,M), optional
    :param start: start position :math:`(x, y)` or configuration :math:`(x, y, \theta)`, defaults to None
    :type start: array_like(2) or array_like(3), optional
    :param goal: goal position :math:`(x, y)` or configuration :math:`(x, y, \theta)`, defaults to None
    :type goal: array_like(2) or array_like(3), optional
    :param inflate: obstacle inflation, defaults to 0
    :type inflate: float, optional
    :param ndims: dimensionality of the planning, either 2 for :math:`\mathbb{R}^2` or
        3 for :math:`\SE{2}`
    :param ndims: int, optional
    :param verbose: verbosity, defaults to False
    :type verbose: bool, optional
    :param msgcolor: color for message channel printing
    :type msgcolor: str, defaults to yellow
    :param seed: seed provided to private random number generator, defaults to None
    :type seed: int, optional

    Superclass for all mobile robot motion planners.  Key functionality
    includes:

    - encapsulates an occupancy grid and optionally inflates it
    - encapsulates a private random number generator with specifiable seed
    - encapsulates state such as start, goal, and the plan
    - provides a message channel for diagnostic output

    :seealso: :class:`OccGrid`
    """

    def __init__(self, occgrid=None, inflate=0, ndims=None,
                 verbose=False, msgcolor='yellow',
                 progress=True, marker=None, seed=None, **unused):

        self._occgrid = None
        if ndims is None:
            raise ValueError('ndims must be specified')
        self._ndims = ndims
        self._verbose = verbose
        self._msgcolor = msgcolor
        self._seed = seed
        self._private_random = np.random.default_rng(seed=seed)
        self._inflate = inflate
        self._start = None
        self._goal = None

        self._progress = progress and _progress

        self.marker = marker

        if occgrid is not None:
            if not isinstance(occgrid, BaseOccupancyGrid):
                occgrid = BinaryOccupancyGrid(occgrid)
            self._occgrid0 = occgrid  # original occgrid for reference

            if inflate > 0:
                self._occgrid = occgrid.copy()
                self._occgrid.inflate(inflate)
            else:
                self._occgrid = occgrid

    def __str__(self):
        """
        Compact representation of the planner

        :return: pretty printed representation
        :rtype: str
        """
        s = f"{self.__class__.__name__}: "
        if self._occgrid0 is not None:
            s += str(self.occgrid)
        if self._start is not None:
            s += f"\n  Start: {self.start}"
        if self._goal is not None:
            s += f"\n  Goal: {self.goal}"
        return s

    def __repr__(self):
        return str(self)

    @property
    def occgrid(self):
        """
        Occupancy grid

        :return: occupancy grid used for planning
        :rtype: :class:`OccGrid` instance or None

        Returns the grid that was optionally inflated at constructor time.

        :seealso: :class:`OccGrid`
        """
        return self._occgrid

    @property
    def start(self):
        r"""
        Start point or configuration used for planning

        :return: start point :math:`(x, y)` or configuration :math:`(x, y, \theta)`
        :rtype: ndarray(2) or ndarray(3)
        """
        return self._start

    @start.setter
    def start(self, start):
        r"""
        Set start point or configuration for planning

        :param start: Set start :math:`(x, y)` or configuration :math:`(x, y, \theta)`
        :type start: array_like(2) or array_like(3)
        :raises ValueError: if start point is occupied
        """
        if start is not None:
            if self.isoccupied(start):
                raise ValueError("Start location inside obstacle")
            self._start = base.getvector(start)

    @property
    def goal(self):
        r"""
        Goal point or configuration used for planning

        :return: goal point :math:`(x, y)` or configuration :math:`(x, y, \theta)`
        :rtype: ndarray(2) or ndarray(3)
        """
        return self._goal

    @goal.setter
    def goal(self, goal):
        r"""
        Set goal point or configuration for planning

        :param goal: Set goal :math:`(x, y)` or configuration :math:`(x, y, \theta)`
        :type goal: array_like(2) or array_like(3)
        :raises ValueError: if goal point is occupied
        """
        if goal is not None:
            if self.isoccupied(goal):
                raise ValueError("Goal location inside obstacle")
            self._goal = base.getvector(goal)

    def isoccupied(self, p):
        """
        Test if point is occupied

        :param p: world coordinate (x, y)
        :type p: array_like(2)
        :return: occupancy status of corresponding grid cell
        :rtype: bool

        The world coordinate is transformed and the status of the occupancy
        grid cell is returned.  If the point lies outside the bounds of
        the occupancy grid return True (obstacle)

        If there is no occupancy grid this function always returns False (free).

        :seealso: :meth:`OccGrid.isoccupied`
        """
        if self.occgrid is None:
            return False
        else:
            return self.occgrid.isoccupied(p)

    @property
    def verbose(self):
        """
        Get verbosity

        :return: verbosity
        :rtype: bool

        If ``verbosity`` print more diagnostic messages to the planner's
        message channel.
        """
        return self._verbose

    @verbose.setter
    def verbose(self, v):
        """
        Set verbosity

        :param v: verbosity
        :type v: bool

        If ``verbosity`` print more diagnostic messages to the planner's
        message channel.
        """
        self._verbose = v

    @property
    def random(self):
        """
        Private random number generator

        :return: random number generator
        :rtype: NumPy Generator

        For example::

            planner.random(size, dtype)
            planner.integers(low, high, size, dtype)
            planner.uniform(0, 5)

        :seealso: `Random Generator <https://numpy.org/doc/stable/reference/random/generator.html>`_
        """
        return self._private_random

    def randinit(self):
        if self._seed is not None:
            self._private_random = np.random.default_rng(seed=self._seed)

    # Define abstract classes to be implemented later
    def plan(self):
        r"""
        Plan path (abstract superclass)

        :param start: start position :math:`(x, y)` or configuration :math:`(x, y, \theta)`, defaults to None
        :type start: array_like(2) or array_like(3), optional
        :param goal: goal position :math:`(x, y)` or configuration :math:`(x, y, \theta)`, defaults to None
        :type goal: array_like(2) or array_like(3), optional

        The implementation depends on the particular planner.  Some may have
        no planning phase.  The plan may also depend on just the start or goal.
        """
        pass

    def validate_endpoint(self, p, dtype=None):
        if p is not None:
            p = base.getvector(p, self._ndims, dtype=dtype)
            if self.isoccupied(p):
                raise ValueError("Point is inside obstacle")
        return p

    def progress_start(self, n):
        self._bar = FillingCirclesBar(self.__class__.__name__, max=n,
            suffix = '%(percent).1f%% - %(eta)ds')

    def progress_next(self):
        self._bar.next()

    def progress_end(self):
        self._bar.finish()

    def query(self, start=None, goal=None, dtype=None, next=True, animate=False, movie=None):
        """
        Find a path from start to goal using plan (superclass)

        :param start: start position :math:`(x, y)` or configuration :math:`(x, y, \theta)`, defaults to None
        :type start: array_like(2) or array_like(3), optional
        :param goal: goal position :math:`(x, y)` or configuration :math:`(x, y, \theta)`, defaults to None
        :type goal: array_like(2) or array_like(3), optional
        :param animate: show the vehicle path, defaults to False
        :type animate: bool, optional
        :return: path from start to goal, one point :math:`(x, y)` or configuration :math:`(x, y, \theta)` per row
        :rtype: ndarray(N,2) or ndarray(N,3)

        Find a path from ``start`` to ``goal`` using a previously computed plan.

        The method performs the following steps:
        - Initialize navigation, invoke method N.navigate_init()
        - Visualize the environment, invoke method N.plot()
        - Iterate on the ``next()`` method of the subclass until the ``goal`` is
          achieved.
        """
        # make sure start and goal are set and valid
        self.start = self.validate_endpoint(start, dtype=dtype)
        self.goal = self.validate_endpoint(goal, dtype=dtype)

        # if movie is not None:
        #     animate = True

        if next:
            if animate:
                self.plot()

            # movie = MovieWriter(movie)

            robot = self._start
            path = [robot]

            while True:
                if animate:
                    plt.plot(robot[0], robot[1], 'y.', 12)
                    plt.pause(0.05)

                # get next point on the path
                robot = self.next(robot)

                # are we are done?
                if robot is None:
                    path.append(self._goal)
                    return np.array(path).astype(int)

                path.append(robot)

    def plot(self, path=None,
            style='striped', stripe=(('black', 4), ('yellow', 3)),
            stripe_r=(('black', 4), ('red', 3)),
            line=None, line_r=None,
            configspace=False, unwrap=True,
            direction=None, background=True,
            path_marker=None, path_marker_reverse=None,
            start_marker=None, goal_marker=None,
            start_vehicle=None, goal_vehicle=None,
            start=None, goal=None,
            ax=None, block=False, **kwargs):
        r"""
        Plot vehicle path

        :param path: path, defaults to None
        :type path: ndarray(N, 2) or ndarray(N, 3)
        :param style: line style: 'striped' [default] or 'line'
        :type style: str
        :param stripe: striped line style for forward motion
        :type stripe: tuple of (color, linewidth)
        :param stripe_r: striped line style for reverse motion
        :type stripe: tuple of (color, linewidth)
        :param line: plain line style for forward motion
        :type line: dict of arguments for ``plot``
        :param line_r: plain line style for forward motion
        :type line_r: dict of arguments for ``plot``
        :param direction: travel direction associated with each point on path, is either >0 or <0, defaults to None
        :type direction: ndarray(N,), optional
        :param configspace: plot the path in 3D configuration space, input must be 3xN.
            Start and goal style will be given by ``qstart_marker`` and ``qgoal_marker``, defaults to False
        :type configspace: bool, optional
        :param unwrap: for configuration space plot unwrap :math:`\theta` so
            there are no discontinuities at :math:`\pm \pi`, defaults to True
        :type unwrap: bool, optional
        :param background: plot occupancy grid if present, default True
        :type background: bool, optional
        :param start_marker: style for marking start point
        :type start_marker: dict, optional
        :param goal_marker: style for marking goal point
        :type goal_marker: dict, optional
        :param start_vehicle: style for vehicle animation object at start configuration
        :type start_vehicle: dict
        :param goal_vehicle: style for vehicle animation object at goal configuration
        :type goal_vehicle: dict
        :param start: start position :math:`(x, y)` or configuration :math:`(x, y, \theta)`, defaults to value used for ``plan/query``
        :type start: array_like(2) or array_like(3), optional
        :param goal: goal position :math:`(x, y)` or configuration :math:`(x, y, \theta)`, defaults to value used for ``plan/query``
        :type goal: array_like(2) or array_like(3), optional
        :param ax: axes to plot into
        :type ax: matplotlib axes
        :param block: block after displaying the plot
        :type block: bool, optional

        The path has one column per point and either 2 or 3 rows:

        - 2 rows describes motion in the :math:`x-y` plane and a 2D plot  is created
        - 3 rows describes motion in the :math:`x-y-\theta` configuration space. By
          default only the :math:`x-y` plane is plotted unless ``configspace``
          is True in which case motion in :math:`x-y-\theta` configuration space
          is shown.

        If the planner supports bi-directional motion then the ``direction``
        option gives the direction for each point on the path.

        The default line style is a ``'striped'`` line which comprises a wide first
        line with a slightly narrow dashed line plotted on top.  For example::

                (('black', 4), ('yellow', 3))

        is a blackline of width 4 with a dashed yellow line of width 3 plotted
        on top, giving a line of alternating black and yellow dashes.

        The parameters ``stripe`` and ``stripe_r`` specify the colors and widths for
        forward and reverse motion respectively.

        The ``'line'`` style is a regular Matplotlib and the parameters
        ``line`` and ``line_r`` specify the line and marker styles for
        forward and reverse motion respectively.

        For 2D plots with an :math:`x-y` path or 3D plots, the start and goal markers are specified
        by the dicts ``start_marker`` and ``goal_marker`` respectively.

        For 2D plots with an :math:`x-y-\theta` path the vehicle pose
        is indicated by a vehicle animation object passed to the constructor.

        Markers are specified as dicts using Matplotlib keywords, for example::

            planner.plot(path, path_marker=dict(marker='s', color='b'))

        Default values are provided for all markers:

            - the start point is a circle
            - the goal point is a star
            - the start vehicle style is a ``VehiclePolygon(shape='car')`` as
              an unfilled outline
            - the goal vehicle style is a ``VehiclePolygon(shape='car')`` as
              a transparent filled shape

        If ``configspace`` is True then direction-indicating markers are used to
        display start and goal configuration. These are also given as dicts but
        have two items: ``'shape'`` which is the shape of the polygonal marker
        and is either ``'triangle'`` or ``'car'``.  The second item ``'args'`` is
        passed to :func:`base.plot_poly` and Matplotlib.

        If ``background`` is True then the background of the plot is either or
        both of:

        - the occupancy grid
        - the distance field of the planner

        Additional arguments are passed through to :meth:`plot_bg`

        :seealso: :meth:`plot_bg` :func:`base.plot_poly`
        """
        # create default markers

        # passed to Matplotlib plot()
        if start_marker is None:
            start_marker = {'marker': 'o',
                            'markeredgecolor': 'w',
                            'markerfacecolor': 'y',
                            'markersize': 10,
                            'zorder': 10,
                           }
        if goal_marker is None:
            goal_marker = { 'marker': '*',
                            'markeredgecolor': 'w',
                            'markerfacecolor': 'y',
                            'markersize': 16,
                            'zorder': 10,
                          }

        # passed to VehiclePolygon
        if start_vehicle is None:
            start_vehicle = {'facecolor': 'none', 'edgecolor': 'k', 'linewidth': 2}

        if goal_vehicle is None:
            goal_vehicle = {'alpha': 0.5}

        ndims = self._ndims

        if ndims == 3 and not configspace:
            ndims = 2
            if path is not None:
                path = path[:, :2]

        if configspace and ndims < 3 and path is not None:
            raise ValueError(f"path should have {ndims} rows")

        ax = base.axes_logic(ax, ndims)

        # plot occupancy grid background
        if background:
            self.plot_bg(ax=ax, **kwargs)

        # mark the path
        if path is not None:
            if ndims == 2:
                # 2D case
                if direction is not None:
                    direction = np.array(direction)
                    if direction.shape[0] != path.shape[0]:
                        raise ValueError('direction vector must have same length as path')

                    while len(direction) > 0:
                        dir = direction[0]
                        change = np.argwhere(dir != direction)
                        if len(change) == 0:
                            k = -1
                        else:
                            k = change[0, 0]

                        if style == 'striped':
                            if dir > 0:
                                xstripe = stripe
                            else:
                                xstripe = stripe_r
                            ax.plot(path[:k, 0], path[:k, 1], color=xstripe[0][0], linewidth=xstripe[0][1])
                            ax.plot(path[:k, 0], path[:k, 1], color=xstripe[1][0], linewidth=xstripe[1][1], dashes=(5,5))
                        elif style == 'line':
                            if dir > 0:
                                ax.plot(path[:, 0], path[:, 1], **line)
                            else:
                                ax.plot(path[:, 0], path[:, 1], **line_r)

                        if len(change) == 0:
                            break
                        direction = direction[k-1:]
                        direction[0] = direction[1]
                        path = path[k-1:, :]

                else:
                    if style == 'striped':
                        ax.plot(path[:, 0], path[:, 1], color=stripe[0][0], linewidth=stripe[0][1], zorder=9)
                        ax.plot(path[:, 0], path[:, 1], color=stripe[1][0], linewidth=stripe[1][1], dashes=(5,5), zorder=9)
                    elif style == 'line':
                        ax.plot(path[:, 0], path[:, 1], **kwargs)
            elif ndims == 3:
                # 3D case
                if direction is not None:
                    direction = np.array(direction)
                    if direction.shape[0] != path.shape[0]:
                        raise ValueError('direction vector must have same length as path')
                    theta = path[:, 2]
                    if unwrap:
                        theta = np.unwrap(theta)

                    while len(direction) > 0:
                        dir = direction[0]
                        change = np.argwhere(dir != direction)
                        if len(change) == 0:
                            k = -1
                        else:
                            k = change[0, 0]

                        if style == 'striped':
                            if dir > 0:
                                xstripe = stripe
                            else:
                                xstripe = stripe_r
                            ax.plot(path[:k, 0], path[:k, 1], theta[:k], color=xstripe[0][0], linewidth=xstripe[0][1])
                            ax.plot(path[:k, 0], path[:k, 1], theta[:k], color=xstripe[1][0], linewidth=xstripe[1][1], dashes=(5,5))
                        elif style == 'line':
                            if dir > 0:
                                ax.plot(path[:, 0], path[:, 1], **line)
                            else:
                                ax.plot(path[:, 0], path[:, 1], **line_r)

                        if len(change) == 0:
                            break
                        direction = direction[k-1:]
                        direction[0] = direction[1]
                        path = path[k-1:, :]
                        theta = theta[k-1:]

                else:
                    theta = path[:, 2]
                    if unwrap:
                        theta = np.unwrap(theta)
                    if style == 'striped':
                        ax.plot(path[:, 0], path[:, 1], theta, color=stripe[0][0], linewidth=stripe[0][1])
                        ax.plot(path[:, 0], path[:, 1], theta, color=stripe[1][0], linewidth=stripe[1][1], dashes=(5,5))
                    elif style == 'line':
                        ax.plot(path[:, 0], path[:, 1], **line)

        # mark start and goal if requested
        if start is not None:
            start = self.validate_endpoint(start)
        else:
            start = self.start
        if goal is not None:
            self.goal = self.validate_endpoint(goal)
        else:
            goal = self.goal

        if ndims == 2 and self._ndims == 2:
            # proper 2d plot
            if start is not None:
                ax.plot(start[0], start[1], **start_marker)
            if goal is not None:
                ax.plot(goal[0], goal[1], **goal_marker)

        elif ndims == 2 and self._ndims == 3:
            # 2d projection of 3d plot, show start/goal configuration
            scale = base.axes_get_scale(ax) / 10

            if self.marker is None:
                self.marker = VehiclePolygon(shape='car', scale=scale)

            if start is not None:
                self.marker.plot(start, **start_vehicle)
            if goal is not None:
                self.marker.plot(goal, **goal_vehicle)

        elif ndims == 3:
            # 3d plot

            if start is not None:
                ax.plot(start[0], start[1], start[2], **start_marker)
            if goal is not None:

                if path is not None and unwrap:
                    theta = theta[-1]
                else:
                    theta = goal[2]
                plt.plot(goal[0], goal[1], theta, **goal_marker)

        ax.set_xlabel('x')
        ax.set_ylabel('y')
        if ndims == 2:
            ax.set_aspect('equal')
        else:
            ax.set_zlabel(r'$\theta$')

        plt.show(block=block)

        return ax

    def _qmarker(self, shape):
        h = 0.3
        t = 0.8  # start of head taper
        c = 0.5  # centre x coordinate
        w = 1    # width in x direction
        if shape == 'car':
            return np.array([
                [-c,     h],
                [t - c,  h],
                [w - c,  0],
                [t - c, -h],
                [-c,    -h],
            ]).T
        elif shape == 'triangle':
            return np.array([
                [-c,  h],
                [ w,  0],
                [-c, -h],
                [-c,  h],
            ]).T

    def plot_bg(self, distance=None, cmap='gray',
                ax=None, inflated=True,  **unused):
        """
        Plot background

        :param distance: override distance field, defaults to None
        :type distance: ndarray(N,M), optional
        :param cmap: Specify a colormap for the distance field, defaults to 'gray'
        :type cmap: str or Colormap, optional

        Displays the background which is either the occupancy grid or a distance
        field.  The distance field encodes the distance of a point from the goal, small
        distance is dark, a large distance is bright.

        If the planner has an occupancy grid then that will be displayed with:
            - free cells in white
            - occupied cells in red
            - inflated occupied cells in pink

        If distance is provided, or the planner has a distancemap attribute
        the the distance field will be used as the background and obstacle cells
        (actual or inflated) will be shown in red. A colorbar is added.
        """
        if self._occgrid is None:
            return

        if isinstance(self._occgrid, BaseOccupancyGrid):
            ax = base.plotvol2(dim=self._occgrid.workspace, ax=ax)
        else:
            ax = base.axes_logic(ax, 2)

        # create color map for free space + obstacle:
        #   free space, color index = 1, white, alpha=0 to allow background and grid lines to show
        #   obstacle, color index = 2, red, alpha=1

        if self._inflate > 0 and inflated:
            # 0 background (white, transparent)
            # 1 inflated obstacle (pink)
            # 2 original obstacle (red)
            colors = [(1, 1, 1, 0), (1, 0.75, 0.8, 1), (1, 0, 0, 1)]
            image = self.occgrid.grid.astype(int) + self._occgrid0.grid.astype(int)
        else:
            # 0 background
            # 1 obstacle
            colors = [(1, 1, 1, 0), (1, 0, 0, 1)]
            image = self.occgrid.grid

        if distance is None and hasattr(self, 'distancemap'):
            distance = self.distancemap

        if distance is not None:
            # distance field with obstacles

            # find largest finite value

            v = distance.ravel()
            vmax = max(v[np.isfinite(v)])

            # create a copy of greyscale color map
            c_map = copy.copy(mpl.cm.get_cmap(cmap))
            # c_map.set_bad(color=(1,0,0,1))  # nan and inf are red

            # change all inf to large value, so they are not 'bad' ie. red
            distance[np.isinf(distance)] = 2 * vmax
            c_map.set_over(color=(0,0,1))  # ex-infs are now blue

            # display image
            norm = mpl.colors.Normalize(vmin=0, vmax=vmax, clip=False)
            ax.imshow(distance, origin='lower',
                interpolation=None,
                cmap=c_map,
                norm=norm,
                )
            ax.grid(True, alpha=0.1, color=(1,1,1))

            # add colorbar
            scalar_mappable_c_map = cm.ScalarMappable(cmap=c_map, norm=norm)
            plt.colorbar(scalar_mappable_c_map, label='Distance', shrink=0.7, aspect=20*0.7)

            # overlay obstacles
            c_map = mpl.colors.ListedColormap(colors)
            self.occgrid.plot(image, cmap=c_map, zorder=1)

        else:
            # occupancy grid only

            # overlay obstacles
            c_map = mpl.colors.ListedColormap(colors)
            self.occgrid.plot(image, cmap=c_map, zorder=1)

        ax.set_facecolor((1, 1, 1)) # create white background
        ax.set_xlabel('x (cells)')
        ax.set_ylabel('y (cells)')
        ax.grid(True, zorder=0)

        # lock axis limits to current value
        # ax.set_xlim(ax.get_xlim())
        # ax.set_ylim(ax.get_ylim())

        plt.draw()
        plt.show(block=False)

    def message(self, s, color=None):
        """
        Print message to message channel

        :param s: message to print
        :type s: str
        :param color: color to print it, defaults to color specified at
            constructor time.
        :type color: str, optional

        """
        if self.verbose:
            if color is None:
                color = self._msgcolor
            print(fg(color), "Planner:: " + s, attr(0))

    # @staticmethod
    # def show_distance(d):
    #     d[np.isinf(d)] = None
    #     ax = plt.gca()
    #     c_map = plt.get_cmap("Greys")
    #     plt.clim(0, np.max(d[:]))
    #     plt.figimage(d)
    #     plt.xlabel('X')
    #     plt.ylabel('Y')
    #     plt.show()

class MovieWriter:

    def __init__(self, filename=None, interval=0.1, fig=None):
        """
        Save animation as a movie file

        :param filename: name of movie file, or tuple containing filename and
            frame interval
        :type filename: str or tuple(str, float)
        :param interval: frame interval, defaults to 0.1
        :type interval: float, optional
        :param fig: figure to record for the movie
        :type fig: figure reference

        Example::

            movie = MovieWriter(filename)

            while ...
                movie.add()

            movie.done()

        To avoid extra user-logic, if ``MovieWriter`` is called with ``filename`` equal to None,
        then the writer will do nothing when the ``add`` and ``done`` methods are called.
        """
        # Set up formatting for the movie files
        if filename is None:
            self.writer = None
            return

        if isinstance(filename, (tuple, list)):
            filename, interval = filename

        if os.path.exists(filename):
            print("overwriting movie", filename)
        else:
            print("creating movie", filename)
        self.writer = animation.FFMpegWriter(
            fps=round(1 / interval), extra_args=["-vcodec", "libx264"]
        )
        if fig is None:
            fig = plt.gcf()
        self.writer.setup(fig, filename)
        self.filename = filename

    def add(self):
        """
        Add frame to the movie
        """
        if self.writer is not None:
            self.writer.grab_frame()

    def done(self):
        if self.writer is not None:
            self.writer.finish()
            self.writer = None


import numpy as np
from scipy.interpolate import interp1d

from .. import constants as c
from . import scatmodels as sms
from ..distlib.composition import cmindex as cmi
from .. import distlib

__all__ = ['ScatModel','DiffScat','SigmaExt','SigmaScat','KappaExt','KappaScat']

#----------------------------------------------------------
# evals( emin=1.0, emax=2.0, de=0.1 ) : np.array [keV]
# angles( thmin=5.0, thmax=100.0, dth=5.0 ) : np.array [arcsec]
#

def evals( emin=1.0, emax=2.0, de=0.1 ):
    """
    FUNCTION evals( emin=1.0, emax=2.0, de=0.1 )
    RETURNS : np.array
    Distribution of energies [keV]
    """
    return np.arange( emin, emax+de, de )

def angles( thmin=5.0, thmax=100.0, dth=5.0 ):
    """
    FUNCTION angles( thmin=5.0, thmax=100.0, dth=5.0 )
    RETURNS : np.array
    Distribution of angles [arcsec]
    """
    return np.arange( thmin, thmax+dth, dth )

#-------------- Tie scattering mechanism to an index of refraction ------------------

class ScatModel(object):
    """
    | **ATTRIBUTES**
    | smodel : scattering model object : RGscat(), Mie()
    | cmodel : cmindex object : CmDrude(), CmGraphite(), CmSilicate()
    | stype  : string : 'RGscat', 'Mie'
    | cmtype : 'Drude', 'Silicate', 'Graphite'
    """
    def __init__( self, smodel=sms.RGscat(), cmodel=cmi.CmDrude() ):
        self.smodel = smodel
        self.cmodel = cmodel
        self.stype  = smodel.stype
        self.cmtype = cmodel.cmtype
        # cmtype choices : 'Drude' (requires rho term only)
        #                  'Graphite' (Carbonaceous grains)
        #                  'Silicate' (Astrosilicate)
        #                  --- Graphite and Silicate values come from Draine (2003)

#-------------- Quickly make a common ScatModel object ---------------------------

def makeScatModel( model_name, material_name ):
    """
    | **INPUTS**
    | model_name    : string : 'RG' or 'Mie'
    | material_name : string : 'Drude', 'Silicate', 'Graphite', 'SmallGraphite'
    |
    | **RETURNS**
    | ScatModel object
    """

    if model_name == 'RG':
        sm = sms.RGscat()
    elif model_name == 'Mie':
        sm = sms.Mie()
    else:
        print('Error: Model name not recognized')
        return

    if material_name == 'Drude':
        cm = cmi.CmDrude()
    elif material_name == 'Silicate':
        cm = cmi.CmSilicate()
    elif material_name == 'Graphite':
        cm = cmi.CmGraphite()
    elif material_name == 'SmallGraphite': # Small Graphite ~ 0.01 um
        cm = cmi.CmGraphite( size='small' )
    else:
        print('Error: CM name not recognized')
        return

    return ScatModel(sm, cm)

#-------------- Various Types of Scattering Cross-sections -----------------------

class DiffScat(object):
    """
    | A differential scattering cross-section [cm^2 ster^-1] integrated
    | over dust grain size distribution
    |
    | **ATTRIBUTES**
    | scatm : ScatModel
    | theta : np.array : arcsec
    | E     : scalar or np.array : Note, must match number of theta values if size > 1
    | a     : scalar : um
    | dsig  : np.array : cm^2 ster^-1
    """
    def __init__(self, scatm=ScatModel(), theta=angles(), E=1.0, a=1.0):
        self.scatm  = scatm
        self.theta  = theta
        self.E      = E
        self.a      = a

        cm   = scatm.cmodel
        scat = scatm.smodel
        # Do not print citation here, because this function is called multiple times by other modules

        if cm.cmtype == 'Graphite':
            dsig_pe = scat.Diff(theta=theta, a=a, E=E, cm=cmi.CmGraphite(size=cm.size, orient='perp'))
            dsig_pa = scat.Diff(theta=theta, a=a, E=E, cm=cmi.CmGraphite(size=cm.size, orient='para'))
            self.dsig = (dsig_pa + 2.0 * dsig_pe) / 3.0
        else:
            self.dsig   = scat.Diff(theta=theta, a=a, E=E, cm=cm)

class SigmaScat(object):
    """
    | Total scattering cross-section [cm^2] integrated over a dust grain
    | size distribution
    |
    | **ATTRIBUTES**
    | scatm : ScatModel
    | E     : scalar or np.array : keV
    | a     : scalar : um
    | qsca  : scalar or np.array : unitless scattering efficiency
    | sigma : scalar or np.array : cm^2
    """
    def __init__(self, scatm=ScatModel(), E=1.0, a=1.0):
        self.scatm  = scatm
        self.E      = E
        self.a      = a

        cm   = scatm.cmodel
        scat = scatm.smodel
        print(cm.citation)

        cgeo  = np.pi * np.power(a*c.micron2cm, 2)

        if cm.cmtype == 'Graphite':
            qsca_pe = scat.Qsca(a=a, E=E, cm=cmi.CmGraphite(size=cm.size, orient='perp'))
            qsca_pa = scat.Qsca(a=a, E=E, cm=cmi.CmGraphite(size=cm.size, orient='para'))
            self.qsca = (qsca_pa + 2.0*qsca_pe) / 3.0
        else:
            self.qsca = scat.Qsca(a=a, E=E, cm=cm)

        self.sigma = self.qsca * cgeo

class SigmaExt(object):
    """
    | Total EXTINCTION cross-section [cm^2] integrated over a dust grain
    | size distribution
    |
    | **ATTRIBUTES**
    | scatm : ScatModel
    | E     : scalar or np.array : keV
    | a     : scalar : um
    | qext  : scalar or np.array : unitless extinction efficiency
    | sigma : scalar or np.array : cm^2
    """
    def __init__(self, scatm=ScatModel(), E=1.0, a=1.0):
        self.scatm  = scatm
        self.E      = E
        self.a      = a

        if scatm.stype == 'RGscat':
            print('Rayleigh-Gans cross-section not currently supported for KappaExt')
            self.qext = None
            self.sigma = None
            return

        cm   = scatm.cmodel
        scat = scatm.smodel
        print(cm.citation)

        cgeo  = np.pi * np.power(a*c.micron2cm, 2)
        if cm.cmtype == 'Graphite':
            qext_pe = scat.Qext(a=a, E=E, cm=cmi.CmGraphite(size=cm.size, orient='perp'))
            qext_pa = scat.Qext(a=a, E=E, cm=cmi.CmGraphite(size=cm.size, orient='para'))
            self.qext = (qext_pa + 2.0*qext_pe) / 3.0
        else:
            self.qext = scat.Qext(a=a, E=E, cm=cm)
        self.sigma = self.qext * cgeo

class KappaScat(object):
    """
    | Opacity to scattering [g^-1 cm^2] integrated over dust grain size distribution.
    |
    | **ATTRIBUTES**
    | scatm : ScatModel
    | E     : scalar or np.array : keV
    | dist  : distlib.DustSpectrum
    | kappa : scalar or np.array : cm^2 g^-1, typically
    """
    def __init__(self, E=1.0, scatm=ScatModel(), dist=distlib.MRN_dist()):
        self.scatm  = scatm
        self.E      = E
        self.dist   = dist

        cm   = scatm.cmodel
        scat = scatm.smodel
        print(cm.citation)

        cgeo = np.pi * np.power(dist.a * c.micron2cm, 2)

        qsca    = np.zeros(shape=(np.size(E),np.size(dist.a)))
        qsca_pe = np.zeros(shape=(np.size(E),np.size(dist.a)))
        qsca_pa = np.zeros(shape=(np.size(E),np.size(dist.a)))

        # Test for graphite case
        if cm.cmtype == 'Graphite':
            cmGraphitePerp = cmi.CmGraphite(size=cm.size, orient='perp')
            cmGraphitePara = cmi.CmGraphite(size=cm.size, orient='para')

            if np.size(dist.a) > 1:
                for i in range(np.size(dist.a)):
                    qsca_pe[:,i] = scat.Qsca(E, a=dist.a[i], cm=cmGraphitePerp)
                    qsca_pa[:,i] = scat.Qsca(E, a=dist.a[i], cm=cmGraphitePara)
            else:
                qsca_pe = scat.Qsca(E, a=dist.a, cm=cmGraphitePerp)
                qsca_pa = scat.Qsca(E, a=dist.a, cm=cmGraphitePara)

            qsca    = (qsca_pa + 2.0 * qsca_pe) / 3.0

        else:
            if np.size(dist.a) > 1:
                for i in range(np.size(dist.a)):
                    qsca[:,i] = scat.Qsca(E, a=dist.a[i], cm=cm)
            else:
                qsca = scat.Qsca(E, a=dist.a, cm=cm)

        if np.size(dist.a) == 1:
            kappa = dist.nd * qsca * cgeo / dist.md
        else:
            kappa = np.array([])
            for j in range(np.size(E)):
                kappa = np.append(kappa,
                                  c.intz(dist.a, dist.nd * qsca[j,:] * cgeo) / dist.md)

        self.kappa = kappa

class KappaExt(object):
    """
    | Opacity to EXTINCTION [g^-1 cm^2] integrated over dust grain size
    | distribution
    |
    | **ATTRIBUTES**
    | scatm : ScatModel
    | E     : scalar or np.array : keV
    | dist  : distlib.DustSpectrum
    | kappa : scalar or np.array : cm^2 g^-1, typically
    """
    def __init__(self, E=1.0, scatm=ScatModel(), dist=distlib.MRN_dist()):
        self.scatm  = scatm
        self.E      = E
        self.dist   = dist

        if scatm.stype == 'RGscat':
            print 'Rayleigh-Gans cross-section not currently supported for KappaExt'
            self.kappa = None
            return

        cm   = scatm.cmodel
        scat = scatm.smodel
        print(cm.citation)

        cgeo = np.pi * np.power(dist.a * c.micron2cm, 2)

        qext    = np.zeros(shape=(np.size(E),np.size(dist.a)))
        qext_pe = np.zeros(shape=(np.size(E),np.size(dist.a)))
        qext_pa = np.zeros(shape=(np.size(E),np.size(dist.a)))

        # Test for graphite case
        if cm.cmtype == 'Graphite':
            cmGraphitePerp = cmi.CmGraphite(size=cm.size, orient='perp')
            cmGraphitePara = cmi.CmGraphite(size=cm.size, orient='para')

            if np.size(dist.a) > 1:
                for i in range(np.size(dist.a)):
                    qext_pe[:,i] = scat.Qext(E, a=dist.a[i], cm=cmGraphitePerp)
                    qext_pa[:,i] = scat.Qext(E, a=dist.a[i], cm=cmGraphitePara)
            else:
                qext_pe = scat.Qext(E, a=dist.a, cm=cmGraphitePerp)
                qext_pa = scat.Qext(E, a=dist.a, cm=cmGraphitePara)

            qext    = (qext_pa + 2.0 * qext_pe) / 3.0

        else:
            if np.size(dist.a) > 1:
                for i in range(np.size(dist.a)):
                    qext[:,i] = scat.Qext(E, a=dist.a[i], cm=cm)
            else:
                qext = scat.Qext(E, a=dist.a, cm=cm)

        if np.size(dist.a) == 1:
            kappa = dist.nd * qext * cgeo / dist.md
        else:
            kappa = np.array([])
            for j in range(np.size(E)):
                kappa = np.append(kappa,
                                  c.intz(dist.a, dist.nd * qext[j,:] * cgeo) / dist.md)

        self.kappa = kappa


# -*- coding: utf-8 -*-
"""
StationSim
Created on Tue Nov 20 15:25:27 2018
@author: medkmin
"""

# sspmm.py
'''
StationSim (aka Mike's model) converted into python.
'''
# Note cm (classmethods): Class Methods seem to have caused an issue.. Is so we can take these methods outside the class (excluding __init__).  This will reduce reproducing methods in the PF.

#%% INIT
import numpy as np
from scipy.spatial import cKDTree
import matplotlib.pyplot as plt
sqrt2 = np.sqrt(2)  # required for Agent.lerp()

def two_element_norm(arr):
    """
    A helpful function to calculate the norm for an array of two elements.
    This simply takes the square root of the sum of the square of the elements.
    This appears to be faster than np.linalg.norm.
    No doubt the numpy implementation would be faster for large arrays.
    Fortunately, all of our norms are of two-element arrays.

    :param arr:     A numpy array (or array-like DS) with length two.
    :return norm:   The norm of the array.
    """
    return np.sqrt(arr[0] * arr[0] + arr[1] * arr[1])

#%% MODEL
class Agent:
    """
    A class representing a generic agent for the StationSim ABM.
    """
    def __init__(self, model, unique_id):
        """
        Initialise a new agent.

        Creates a new agent and gives it a randomly chosen entrance, exit, and
        desired speed. All agents start with active state 0 ('not started').
        Their initial location (** HOW IS LOCATION REPRESENTED?? Answer: With (x,y) tuple-floats **) is set
        to the location of the entrance that they are assigned to.

        :param model: a pointer to the station sim model that is creating this agent
        """
        # Required
        self.unique_id = unique_id
        self.active = 0  # 0 Not Started, 1 Active, 2 Finished
        model.pop_active += 1

        # Choose at random at which of the entrances the agent starts
        self.location = model.loc_entrances[np.random.randint(model.entrances)]
        self.location[1] += model.entrance_space * (np.random.uniform() - .5)
        self.loc_desire = model.loc_exits[np.random.randint(model.exits)]

        # Parameters
        # model.entrance_speed -> the rate at which agents enter
        # self.time_activate -> the time at which the agent should become active
        # time_activate is exponentially distributed based on entrance_speed
        self.time_activate = np.random.exponential(model.entrance_speed)
        # The maximum speed that this agent can travel at:
        # self.speed_desire = max(np.random.normal(model.speed_desire_mean, model.speed_desire_std), 2*model.speed_min)  # this is not a truncated normal distribution
        self.speed_desire = model.speed_min - 1
        while self.speed_desire <= model.speed_min:
            self.speed_desire = np.random.normal(model.speed_desire_mean, model.speed_desire_std)
        self.wiggle = min(self.speed_desire, model.wiggle)  # if they can wiggle faster than they can move they may beat the expected time
        # A few speeds to check; used if a step at the max speed would cause a collision
        self.speeds = np.arange(self.speed_desire, model.speed_min, -model.speed_step)
        if model.do_save:
            self.history_loc = []
        self.time_expected = None
        self.time_start = None

    def step(self, model):
        """
        Iterate the agent. If they are inactive then it checks to see if they
        should become active. If they are active then then move (see
        self.move()) and, possibly, leave the model (see exit_query())).
        """
        if self.active == 0:
            self.activate(model)
        elif self.active == 1:
            self.move(model)
            self.exit_query(model)
            self.save(model)

    def activate(self, model):
        """
        Test whether an agent should become active. This happens when the model
        time is greater than the agent's activate time.
        """
        if not self.active and model.time_id > self.time_activate:
            self.active = 1
            self.time_start = model.time_id
            self.time_expected = (np.linalg.norm(self.location - self.loc_desire) - model.exit_space) / self.speed_desire

    @staticmethod
    def is_within_bounds(boundaries, new_location):
        """
        Check if new location is within the bounds of the model.
        :param boundaries      The boundaries of the model
        :param new_location    The proposed location for the agent
        :return                Is new location within boundaries, boolean
        """
        within0 = all(boundaries[0] <= new_location)
        within1 = all(boundaries[1] <= new_location)
        return within0 and within1

    def move(self, model):
        """
        Move the agent towards their destination. If the way is clear then the
        agent moves the maximum distance they can given their maximum possible
        speed (self.speed_desire). If not, then they iteratively test smaller
        and smaller distances until they find one that they can travel to
        without causing a collision with another agent.
        """
        for speed in self.speeds:
            # Direct
            new_location = Agent.lerp(self.loc_desire, self.location, speed)
            if not Agent.collision(model, new_location):
                break
            elif speed == self.speeds[-1]:
                # Wiggle
                # Why 1+1? Answer: randint(1)=0, randint(2)=0 or 1 - i think it is the standard pythonic idea of up to that number like array[0:2] is elements (array[0],array[1])
                # randint is upper-bound exclusive, so 2 instead of 1
                new_location = self.location + self.wiggle*np.random.randint(-1, 1+1, 2)
        # Rebound
        if not self.is_within_bounds(model.boundaries, new_location):
            new_location = np.clip(new_location, model.boundaries[0], model.boundaries[1])
        # Move
        self.location = new_location

    @classmethod
    def collision(cls, model, new_location):
        """
        Detects whether a move to the new_location will cause a collision
        (either with the model boundary or another agent).
        """
        within_bounds = all(model.boundaries[0] <= new_location) and all(new_location <= model.boundaries[1])
        if not within_bounds:
            collide = True
        elif Agent.neighbourhood(model, new_location):
            collide = True
        else:
            collide = False
        return collide

    @classmethod
    def neighbourhood(cls, model, new_location):
        """
        XXXX WHAT DOES THIS DO??  Answer: This method finds whether or not nearby neighbours are a collision.

         :param model:        the model that this agent is part of
         :param new_location: the proposed new location that the agent will move to
                         (a XXXX - what kind of object/data is the location?  Answer: the standard (x,y) floats-tuple)
         :param do_kd_tree    whether to use a spatial index (kd_tree) (default true)
        """
        neighbours = False
        neighbouring_agents = model.tree.query_ball_point(new_location, model.separation)
        for neighbouring_agent in neighbouring_agents:
            agent = model.agents[neighbouring_agent]
            if agent.active == 1 and new_location[0] <= agent.location[0]:
                neighbours = True
                break
        return neighbours

    @classmethod
    def lerp(cls, loc1, loc2, speed):
        """
        lerp - linear extrapolation
        Find the new position of after moving 'speed' distance from loc2 towards loc1.
            :param loc1: desired location
            :param loc2: current location
            :param speed: distance that can be covered in an iteration
            :return: The new location
        lerp is a intensively used method hence profiling and adjustments have been made, see 'github dust/Projects/awest/code/experiments/lerp.py' for more understanding.
        """
        #distance = np.linalg.norm(loc1 - loc2)     # frobenius norm: profiled at 8.05μs
        #distance = np.sqrt(sum((loc1 - loc2)**2))  # euclidean norm: profiled at 7.82μs
        #distance = sum(abs(loc1 - loc2))           # manhattan norm: profiled at 6.19μs
        reciprocal_distance = sqrt2 / sum(abs(loc1 - loc2))  # lerp5: profiled at 6.41μs
        loc = loc2 + speed * (loc1 - loc2) * reciprocal_distance
        return loc

    def exit_query(self, model):
        """
        Determine whether the agent should leave the model and, if so,
        remove them. Otherwise do nothing.
        """
        if sum(abs(self.location - self.loc_desire)) / sqrt2 < model.exit_space:
#        if two_element_norm(self.location - self.loc_desire) < model.exit_space:
            self.active = 2
            model.pop_active -= 1
            model.pop_finished += 1
            if model.do_save:
                time_delta = model.time_id - self.time_start
                model.time_taken.append(time_delta)
                time_delta -= self.time_expected
                model.time_delay.append(time_delta)

    def save(self, model):
        """
        Save agent location.
        """
        if model.do_save:
            self.history_loc.append(self.location)

class Model:
    """
    A class to represent the StationSim model.
    """
    def __init__(self, params):
        """
        Create a new model, reading parameters from a dictionary.
        XXXX Need to document the required parameters.
        """
        self.params = params
        # There are a lot of required attributes here that we hope are in params
        # Perhaps we should have a way to ensure we get what we require?
        # Also, consider using **kwargs
        [setattr(self, key, value) for key, value in params.items()]
        # Average number of speeds to check
        self.speed_step = (self.speed_desire_mean - self.speed_min) / 3
        # Batch Details
        self.time_id = 0
        self.step_id = 0
        if self.do_save:
            self.time_taken = []
            self.time_delay = []
        # Model Parameters
        self.boundaries = np.array([[0, 0], [self.width, self.height]])
        self.pop_active = 0
        self.pop_finished = 0
        # Initialise
        self.initialise_gates()
        self.agents = [Agent(self, unique_id) for unique_id in range(self.pop_total)]

    def step(self):
        """
        Iterate model forward one step.
        """
        if self.pop_finished < self.pop_total and self.step:
            self.kdtree_build()
            [agent.step(self) for agent in self.agents]
        self.time_id += 1
        self.step_id += 1

    def initialise_gates(self):
        """
        Initialise the locations of the entrances and exits.
        """
        self.loc_entrances = self.initialise_gates_generic(self.entrances, 0)
        self.loc_exits = self.initialise_gates_generic(self.exits, self.width)

    def initialise_gates_generic(self, n_gates, x):
        """
        General method for initialising gates.
        Note: This method relies on a lot of class attributes, many of which are
        not explicitly required in the init method - perhaps we should be
        careful of this?  Answer: see note cm at top
        """
        gates = np.zeros((n_gates, 2))
        gates[:, 0] = x
        if n_gates == 1:
            gates[0, 1] = self.height / 2
        else:
            gates[:, 1] = np.linspace(self.height / 4, 3 * self.height / 4,
                                      n_gates)
        return gates

    def kdtree_build(self):
        """
        Build kdtree for the model.
        """
        state = self.agents2state(do_ravel=False)
        self.tree = cKDTree(state)

    def agents2state(self, do_ravel=True):
        """
        Convert list of agents in model to state vector.
        """
        state = [agent.location for agent in self.agents]
        state = np.ravel(state) if do_ravel else np.array(state)
        return state

    def state2agents(self, state):
        """
        Use state vector to set agent locations.
        """
        for i in range(len(self.agents)):
            self.agents[i].location = state[2 * i:2 * i + 2]

    def batch(self):
        """
        Run the model.
        """
        print("Starting batch mode with following parameters:")
        print('\tParameter\tValue')
        for k, v in self.params.items():
            print('\t{0}:\t{1}'.format(k, v))
        print('')
        for i in range(self.batch_iterations):
            self.step()
            if i % 100 == 0:
                print("\tIterations: ", i)
            if self.do_ani:
                self.ani()
            if self.pop_finished == self.pop_total:
                print('Everyone made it!')
                break
        print("Finished at iteration", i)
        if self.do_save:
            self.save_stats()
            if self.do_plot: self.save_plot()

    def ani(self, agents=None, colour='k', alpha=1, show_separation=True):
        # Design for use in PF
        wid = 8  # image size
        hei = wid * self.height / self.width
        if show_separation:
            # the magic formular for marksize scaling
            magic = 1.8  # dependant on the amount of figure space used
            markersizescale = magic*72*hei/self.height
        plt.figure(1, figsize=(wid, hei))
        plt.clf()
        plt.axis(np.ravel(self.boundaries, 'F'))
        plt.axes().set_aspect('equal')
        for agent in self.agents[:agents]:
            if agent.active == 1:
                if show_separation:
                    plt.plot(*agent.location, marker='s', markersize=markersizescale*self.separation, color=colour, alpha=.05)
                plt.plot(*agent.location, marker='.', markersize=2, color=colour, alpha=alpha)
        plt.xlabel('Corridor Width')
        plt.ylabel('Corridor Height')
        plt.pause(1 / 30)
        return

    def save_ani(self):
        return

    def save_plot(self):
        """
        Produce plots for model.
        """
        self.plot_trails()
        self.plot_agent_times()

    def plot_trails(self):
        """
        Produce a plot of the trails of each agent in the 2-d corridor.
        """
        # Trails
        plt.figure()
        for agent in self.agents:
            if agent.active == 0:
                colour = 'r'
            elif agent.active == 1:
                colour = 'b'
            else:
                colour = 'm'
            locs = np.array(agent.history_loc).T
            plt.plot(locs[0], locs[1], color=colour, linewidth=.5)
        plt.axis(np.ravel(self.boundaries, 'F'))
        plt.xlabel('Corridor Width')
        plt.ylabel('Corridor Height')
        plt.legend(['Agent trails'])
        plt.show()

    def plot_agent_times(self):
        """
        Produce a plot of the time taken by each agent, and the delay of each
        agent.
        """
        # Time Taken, Delay Amount
        plt.figure()
        plt.hist(self.time_taken, alpha=.5, label='Time taken')
        plt.hist(self.time_delay, alpha=.5, label='Time delay')
        plt.xlabel('Time')
        plt.ylabel('Number of Agents')
        plt.legend()
        plt.show()

    def save_stats(self):
        """
        Print model run stats to console.
        """
        print()
        print('Stats:')
        print('Finish Time: ' + str(self.time_id))
        print('Active / Finished / Total agents: ' +
              str(self.pop_active) + '/' + str(self.pop_finished) +
              '/' + str(self.pop_total))
        print('Average time taken: {:.2f} steps'.format(np.mean(self.time_taken)))
        print('Average time delay: {:.2f} steps'.format(np.mean(self.time_delay)))

    def __repr__(self):
        """Print this model's ID and its memory location"""
        return "StationSim [{}]".format(hex(id(self)))

    @classmethod
    def run_defaultmodel(cls):
        """
        Run a model with some common parameters. Mostly used for testing.
        """
        np.random.seed(42)
        model_params = {
            'width': 200,
            'height': 100,
            'pop_total': 100,
            'entrances': 3,
            'entrance_space': 2,
            'entrance_speed': 1,
            'exits': 2,
            'exit_space': 1,
            'speed_min': .1,
            'speed_desire_mean': 1,
            'speed_desire_std': 1,
            'separation': 4,
            'wiggle': 1,
            'batch_iterations': 10_000,
            'do_save': True,
            'do_plot': False,
            'do_ani': True
        }
        # Run the model
        Model(model_params).batch()

# If this is called from the command line then run a default model.
if __name__ == '__main__':
    Model.run_defaultmodel()


#   Copyright 2020 The PyMC Developers
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
import warnings

from collections.abc import Mapping
from functools import singledispatch
from typing import Dict, Optional, Union

import aesara.tensor as at
import numpy as np

from aeppl import factorized_joint_logprob
from aeppl.transforms import TransformValuesOpt
from aesara import config
from aesara.graph.basic import graph_inputs, io_toposort
from aesara.graph.op import Op, compute_test_value
from aesara.tensor.random.op import RandomVariable
from aesara.tensor.subtensor import (
    AdvancedIncSubtensor,
    AdvancedIncSubtensor1,
    AdvancedSubtensor,
    AdvancedSubtensor1,
    IncSubtensor,
    Subtensor,
)
from aesara.tensor.var import TensorVariable

from pymc.aesaraf import extract_rv_and_value_vars, floatX, rvs_to_value_vars

@singledispatch
def logp_transform(op: Op):
    return None

def _get_scaling(total_size, shape, ndim):
    """
    Gets scaling constant for logp

    Parameters
    ----------
    total_size: int or list[int]
    shape: shape
        shape to scale
    ndim: int
        ndim hint

    Returns
    -------
    scalar
    """
    if total_size is None:
        coef = floatX(1)
    elif isinstance(total_size, int):
        if ndim >= 1:
            denom = shape[0]
        else:
            denom = 1
        coef = floatX(total_size) / floatX(denom)
    elif isinstance(total_size, (list, tuple)):
        if not all(isinstance(i, int) for i in total_size if (i is not Ellipsis and i is not None)):
            raise TypeError(
                "Unrecognized `total_size` type, expected "
                "int or list of ints, got %r" % total_size
            )
        if Ellipsis in total_size:
            sep = total_size.index(Ellipsis)
            begin = total_size[:sep]
            end = total_size[sep + 1 :]
            if Ellipsis in end:
                raise ValueError(
                    "Double Ellipsis in `total_size` is restricted, got %r" % total_size
                )
        else:
            begin = total_size
            end = []
        if (len(begin) + len(end)) > ndim:
            raise ValueError(
                "Length of `total_size` is too big, "
                "number of scalings is bigger that ndim, got %r" % total_size
            )
        elif (len(begin) + len(end)) == 0:
            return floatX(1)
        if len(end) > 0:
            shp_end = shape[-len(end) :]
        else:
            shp_end = np.asarray([])
        shp_begin = shape[: len(begin)]
        begin_coef = [floatX(t) / shp_begin[i] for i, t in enumerate(begin) if t is not None]
        end_coef = [floatX(t) / shp_end[i] for i, t in enumerate(end) if t is not None]
        coefs = begin_coef + end_coef
        coef = at.prod(coefs)
    else:
        raise TypeError(
            "Unrecognized `total_size` type, expected int or list of ints, got %r" % total_size
        )
    return at.as_tensor(floatX(coef))

subtensor_types = (
    AdvancedIncSubtensor,
    AdvancedIncSubtensor1,
    AdvancedSubtensor,
    AdvancedSubtensor1,
    IncSubtensor,
    Subtensor,
)

def logpt(
    var: TensorVariable,
    rv_values: Optional[Union[TensorVariable, Dict[TensorVariable, TensorVariable]]] = None,
    *,
    jacobian: bool = True,
    scaling: bool = True,
    transformed: bool = True,
    sum: bool = True,
    **kwargs,
) -> TensorVariable:
    """Create a measure-space (i.e. log-likelihood) graph for a random variable
    or a list of random variables at a given point.

    The input `var` determines which log-likelihood graph is used and
    `rv_value` is that graph's input parameter.  For example, if `var` is
    the output of a ``NormalRV`` ``Op``, then the output is a graph of the
    density function for `var` set to the value `rv_value`.

    Parameters
    ==========
    var
        The `RandomVariable` output that determines the log-likelihood graph.
        Can also be a list of variables. The final log-likelihood graph will
        be the sum total of all individual log-likelihood graphs of variables
        in the list.
    rv_values
        A variable, or ``dict`` of variables, that represents the value of
        `var` in its log-likelihood.  If no `rv_value` is provided,
        ``var.tag.value_var`` will be checked and, when available, used.
    jacobian
        Whether or not to include the Jacobian term.
    scaling
        A scaling term to apply to the generated log-likelihood graph.
    transformed
        Apply transforms.
    sum
        Sum the log-likelihood.

    """
    # TODO: In future when we drop support for tag.value_var most of the following
    # logic can be removed and logpt can just be a wrapper function that calls aeppl's
    # joint_logprob directly.

    # If var is not a list make it one.
    if not isinstance(var, list):
        var = [var]

    # If logpt isn't provided values and the variable (provided in var)
    # is an RV, it is assumed that the tagged value var or observation is
    # the value variable for that particular RV.
    if rv_values is None:
        rv_values = {}
        for _var in var:
            if isinstance(_var.owner.op, RandomVariable):
                rv_value_var = getattr(
                    _var.tag, "observations", getattr(_var.tag, "value_var", _var)
                )
                rv_values = {_var: rv_value_var}
    elif not isinstance(rv_values, Mapping):
        # Else if we're given a single value and a single variable we assume a mapping among them.
        rv_values = (
            {var[0]: at.as_tensor_variable(rv_values).astype(var[0].type)} if len(var) == 1 else {}
        )

    # Since the filtering of logp graph is based on value variables
    # provided to this function
    if not rv_values:
        warnings.warn("No value variables provided the logp will be an empty graph")

    if scaling:
        rv_scalings = {}
        for _var in var:
            rv_value_var = getattr(_var.tag, "observations", getattr(_var.tag, "value_var", _var))
            rv_scalings[rv_value_var] = _get_scaling(
                getattr(_var.tag, "total_size", None), rv_value_var.shape, rv_value_var.ndim
            )

    # Unlike aeppl, PyMC's logpt is expected to plug in the values variables to corresponding
    # RVs automatically unless the values are explicity set to None. Hence we iterate through
    # the graph to find RVs and construct a new RVs to values dictionary.
    tmp_rvs_to_values = rv_values.copy()
    transform_map = {}
    for node in io_toposort(graph_inputs(var), var):
        if isinstance(node.op, RandomVariable):
            curr_var = node.out
            rv_value_var = getattr(
                curr_var.tag, "observations", getattr(curr_var.tag, "value_var", curr_var)
            )
            rv_value = rv_values.get(curr_var, rv_value_var)
            tmp_rvs_to_values[curr_var] = rv_value
            # Along with value variables we also check for transforms if any.
            if hasattr(rv_value_var.tag, "transform") and transformed:
                transform_map[rv_value] = rv_value_var.tag.transform
        # The condition below is a hackish way of excluding the value variable for the
        # RV being indexed in case of Advanced Indexing of RVs. It gets added by the
        # logic above but aeppl does not expect us to include it in the dictionary of
        # {RV:values} given to it.
        if isinstance(node.op, subtensor_types):
            curr_var = node.out
            if (
                curr_var in tmp_rvs_to_values.keys()
                and curr_var.owner.inputs[0] in tmp_rvs_to_values.keys()
            ):
                tmp_rvs_to_values.pop(curr_var.owner.inputs[0])

    transform_opt = TransformValuesOpt(transform_map)
    temp_logp_var_dict = factorized_joint_logprob(
        tmp_rvs_to_values, extra_rewrites=transform_opt, use_jacobian=jacobian, **kwargs
    )

    # aeppl returns the logpt for every single value term we provided to it. This includes
    # the extra values we plugged in above so we need to filter those out.
    logp_var_dict = {}
    for value_var, _logp in temp_logp_var_dict.items():
        if value_var in rv_values.values():
            logp_var_dict[value_var] = _logp

    # If it's an empty dictionary the logp is None
    if not logp_var_dict:
        logp_var = None
    else:
        # Otherwise apply appropriate scalings and at.add and/or at.sum the
        # graphs accordingly.
        if scaling:
            for _value in logp_var_dict.keys():
                if _value in rv_scalings:
                    logp_var_dict[_value] *= rv_scalings[_value]

        if len(logp_var_dict) == 1:
            logp_var_dict = tuple(logp_var_dict.values())[0]
            if sum:
                logp_var = at.sum(logp_var_dict)
            else:
                logp_var = logp_var_dict
        else:
            if sum:
                logp_var = at.sum([at.sum(factor) for factor in logp_var_dict.values()])
            else:
                logp_var = at.add(*logp_var_dict.values())

        # Recompute test values for the changes introduced by the replacements
        # above.
        if config.compute_test_value != "off":
            for node in io_toposort(graph_inputs((logp_var,)), (logp_var,)):
                compute_test_value(node)

    return logp_var

def logcdfpt(
    var: TensorVariable,
    rv_values: Optional[Union[TensorVariable, Dict[TensorVariable, TensorVariable]]] = None,
    *,
    scaling: bool = True,
    sum: bool = True,
    **kwargs,
) -> TensorVariable:
    """Create a measure-space (i.e. log-cdf) graph for a random variable at a given point.

    Parameters
    ==========
    var
        The `RandomVariable` output that determines the log-likelihood graph.
    rv_values
        A variable, or ``dict`` of variables, that represents the value of
        `var` in its log-likelihood.  If no `rv_value` is provided,
        ``var.tag.value_var`` will be checked and, when available, used.
    jacobian
        Whether or not to include the Jacobian term.
    scaling
        A scaling term to apply to the generated log-likelihood graph.
    transformed
        Apply transforms.
    sum
        Sum the log-likelihood.

    """
    if not isinstance(rv_values, Mapping):
        rv_values = {var: rv_values} if rv_values is not None else {}

    rv_var, rv_value_var = extract_rv_and_value_vars(var)

    rv_value = rv_values.get(rv_var, rv_value_var)

    if rv_var is not None and rv_value is None:
        raise ValueError(f"No value variable specified or associated with {rv_var}")

    if rv_value is not None:
        rv_value = at.as_tensor(rv_value)

        if rv_var is not None:
            # Make sure that the value is compatible with the random variable
            rv_value = rv_var.type.filter_variable(rv_value.astype(rv_var.dtype))

        if rv_value_var is None:
            rv_value_var = rv_value

    rv_node = rv_var.owner

    rng, size, dtype, *dist_params = rv_node.inputs

    # Here, we plug the actual random variable into the log-likelihood graph,
    # because we want a log-likelihood graph that only contains
    # random variables.  This is important, because a random variable's
    # parameters can contain random variables themselves.
    # Ultimately, with a graph containing only random variables and
    # "deterministics", we can simply replace all the random variables with
    # their value variables and be done.
    tmp_rv_values = rv_values.copy()
    tmp_rv_values[rv_var] = rv_var

    logp_var = _logcdf(rv_node.op, rv_var, tmp_rv_values, *dist_params, **kwargs)

    transform = getattr(rv_value_var.tag, "transform", None) if rv_value_var else None

    # Replace random variables with their value variables
    replacements = rv_values.copy()
    replacements.update({rv_var: rv_value, rv_value_var: rv_value})

    (logp_var,), _ = rvs_to_value_vars(
        (logp_var,),
        apply_transforms=False,
        initial_replacements=replacements,
    )

    if sum:
        logp_var = at.sum(logp_var)

    if scaling:
        logp_var *= _get_scaling(
            getattr(rv_var.tag, "total_size", None), rv_value.shape, rv_value.ndim
        )

    # Recompute test values for the changes introduced by the replacements
    # above.
    if config.compute_test_value != "off":
        for node in io_toposort(graph_inputs((logp_var,)), (logp_var,)):
            compute_test_value(node)

    if rv_var.name is not None:
        logp_var.name = f"__logp_{rv_var.name}"

    return logp_var

def logp(var, rv_values, **kwargs):
    """Create a log-probability graph."""

    # Attach the value_var to the tag of var when it does not have one
    if not hasattr(var.tag, "value_var"):
        if isinstance(rv_values, Mapping):
            value_var = rv_values[var]
        else:
            value_var = rv_values
        var.tag.value_var = at.as_tensor_variable(value_var, dtype=var.dtype)

    return logpt(var, rv_values, **kwargs)

def logcdf(var, rv_values, **kwargs):
    """Create a log-CDF graph."""

    # Attach the value_var to the tag of var when it does not have one
    if not hasattr(var.tag, "value_var"):
        if isinstance(rv_values, Mapping):
            value_var = rv_values[var]
        else:
            value_var = rv_values
        var.tag.value_var = at.as_tensor_variable(value_var, dtype=var.dtype)

    return logcdfpt(var, rv_values, **kwargs)

@singledispatch
def _logcdf(op, values, *args, **kwargs):
    """Create a log-CDF graph.

    This function dispatches on the type of `op`, which should be a subclass
    of `RandomVariable`.  If you want to implement new log-CDF graphs
    for a `RandomVariable`, register a new function on this dispatcher.

    """
    raise NotImplementedError()

def logpt_sum(*args, **kwargs):
    """Return the sum of the logp values for the given observations.

    Subclasses can use this to improve the speed of logp evaluations
    if only the sum of the logp values is needed.
    """
    return logpt(*args, sum=True, **kwargs)