text: '''
  StationSim
      author: aw-west
      created: 19/06/18
      version: 0.8.1
  Changelog:
      .0 merged models
      .0 batch() moved to experiments
      .0 separated figures
      .1 checked figures and created experiment notebooks
      .1 fixed _heightmap
  Todo:
      new entering
      speeds[-1] separation drawn
      set_pickle?
      get_pickle?
      Look at model get_state and set_state - should comparisons be '==' or 'is'?
  '''
  
  import warnings
  import numpy as np
  import os
  from scipy.spatial import cKDTree
  import matplotlib.pyplot as plt
  from matplotlib.animation import FuncAnimation
  # Dont automatically load seaborn as it isn't needed on the HPC
  try: 
      from seaborn import kdeplot as sns_kdeplot
  except ImportError as e:
      warnings.warn("The seaborn module is not available. If you try to create kde plots for this model (i.e. a wiggle map or density map) then it will fail.")
  
  
  class Agent:
      '''
      A class representing a generic agent for the StationSim ABM.
      '''
      def __init__(self, model, unique_id):
          '''
          Initialise a new agent.
  
          Desctiption:
              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 (** (x,y) tuple-floats **) is set to the location
              of the entrance that they are assigned to.
  
          Parameters:
              model - a pointer to the StationSim model that is creating this agent
          '''
          self.unique_id = unique_id
          # Required
          self.status = 0  # 0 Not Started, 1 Active, 2 Finished
          # Location
          perturb = model.gates_space * np.random.uniform(-1, +1)
          gate_in = np.random.randint(model.gates_in)
          self.loc_start = model.gates_locations[gate_in] + [0, perturb]
          gate_out = np.random.randint(model.gates_out) + model.gates_in
          self.loc_desire = model.gates_locations[gate_out]
          self.location = self.loc_start
          # Speed
          speed_max = 0
          while speed_max <= model.speed_min:
              speed_max = np.random.normal(model.speed_mean, model.speed_std)
          self.speeds = np.arange(speed_max, model.speed_min, -model.speed_step)
          self.speed = None
          # Others
          self.steps_activate = np.random.exponential(model.gates_speed)
          self.wiggle = min(model.max_wiggle, speed_max)
          # History
          if model.do_history:
              self.history_locations = []
              self.history_speeds = []
              self.history_wiggles = 0
              self.history_collisions = 0
              self.step_start = None
  
      def step(self, model):
          '''
          Iterate the agent.
  
          Description:
              If they are inactive then it checks to see if they should become active.
              If they are active then they move or leave the model.
          '''
          if self.status == 0:
              self.activate(model)
          elif self.status == 1:
              self.move(model)
              self.deactivate(model)
          self.history(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 model.step_id > self.steps_activate:
              self.status = 1
              model.pop_active += 1
              self.step_start = model.step_id
  
      @staticmethod
      def distance(loc1, loc2):
          '''
          A helpful function to calculate the distance between two points.
          This simply takes the square root of the sum of the square of the elements.
          This appears to be faster than using 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.
          '''
          x = loc1[0] - loc2[0]
          y = loc1[1] - loc2[1]
          norm = (x*x + y*y)**.5
          return norm
  
      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 colision with another agent.
          '''
          direction = (self.loc_desire - self.location) / self.distance(self.loc_desire, self.location)
          for speed in self.speeds:
              # Direct. Try to move forwards by gradually smaller and smaller amounts
              new_location = self.location + speed * direction
              if self.collision(model, new_location):
                  if model.do_history:
                      self.history_collisions += 1
                      model.history_collision_locs.append(new_location)
                      model.history_collision_times.append(model.step_id)
  
              else:
                  break
              # If even the slowest speed results in a colision, then wiggle.
              if speed == self.speeds[-1]:
                  new_location = self.location + [0, self.wiggle*np.random.randint(-1, 1+1)]
                  if model.do_history:
                      self.history_wiggles += 1
                      model.history_wiggle_locs.append(new_location)
          # Rebound
          if not model.is_within_bounds(new_location):
              new_location = model.re_bound(new_location)
          # Move
          self.location = new_location
          self.speed = speed
  
      def collision(self, model, new_location):
          '''
          Detects whether a move to the new_location will cause a collision
          (either with the model boundary or another agent).
          '''
          if not model.is_within_bounds(new_location):
              collide = True
          elif self.neighbourhood(model, new_location):
              collide = True
          else:
              collide = False
          return collide
  
      def neighbourhood(self, model, new_location):
          '''
          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 standard (x,y) floats-tuple)
          '''
          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.status == 1 and self.unique_id != agent.unique_id and new_location[0] <= agent.location[0]:
                  neighbours = True
                  break
          return neighbours
  
      def deactivate(self, model):
          '''
          Determine whether the agent should leave the model and, if so,
          remove them. Otherwise do nothing.
          '''
          if self.distance(self.location, self.loc_desire) < model.gates_space:
              self.status = 2
              model.pop_active -= 1
              model.pop_finished += 1
              if model.do_history:
                  steps_exped = (self.distance(self.loc_start, self.loc_desire) - model.gates_space) / self.speeds[0]
                  model.steps_exped.append(steps_exped)
                  steps_taken = model.step_id - self.step_start
                  model.steps_taken.append(steps_taken)
                  steps_delay = steps_taken - steps_exped
                  model.steps_delay.append(steps_delay)
  
      def history(self, model):
          '''
          Save agent location.
          '''
          if model.do_history:
              if self.status == 1:
                  self.history_locations.append(self.location)
              else:
                  self.history_locations.append((None, None))
  
  
  class Model:
      '''
      StationSim Model
  
      Description:
          An Agent-Based Model (ABM) that synchronously `steps`
          step()
  
      Params:
          unique_id
          **kwargs    # check `params`, and `params_changed`
          do_history  # save memory
          do_print    # mute printing
  
      Returns:
          step_id
          params
          params_changed
  
          get_state()
          set_state()
  
          get_analytics()
          get_trails()
          get_timehist()
          get_location_map()
          get_wiggle_map()
          get_ani()
      '''
      def __init__(self, unique_id=None, **kwargs):
          '''
          Create a new model, reading parameters from a keyword arguement dictionary.
          '''
          self.unique_id = unique_id
          self.status = 1
          # Default Parameters (usually overridden by the caller)
          params = {
              'pop_total': 100,
  
              'width': 400,
              'height': 200,
  
              'gates_in': 3,
              'gates_out': 2,
              'gates_space': 1, # Distance around gate that will cause the agent to leave the simulation
              'gates_speed': 1,
  
              # Agent maximum speed is chosen when the agent is created and drawn from a normal distribution
              'speed_min': .2, # No speeds can be below this
              'speed_mean': 1,
              'speed_std': 1,
              'speed_steps': 3, #
  
              'separation': 5,
              'max_wiggle': 1,
  
              'step_limit': 3600,
  
              'do_history': True,
              'do_print': True,
  
              'random_seed': int.from_bytes(os.urandom(4), byteorder='little')
  
          }
          if len(kwargs) == 0:
              warnings.warn(
                  "No parameters have been passed to the model; using the default parameters: {}".format(params),
                  RuntimeWarning
              )
          self.params, self.params_changed = Model._init_kwargs(params, kwargs)
          [setattr(self, key, value) for key, value in self.params.items()]
          # Set the random seed
          np.random.seed(self.random_seed)
          # Constants
          self.speed_step = (self.speed_mean - self.speed_min) / self.speed_steps
          self.boundaries = np.array([[0, 0], [self.width, self.height]])
          # Following replaced with a normal function
          #gates_init = lambda x, y, n: np.array([np.full(n, x), np.linspace(0, y, n + 2)[1:-1]]).T
          self.gates_locations = np.concatenate([Model._gates_init(0, self.height, self.gates_in), Model._gates_init(self.width, self.height, self.gates_out)])
          # Variables
          self.step_id = 0
          self.pop_active = 0
          self.pop_finished = 0
          # Initialise
          self.agents = [Agent(self, unique_id) for unique_id in range(self.pop_total)]
          # Following replaced with a normal function
          #self.is_within_bounds = lambda loc: all(self.boundaries[0] <= loc) and all(loc <= self.boundaries[1])
          #self.re_bound = lambda loc: np.clip(loc, self.boundaries[0], self.boundaries[1])
          if self.do_history:
              self.history_state = []
              self.history_wiggle_locs = []
              self.history_collision_locs = []
              self.history_collision_times = []
              self.steps_taken = []
              self.steps_exped = []
              self.steps_delay = []
              # Figure Shape Stuff
              self._wid = 8
              self._rel = self._wid / self.width
              self._hei = self._rel * self.height
              self._figsize = (self._wid, self._hei)
              self._dpi = 160
  
      @staticmethod
      def _gates_init(x, y, n):
          return np.array([np.full(n, x), np.linspace(0, y, n+2)[1:-1]]).T
  
      def is_within_bounds(self, loc):
          return all(self.boundaries[0] <= loc) and all(loc <= self.boundaries[1])
  
      def re_bound(self, loc):
          return np.clip(loc, self.boundaries[0], self.boundaries[1])
  
      @staticmethod
      def _init_kwargs(dict0, dict1):
          '''
          Internal dictionary update tool
  
          dict0 is updated by dict1 adding no new keys.
          dict2 is the changes excluding 'do_' keys.
          '''
          dict2 = dict()
          for key in dict1.keys():
              if key in dict0:
                  if dict0[key] is not dict1[key]:
                      dict0[key] = dict1[key]
                      if 'do_' not in key:
                          dict2[key] = dict1[key]
              else:
                  print(f'BadKeyWarning: {key} is not a model parameter.')
          return dict0, dict2
  
      def step(self):
          '''
          Iterate model forward one step.
          '''
          if self.pop_finished < self.pop_total and self.step_id < self.step_limit and self.status==1:
              if self.do_print and self.step_id%100==0:
                  print(f'\tIteration: {self.step_id}/{self.step_limit}')
              state = self.get_state('location2D')
              self.tree = cKDTree(state)
              [agent.step(self) for agent in self.agents]
              if self.do_history:
                  self.history_state.append(state)
              self.step_id += 1
          else:
              if self.do_print and self.status==1:
                  print(f'StationSim {self.unique_id} - Everyone made it!')
                  self.status = 0
  
      # State
      def get_state(self, sensor=None):
          '''
          Convert list of agents in model to state vector.
          '''
          if sensor is None:
              state = [(agent.status, *agent.location, agent.speed) for agent in self.agents]
              state = np.append(self.step_id, np.ravel(state))
          elif sensor is 'location':
              state = [agent.location for agent in self.agents]
              state = np.ravel(state)
          elif sensor is 'location2D':
              state = [agent.location for agent in self.agents]
          return state
  
      def set_state(self, state, sensor=None):
          '''
          Use state vector to set agent locations.
          '''
          if sensor is None:
              self.step_id = int(state[0])
              state = np.reshape(state[1:], (self.pop_total, 3))
              for i, agent in enumerate(self.agents):
                  agent.status = int(state[i, 0])
                  agent.location = state[i, 1:]
          elif sensor is 'location':
              state = np.reshape(state, (self.pop_total, 2))
              for i, agent in enumerate(self.agents):
                  agent.location = state[i, :]
          elif sensor is 'location2D':
              for i, agent in enumerate(self.agents):
                  agent.location = state[i, :]
  
      # TODO: Deprecated, update PF
      def agents2state(self, do_ravel=True):
          warnings.warn("Replace 'state = agents2state()' with 'state = get_state(sensor='location')'", DeprecationWarning)
          return self.get_state(sensor='location')
      def state2agents(self, state):
          warnings.warn("Replace 'state2agents(state)' with 'set_state(state, sensor='location')'", DeprecationWarning)
          return self.set_state(state, sensor='location')
  
      # Analytics
      def get_analytics(self, sig_fig=None):
          '''
          A collection of analytics.
          '''
          analytics = {
              'Finish Time': self.step_id,
              'Total': self.pop_total,
              'Active': self.pop_active,
              'Finished': self.pop_finished,
              'Mean Time Taken': np.mean(self.steps_taken),
              'Mean Time Expected': np.mean(self.steps_exped),
              'Mean Time Delay': np.mean(self.steps_delay),
              'Mean Collisions': np.mean([agent.history_collisions for agent in self.agents]),
              'Mean Wiggles': np.mean([agent.history_wiggles for agent in self.agents]),
              # 'GateWiggles': sum(wig[0]<self.gates_space for wig in self.history_wiggle_locs)/self.pop_total
              }
          return analytics
  
      def get_trails(self, plot_axis=False, plot_legend=True, colours=('b','g','r'), xlim=None, ylim=None):
          """
          Make a figure showing the trails of the agents.
  
          :param plot_axis: Whether to show the axis (default False)
          :param plot_legend: Whether to show the legend (default False)
          :param colours: Optional tuple with three values representing the colours of agents in states
          1 (no started), 2 (active), 3 (finished). Default: ('b','g','r')
          :param xlim Optional x axis limits (usually a tuple of (xmin,xmax)).
          :param ylim Optional y axis limits (usually a tuple of (ymin,ymax)).
          :return: The matplotlib Figure object.
          """
          fig = plt.figure(figsize=self._figsize, dpi=self._dpi)
          plt.axis(np.ravel(self.boundaries, 'f'))
          if not plot_axis:
              plt.axis('off')
          else:
              plt.ylabel("Y position")
              plt.xlabel("X position")
          plt.plot([], 'b')
          plt.plot([], 'g')
          plt.title('Agent Trails')
          if plot_legend:
              plt.legend(['Active', 'Finished'])
          plt.tight_layout(pad=0)
          for agent in self.agents:
              if agent.status == 1:
                  alpha = 1
                  colour = colours[0]
              elif agent.status == 2:
                  alpha = .5
                  colour = colours[1]
              else:
                  alpha = 1
                  colour = colours[2]
              locs = np.array(agent.history_locations).T
              plt.plot(*locs, color=colour, alpha=alpha, linewidth=.5)
          if xlim != None: # Optionally set the x limits
              plt.xlim(xlim)
          if ylim != None:  # Optionally set the x limits
              plt.xlim(ylim)
          return fig
  
      def get_histogram(self):
          fig = plt.figure(figsize=self._figsize, dpi=self._dpi)
          fmax = max(np.amax(self.steps_exped), np.amax(self.steps_taken), np.amax(self.steps_delay))
          sround = lambda x, p: float(f'%.{p-1}e'%x)
          bins = np.linspace(0, sround(fmax, 2), 20)
          plt.hist(self.steps_exped, bins=bins+4, alpha=.5, label='Expected')
          plt.hist(self.steps_taken, bins=bins+2, alpha=.5, label='Taken')
          plt.hist(self.steps_delay, bins=bins+0, alpha=.5, label='Delayed')
          plt.xlabel('Time')
          plt.ylabel('Number of Agents')
          plt.grid(False)
          plt.legend()
          plt.tight_layout(pad=0)
          return fig
  
      @staticmethod
      def _heightmap(data, ax=None, kdeplot=True, cmap=None, alpha=.7, cbar=False):
          if kdeplot:
              from seaborn import kdeplot as sns_kdeplot
              sns_kdeplot(*data, ax=ax, cmap=cmap, alpha=alpha, shade=True, shade_lowest=False, cbar=cbar)
          else:
              hdata, binx, biny = np.histogram2d(*data, (20, 10))
              ax.contourf(hdata.T, cmap=cmap, alpha=alpha, extend='min', extent=(binx[0],binx[-1],biny[0],biny[-1]))
          return ax
  
      def get_wiggle_map(self, do_kdeplot=True, title="Collision Map"):
          """ Show where wiggles and collisions took place
  
          :param do_kdeplot:
          :param title: (optional) title for the graph
          :return: The figure object
          """
          fig, ax = plt.subplots(1, figsize=self._figsize, dpi=self._dpi)
          fig.tight_layout(pad=0)
          self._heightmap(np.array(self.history_collision_locs).T, ax=ax, kdeplot=do_kdeplot)
          self._heightmap(np.array(self.history_wiggle_locs).T, ax=ax)
          ax.set(frame_on=False, aspect='equal', xlim=self.boundaries[:,0], xticks=[],
                 ylim=self.boundaries[:,1], yticks=[], title=title)
          return fig
  
      def get_collision_map(self, *args, **kwargs):
          """For making a map of collisions and wiggles. Just calls get_wiggle_map()"""
          self.get_wiggle_map(*args, **kwargs)
  
      def get_location_map(self, do_kdeplot=True, title="Location Map", color_bar = False, plot_axis=False):
          """
          Create a density plot of the agents' locations
  
          :param do_kdeplot:
          :param title: (optional) title for the plot
          :return:
          """
          history_locs = []
          for agent in self.agents:
              for loc in agent.history_locations:
                  if None not in loc:
                      history_locs.append(loc)
          history_locs = np.array(history_locs).T
          fig, ax = plt.subplots(1, figsize=self._figsize, dpi=self._dpi)
          fig.tight_layout(pad=0)
          self._heightmap(data=history_locs, ax=ax, kdeplot=do_kdeplot, cmap='gray_r', cbar=color_bar)
          ax.set(frame_on=plot_axis, aspect='equal', xlim=self.boundaries[:,0], xticks=[],
                 ylim=self.boundaries[:,1], yticks=[], title=title)
          if plot_axis:
              ax.set_ylabel("Y position")
              ax.set_xlabel("X position")
          return fig
  
      def get_ani(self, agents=None, colour='k', alpha=.5, show_separation=False, wiggle_map=False):
          # Load Data
          locs = np.array([agent.history_locations for agent in self.agents[:agents]]).transpose((1,2,0))
          markersize = self.separation * 216*self._rel  # 3*72px/in=216
          #
          fig, ax = plt.subplots(figsize=self._figsize, dpi=self._dpi)
          if wiggle_map:
              sns.kdeplot(*np.array(self.collision_map).T, ax=ax, cmap='gray_r', alpha=.3, shade=True, shade_lowest=False)
          ln0, = plt.plot([], [], '.', alpha=.05, color=colour, markersize=markersize)
          ln1, = plt.plot([], [], '.', alpha=alpha, color=colour)
          def init():
              fig.tight_layout(pad=0)
              ax.set(frame_on=False, aspect='equal', xlim=self.boundaries[:,0], xticks=[], ylim=self.boundaries[:,1], yticks=[])
              return ln0, ln1,
          def func(frame):
              if show_separation:
                  ln0.set_data(*locs[frame])
              ln1.set_data(*locs[frame])
              return ln0, ln1,
          frames = self.step_id
          ani = FuncAnimation(fig, func, frames, init, interval=100, blit=True)
          return ani
  
      @classmethod
      def set_random_seed(cls, seed=None):
          """Set a new numpy random seed
          :param seed: the optional seed value (if None then get one from os.urandom)
          """
          new_seed = int.from_bytes(os.urandom(4), byteorder='little') if seed == None else seed
          np.random.seed(new_seed)
  
  if __name__ == '__main__':
      warnings.warn("The stationsim_model.py code should not be run directly. Create a separate script and use that "
                    "to run experimets (e.g. see ABM_DA/experiments/StationSim basic experiment.ipynb )")
      print("Nothing to do")
  