text: import argparse
  import gc
  import json
  import os
  import pickle
  import pprint
  import sys
  import time
  from datetime import datetime
  from timeit import default_timer as timer
  from pathlib import Path
  from pdb import set_trace as TT
  from random import randint
  
  import cv2
  from typing import Tuple
  
  import gym
  import matplotlib
  import matplotlib.pyplot as plt
  import numpy as np
  import pandas as pd
  import psutil
  import ray
  import scipy
  import torch as th
  import torch.nn.functional as F
  from gym import envs
  from numba import njit
  from qdpy.phenotype import Fitness, Features
  from ribs.archives import GridArchive
  from ribs.archives._add_status import AddStatus
  from ribs.emitters import (
      GradientImprovementEmitter,
      ImprovementEmitter,
      OptimizingEmitter,
  )
  from ribs.emitters.opt import CMAEvolutionStrategy
  from ribs.optimizers import Optimizer
  from ribs.visualize import grid_archive_heatmap
  from torch import ByteTensor, Tensor, nn
  from torch.nn import Conv2d, CrossEntropyLoss, Linear
  from torch.utils.tensorboard import SummaryWriter
  import deap
  import deap.tools
  import deap.algorithms
  import qdpy
  from qdpy import algorithms, containers, benchmarks, plots, tools
  from deap.base import Toolbox
  import graphviz
  import warnings
  import copy
  
  
  # Use for .py file
  from tqdm import tqdm
  
  import gym_pcgrl
  from evo_args import get_args
  from gym_pcgrl.envs.helper import get_int_prob, get_string_map
  
  # from example_play_call import random_player
  # gvgai_path = '/home/sme/GVGAI_GYM/'
  # sys.path.insert(0,gvgai_path)
  # from play import play
  
  # Use for notebook
  # from tqdm.notebook import tqdm
  
  # Use print to confirm access to local pcgrl gym
  # print([env.id for env in envs.registry.all() if "gym_pcgrl" in env.entry_point])
  """
  /// Required Environment ///
  conda create -n ribs-pt python=3.7
  pip install scipy==1.2.0  # must use this version with GVGAI_GYM
  conda install -c conda-forge notebook
  conda install pytorch torchvision torchaudio -c pyth
  conda install tensorboard
  pip install 'ribs[all]' gym~=0.17.0 Box2D~=2.3.10 tqdm
  git clone https://github.com/amidos2006/gym-pcgrl.git
  cd gym-pcgrl  # Must run in project root folder for access to pcgrl modules
  
  /// Instructions ///
  To start TensorBoard run the following command:
  $ tensorboard --logdir=runs
  
  Then go to:
  http://localhost:6006
  
  /// Resources ///
  
  Sam's example code:
  https://github.com/smearle/gol-cmame/blob/master/gol_cmame.py
  
  PCGRL Repo:
  https://github.com/amidos2006/gym-pcgrl
  
  Neural CA Paper:
  https://arxiv.org/pdf/2009.01398.pdf
  
  RIBS examples:
  https://docs.pyribs.org/en/stable/tutorials/lunar_lander.html
  """
  TARGETS_PENALTY_WEIGHT = 10
  
  def draw_net(config: object, genome: object, view: object = False, filename: object = None, node_names: object = None, show_disabled: object = True,
               prune_unused: object = False,
               node_colors: object = None, fmt: object = 'svg') -> object:
      """ Receives a genome and draws a neural network with arbitrary topology. """
      # Attributes for network nodes.
      if graphviz is None:
          warnings.warn("This display is not available due to a missing optional dependency (graphviz)")
          return
  
      if node_names is None:
          node_names = {}
  
      assert type(node_names) is dict
  
      if node_colors is None:
          node_colors = {}
  
      assert type(node_colors) is dict
  
      node_attrs = {
          'shape': 'circle',
          'fontsize': '9',
          'height': '0.2',
          'width': '0.2'}
  
      dot = graphviz.Digraph(format=fmt, node_attr=node_attrs)
  
      inputs = set()
      for k in config.genome_config.input_keys:
          inputs.add(k)
          name = node_names.get(k, str(k))
          input_attrs = {'style': 'filled', 'shape': 'box', 'fillcolor': node_colors.get(k, 'lightgray')}
          dot.node(name, _attributes=input_attrs)
  
      outputs = set()
      for k in config.genome_config.output_keys:
          outputs.add(k)
          name = node_names.get(k, str(k))
          node_attrs = {'style': 'filled', 'fillcolor': node_colors.get(k, 'lightblue')}
  
          dot.node(name, _attributes=node_attrs)
  
      if prune_unused:
          connections = set()
          for cg in genome.connections.values():
              if cg.enabled or show_disabled:
                  connections.add(cg.key)
  
          used_nodes = copy.copy(outputs)
          pending = copy.copy(outputs)
          while pending:
              new_pending = set()
              for a, b in connections:
                  if b in pending and a not in used_nodes:
                      new_pending.add(a)
                      used_nodes.add(a)
              pending = new_pending
      else:
          used_nodes = set(genome.nodes.keys())
  
      for n in used_nodes:
          if n in inputs or n in outputs:
              continue
  
          attrs = {'style': 'filled', 'fillcolor': node_colors.get(n, 'white')}
          dot.node(str(n), _attributes=attrs)
  
      for cg in genome.connections.values():
          if cg.enabled or show_disabled:
              #if cg.input not in used_nodes or cg.output not in used_nodes:
              #    continue
              input, output = cg.key
              a = node_names.get(input, str(input))
              b = node_names.get(output, str(output))
              style = 'solid' if cg.enabled else 'dotted'
              color = 'green' if cg.weight > 0 else 'red'
              width = str(0.1 + abs(cg.weight / 5.0))
              dot.edge(a, b, _attributes={'style': style, 'color': color, 'penwidth': width})
  
      dot.render(filename, view=view)
  
      return dot
  
  def save_level_frames(level_frames, model_name):
  
      renders_dir = os.path.join(SAVE_PATH, "renders")
      if not os.path.isdir(renders_dir):
          os.mkdir(renders_dir)
      model_dir = os.path.join(renders_dir, "model_{}".format(model_name))
      if not os.path.isdir(model_dir):
          os.mkdir(model_dir)
      for j, im in enumerate(level_frames):
          im.save(
              os.path.join(
                  model_dir, "frame_{:0>4d}.png".format(j)
              )
          )
  
  
  def get_qd_score(archive, env, bc_names):
      max_loss = env.get_max_loss(ctrl_metrics=bc_names)
      max_loss = max_loss * TARGETS_PENALTY_WEIGHT
      if ALGO == 'ME':
          # qd_score = archive.qd_score()  # we need to specify lower *and upper* bounds for this
          # TODO: work out max diversity bonus to make this possible ?? Would this bias scores between n. latent seeds
          #   though?
          qd_score = np.nansum(archive.quality_array + max_loss)
      else:
          df = archive.as_pandas(include_solutions=False)
          qd_score = (df['objective'] + max_loss).sum()
      return qd_score
  
  
  def save_train_stats(objs, archive, env, bc_names, itr=None):
      train_time_stats = {
              "qd_score": get_qd_score(archive, env, bc_names),
              "objective": get_stats(objs),
              }
  
      if itr is not None:
          save_path = os.path.join(SAVE_PATH, "checkpoint_{}".format(itr))
      else:
          save_path = SAVE_PATH
      json.dump(
          train_time_stats,
          open(os.path.join(save_path, "train_time_stats.json"), "w"),
          indent=4,
      )
  
  
  def get_stats(stats):
      """Take 1D numpy array of data and return some fun facts in the form of a dictionary."""
  
      return {
          "mean": np.nanmean(stats),
          "std": np.nanstd(stats),
          "max": np.nanmax(stats),
          "min": np.nanmin(stats),
      }
  
  
  def save_grid(csv_name="levels", d=4):
      fontsize = 32
      if "zelda" in PROBLEM:
          d = 3
          fontsize = int(fontsize * d / 4)
      elif "smb" in PROBLEM:
          d = 4
      if CMAES:
          # TODO: implement me
          return
      # save grid using csv file
      # get path to CSV
      levels_path = os.path.join(SAVE_PATH, csv_name + ".csv")
      # get env name
      env_name = "{}-{}-v0".format(PROBLEM, REPRESENTATION)
      # create env
      env = gym.make(env_name)
      map_width = env._prob._width
  
      df = pd.read_csv(levels_path, header=0, skipinitialspace=True)
      #   .rename(
  #       index=str,
  #       header=0,
  #       columns={
  #           0: "level",
  #           1: "batch_reward",
  #           2: "variance",
  #           3: "diversity",
  #           4: "targets",
  #       },
  #   )
  
      bc_names = []
      for i in range(5, 7):  # assume 2 BCs
          bc_names.append(df.columns[i])
      # look for the most valid levels
      targets_thresh = 0.0
      og_df = df
      df = og_df[og_df['targets'] == targets_thresh]
      last_len = len(df)
      while len(df) < d**2 and targets_thresh > og_df['targets'].min():
          last_len = len(df)
          # Raise the threshold so it includes at least one more individual
          targets_thresh = og_df[og_df['targets'] < targets_thresh]['targets'].max()
          df = og_df[og_df['targets'] >= targets_thresh]
      # d = 6  # dimension of rows and columns
      figw, figh = 16.0, 16.0
      fig = plt.figure()
      fig, axs = plt.subplots(ncols=d, nrows=d, figsize=(figw, figh))
  
      df_g = df.sort_values(by=bc_names, ascending=False)
  
      df_g["row"] = np.floor(np.linspace(0, d, len(df_g), endpoint=False)).astype(int)
  
      for row_num in range(d):
          row = df_g[df_g["row"] == row_num]
          row = row.sort_values(by=[bc_names[1]], ascending=True)
          row["col"] = np.arange(0, len(row), dtype=int)
          idx = np.floor(np.linspace(0, len(row) - 1, d)).astype(int)
          row = row[row["col"].isin(idx)]
          row = row.drop(["row", "col"], axis=1)
          # grid_models = np.array(row.loc[:,'solution_0':])
          grid_models = row["level"].tolist()
  
          for col_num in range(len(row)):
              axs[row_num, col_num].set_axis_off()
              level = np.zeros((map_width, map_width), dtype=int)
  
              for i, l_rows in enumerate(grid_models[col_num].split("], [")):
                  for j, l_col in enumerate(l_rows.split(",")):
                      level[i, j] = int(
                          l_col.replace("[", "").replace("]", "").replace(" ", "")
                      )
  
              # Set map
              env._rep._x = env._rep._y = 0
              env._rep._map = level
              img = env.render(mode="rgb_array")
  #           axs[row_num, col_num].imshow(img, aspect="auto")
              axs[-col_num-1, -row_num-1].imshow(img, aspect="auto")
  
      fig.subplots_adjust(hspace=0.01, wspace=0.01)
      levels_png_path = os.path.join(SAVE_PATH, "{}_grid.png".format(csv_name))
      fig.text(0.5, 0.01, bc_names[0], ha='center', va='center',fontsize=fontsize)
      fig.text(0.01, 0.5, bc_names[1], ha='center', va='center', rotation='vertical', fontsize=fontsize)
      plt.tight_layout(rect=[0.025, 0.025, 1, 1])
      fig.savefig(levels_png_path, dpi=300)
      plt.close()
  
  
  def auto_garbage_collect(pct=80.0):
      if psutil.virtual_memory().percent >= pct:
          gc.collect()
  
  def tran_action(action, **kwargs):
      skip = False
      # return action, skip
      return action.swapaxes(1, 2), skip
  
  # usually, if action does not turn out to change the map, then the episode is terminated
  # the skip boolean tells us whether, for some representation-specific reason, the agent has chosen not to act, but
  # without ending the episode
  @njit
  def id_action(action, int_map=None, n_tiles=None, x=None, y=None, n_dirs=None):
      # the argmax along tile_type dimension is performed inside the representation's update function
      skip = False
  
      return action, skip
  
  
  # @njit
  def wide_action(action, int_map=None, n_tiles=None, x=None, y=None, n_dirs=None):
      # only consider tiles where the generator suggests something different than the existing tile
      act_mask = action.argmax(axis=0) != int_map
      n_new_builds = np.sum(act_mask)
      act_mask = act_mask.reshape((1, *act_mask.shape))
      #   action = action * act_mask
      action = np.where(act_mask == False, action.min() - 10, action)
      coords = np.unravel_index(action.argmax(), action.shape)
  
      if n_new_builds > 0:
          assert act_mask[0, coords[1], coords[2]] == 1
      coords = coords[2], coords[1], coords[0]
      #   assert int_map[coords[0], coords[1]] != coords[2]
      skip = False
  
      return coords, skip
  
  
  @njit
  def narrow_action(action, int_map=None, n_tiles=None, x=None, y=None, n_dirs=None):
      act = action[:, y, x].argmax()
  
      if act == 0:
          skip = True
      else:
          skip = False
  
      return act, skip
  
  
  @njit
  def turtle_action(action, int_map=None, n_tiles=None, x=None, y=None, n_dirs=None):
      act = action[:, y, x].argmax()
      # moving is counted as a skip, so lack of change does not end episode
  
      if act < n_dirs:
          skip = True
      else:
          skip = False
  
      return act, skip
  
  
  @njit
  def flat_to_box(action, int_map=None, n_tiles=None, x=None, y=None, n_dirs=None):
      action = action.reshape((n_tiles, *int_map.shape))
      skip = False
  
      return action, skip
  
  
  @njit
  def flat_to_wide(action, int_map=None, n_tiles=None, x=None, y=None, n_dirs=None):
      w = int_map.shape[0]
      h = int_map.shape[1]
      assert len(action) == int_map.shape[0] + int_map.shape[1] + n_tiles
      action = (action[:w].argmax(), action[w : w + h].argmax(), action[w + h :].argmax())
      skip = False
  
      return action, skip
  
  
  @njit
  def flat_to_narrow(action, int_map=None, n_tiles=None, x=None, y=None, n_dirs=None):
      act = action.argmax()
  
      if act == 0:
          skip = True
      else:
          skip = False
  
      return act, skip
  
  
  @njit
  def flat_to_turtle(action, int_map=None, n_tiles=None, x=None, y=None, n_dirs=None):
      act = action.argmax()
  
      if act < n_dirs:
          skip = True
      else:
          skip = False
  
      return act, skip
  
  
  preprocess_action_funcs = {
      "NCA": {
          "cellular": id_action,
          "wide": wide_action,
          "narrow": narrow_action,
          "turtle": turtle_action,
      },
      "CPPN": {
          "cellular": tran_action,
      },
      "CNN": {
          # will try to build this logic into the model
          "cellular": flat_to_box,
          "wide": flat_to_wide,
          "narrow": flat_to_narrow,
          "turtle": flat_to_turtle,
      },
  }
  
  
  def id_observation(obs, **kwargs):
      return obs
  
  
  def local_observation(obs, **kwargs):
      x = kwargs.get("x")
      y = kwargs.get("y")
      local_obs = np.zeros((1, obs.shape[1], obs.shape[2]))
      # Might be some inconsistencies in ordering of x, y?
      local_obs[0, y, x] = 1
      np.concatenate((obs, local_obs), axis=0)
  
      return obs
  
  
  preprocess_observation_funcs = {
      "NCA": {
          "cellular": id_observation,
          "wide": id_observation,
          "narrow": local_observation,
          "turtle": local_observation,
      },
      "CNN": {
          "cellular": id_observation,
          "wide": id_observation,
          "narrow": local_observation,
          "turtle": local_observation,
      },
  }
  
  @njit
  def archive_init_states(init_states_archive, init_states, index):
      init_states_archive[index] = init_states
  
  
  # @njit
  def get_init_states(init_states_archive, index):
      return init_states_archive[index]
  
  
  
  def mate_individuals(ind_0, ind_1):
      return ind_0.mate(ind_1)
  
  def mutate_individual(ind):
      ind.mutate()
      return (ind,)
  
  class MEOptimizer():
      def __init__(self, grid, ind_cls, batch_size, ind_cls_args, start_time=None, stats=None):
          self.batch_size = batch_size
          self.grid = grid
          self.inds = []
          self.stats=stats
          for _ in range(batch_size):
              self.inds.append(ind_cls(**ind_cls_args))
          toolbox = Toolbox()
          toolbox.register("clone", copy.deepcopy)
          toolbox.register("mutate", mutate_individual)
          toolbox.register("mate", mate_individuals)
          toolbox.register("select", tools.sel_random)
  
          self.cxpb = 0
          self.mutpb = 1.0
          self.toolbox = toolbox
          if start_time == None:
              self.start_time = timer()
          self.logbook = deap.tools.Logbook()
          self.logbook.header = ["iteration", "containerSize", "evals", "nbUpdated"] + (stats.fields if stats else []) + ["elapsed"]
          self.i = 0
  
  
      def tell(self, objective_values, behavior_values):
          # Update individuals' stats with results of last batch of simulations
  #       [(ind.fitness.setValues(obj), ind.fitness.features.setValues(bc)) for
  #        (ind, obj, bc) in zip(self.inds, objective_values, behavior_values)]
          for (ind, obj, bc) in zip(self.inds, objective_values, behavior_values):
              ind.fitness.setValues([obj])
              ind.features.setValues(bc)
          # Replace the current population by the offspring
          nb_updated = self.grid.update(self.inds, issue_warning=True, ignore_exceptions=False)
          # Compile stats and update logs
          record = self.stats.compile(self.grid) if self.stats else {}
          self.logbook.record(iteration=self.i, containerSize=self.grid.size_str(), evals=len(self.inds), nbUpdated=nb_updated, elapsed=timer()-self.start_time, **record)
          self.i += 1
          print(self.logbook.stream)
  
      def ask(self):
  
          if len(self.grid) == 0:
              # Return the initial batch
              return self.inds
  
          elif len(self.grid) < self.batch_size:
              # If few elites, supplement the population with individuals from the last generation
              np.random.shuffle(self.inds)
              breedable = self.grid.items + self.inds[:-len(self.grid)]
  
          else:
              breedable = self.grid
  
          # Select the next batch individuals
          batch = [self.toolbox.select(breedable) for i in range(self.batch_size)]
  
          ## Vary the pool of individuals
          self.inds = deap.algorithms.varAnd(batch, self.toolbox, self.cxpb, self.mutpb)
  
          return self.inds
  
  
  
  class InitStatesArchive(GridArchive):
      """Save (some of) the initial states upon which the elites were evaluated when added to the archive, so that we can
      reproduce their behavior at evaluation time (and compare it to evaluation to other seeds)."""
  
      def __init__(self, bin_sizes, bin_bounds, n_init_states, map_w, map_h, **kwargs):
          super(InitStatesArchive, self).__init__(bin_sizes, bin_bounds, **kwargs)
          self.init_states_archive = np.empty(
              shape=(*bin_sizes, n_init_states, map_w, map_h)
          )
  
      def set_init_states(self, init_states):
          self.init_states = init_states
  
      def add(self, solution, objective_value, behavior_values, meta, index=None):
          status, dtype_improvement = super().add(
              solution, objective_value, behavior_values
          )
  
          # NOTE: for now we won't delete these when popping an elite for re-evaluation
  
          if status != AddStatus.NOT_ADDED:
              if index is None:
                  index = self.get_index(behavior_values)
              archive_init_states(self.init_states_archive, self.init_states, index)
  
          return status, dtype_improvement
  
  
  class MEGrid(containers.Grid):
      def __init__(self, bin_sizes, bin_bounds):
          super(MEGrid, self).__init__(shape=bin_sizes, max_items_per_bin=1,
                                       features_domain=bin_bounds,
                                       fitness_domain=((-np.inf, np.inf),),
                                       )
  
          # pyribs compatibility
          def get_index(self, bcs):
              return self.index_grid(features=bcs)
  
      def add(self, item):
          # We'll clip the feature calues at the extremes
          # TODO: what's happening in this case using pyribs?
          item.features.setValues([np.clip(item.features.values[i], *self.features_domain[i])
                                   for i in range(len(item.features.values))])
          return super(MEGrid, self).add(item)
  
  
  class MEInitStatesArchive(MEGrid):
      """Save (some of) the initial states upon which the elites were evaluated when added to the archive, so that we can
      reproduce their behavior at evaluation time (and compare it to evaluation to other seeds)."""
  
      def __init__(self, bin_sizes, bin_bounds, n_init_states, map_w, map_h, **kwargs):
          super(MEInitStatesArchive, self).__init__(bin_sizes, bin_bounds, **kwargs)
          self.init_states_archive = np.empty(
              shape=(*bin_sizes, n_init_states, map_w, map_h)
          )
  
      def set_init_states(self, init_states):
          self.init_states = init_states
  
      def add(self, item):
          index = super(MEInitStatesArchive, self).add(item)
  
          if index is not None:
              idx = self.index_grid(item.features)
              archive_init_states(self.init_states_archive, self.init_states, idx)
  
          return index
  
  
  
  class FlexArchive(InitStatesArchive):
      """ Subclassing a pyribs archive class to do some funky stuff."""
  
      def __init__(self, *args, **kwargs):
          self.n_evals = {}
          #       self.obj_hist = {}
          #       self.bc_hist = {}
          super().__init__(*args, **kwargs)
          #       # "index of indices", so we can remove them from _occupied_indices when removing
          #       self._index_ranks = {}
          self._occupied_indices = set()
  
      def _add_occupied_index(self, index):
          #       rank = len(self._occupied_indices)
          #       self._index_ranks[index] = rank  # the index of the index in _occupied_indices
  
          return super()._add_occupied_index(index)
  
      def _remove_occupied_index(self, index):
          self._occupied_indices.remove(index)
          self._occupied_indices_cols = tuple(
              [self._occupied_indices[i][j] for i in range(len(self._occupied_indices))]
              for j in range(len(self._storage_dims))
          )
  
      def pop_elite(self, obj, bcs, old_bcs):
          """
          Need to call update_elite after this!
          """
          # Remove it, update it
          old_idx = self.get_index(np.array(old_bcs))
          self._remove_occupied_index(old_idx)
  
          #       rank = self._index_ranks.pop(old_idx)
          #       self._occupied_indices.pop(rank)
          #       [self._occupied_indices_cols[i].pop(rank) for i in range(len(self._storage_dims))]
          n_evals = self.n_evals.pop(old_idx)
          old_obj = self._objective_values[old_idx]
          mean_obj = (old_obj * n_evals + obj) / (n_evals + 1)
          mean_bcs = np.array(
              [
                  (old_bcs[i] * n_evals + bcs[i]) / (n_evals + 1)
                  for i in range(len(old_bcs))
              ]
          )
          #       obj_hist = self.obj_hist.pop(old_idx)
          #       obj_hist.append(obj)
          #       mean_obj = np.mean(obj_hist)
          #       bc_hist = self.bc_hist.pop(old_idx)
          #       bc_hist.append(bcs)
          #       bc_hist_np = np.asarray(bc_hist)
          #       mean_bcs = bc_hist_np.mean(axis=0)
          self._objective_values[old_idx] = np.nan
          self._behavior_values[old_idx] = np.nan
          self._occupied[old_idx] = False
          solution = self._solutions[old_idx].copy()
          self._solutions[old_idx] = np.nan
          self._metadata[old_idx] = np.nan
          #       while len(obj_hist) > 100:
          #           obj_hist = obj_hist[-100:]
          #       while len(bc_hist) > 100:
          #           bc_hist = bc_hist[-100:]
  
          return solution, mean_obj, mean_bcs, n_evals
  
      def update_elite(self, solution, mean_obj, mean_bcs, n_evals):
          """
          obj: objective score from new evaluations
          bcs: behavior characteristics from new evaluations
          old_bcs: previous behavior characteristics, for getting the individuals index in the archive
          """
  
          # Add it back
  
          self.add(solution, mean_obj, mean_bcs, None, n_evals=n_evals)
  
  
      def add(self, solution, objective_value, behavior_values, meta, n_evals=0):
  
          index = self.get_index(behavior_values)
  
          status, dtype_improvement = super().add(
              solution, objective_value, behavior_values, meta, index
          )
  
          if not status == AddStatus.NOT_ADDED:
              if n_evals == 0:
                  self.n_evals[index] = 1
              else:
                  self.n_evals[index] = min(n_evals + 1, 100)
  
          return status, dtype_improvement
  
  
  def unravel_index(
      indices: th.LongTensor, shape: Tuple[int, ...]
  ) -> th.LongTensor:
      r"""Converts flat indices into unraveled coordinates in a target shape.
  
      This is a `th` implementation of `numpy.unravel_index`.
  
      Args:
          indices: A tensor of indices, (*, N).
          shape: The targeted shape, (D,).
  
      Returns:
          unravel coordinates, (*, N, D).
      """
  
      shape = th.tensor(shape)
      indices = indices % shape.prod()  # prevent out-of-bounds indices
  
      coord = th.zeros(indices.size() + shape.size(), dtype=int)
  
      for i, dim in enumerate(reversed(shape)):
          coord[..., i] = indices % dim
          indices = indices // dim
  
      return coord.flip(-1)
  
  
  # TODO: Use the GPU!
  
  #   if CUDA:
  #       m.cuda()
  #       m.to('cuda:0')
  
  class ResettableNN(nn.Module):
  
      def reset(self):
          pass
  
  
  def gauss(x, mean=0, std=1):
      return th.exp((-(x - mean) ** 2)/(2 * std ** 2))
  
  
  class MixActiv(nn.Module):
      def __init__(self):
          super().__init__()
          self.activations = (th.sin, th.tanh, gauss, th.relu)
          self.n_activs = len(self.activations)
  
  
      def forward(self, x):
          n_chan = x.shape[1]
          chans_per_activ = n_chan / self.n_activs
          chan_i = 0
          xs = []
          for i, activ in enumerate(self.activations):
              xs.append(activ(x[:, int(chan_i):int(chan_i + chans_per_activ), :, :]))
              chan_i += chans_per_activ
          x = th.cat(xs, axis=1)
          return x
  
  
  class AuxNCA(ResettableNN):
      def __init__(self, n_in_chans, n_actions, n_aux=3):
          super().__init__()
          self.n_hid_1 = n_hid_1 = 32
          self.n_aux = n_aux
          self.l1 = Conv2d(n_in_chans + self.n_aux, n_hid_1, 3, 1, 1, bias=True)
          self.l2 = Conv2d(n_hid_1, n_hid_1, 1, 1, 0, bias=True)
          self.l3 = Conv2d(n_hid_1, n_actions + self.n_aux, 1, 1, 0, bias=True)
          self.layers = [self.l1, self.l2, self.l3]
          self.apply(init_weights)
          self.last_aux = None
          if RENDER:
              cv2.namedWindow("Auxiliary NCA")
  
      def forward(self, x):
          with th.no_grad():
              if self.last_aux is None:
                  self.last_aux = th.zeros(size=(1, self.n_aux, *x.shape[-2:]))
              x_in = th.cat([x, self.last_aux], axis=1)
              x = self.l1(x_in)
              x = th.nn.functional.relu(x)
              x = self.l2(x)
              x = th.nn.functional.relu(x)
              x = self.l3(x)
              x = th.sigmoid(x)
              self.last_aux = x[:,-self.n_aux:,:,:]
              x = x[:, :-self.n_aux,:,:]
              if RENDER:
  #               im = self.last_aux[0].cpu().numpy().transpose(1,2,0)
                  aux = self.last_aux[0].cpu().numpy()
                  aux = aux / aux.max()
                  im = np.expand_dims(np.vstack(aux), axis=0)
                  im = im.transpose(1, 2, 0)
                  cv2.imshow("Auxiliary NCA", im)
                  cv2.waitKey(1)
  
          # axis 0 is batch
          # axis 1 is the tile-type (one-hot)
          # axis 0,1 is the x value
          # axis 0,2 is the y value
  
          return x, False
  
      def reset(self, init_aux=None):
          self.last_aux = None
  
  
  class DoneAuxNCA(AuxNCA):
      def __init__(self, n_in_chans, n_actions, n_aux=3):
          # Add an extra auxiliary ("done") channel after the others
          n_aux += 1
          super().__init__(n_in_chans, n_actions, n_aux=n_aux)
          done_kernel_size = 3
          self.l_done = Conv2d(1, 1, 7, stride=999)
  
      def forward(self, x):
          with th.no_grad():
              x, done = super().forward(x)
              # retrieve local activation from done channel
              done_x = th.sigmoid(self.l_done(x[:,-1:,:,:])).flatten() - 0.5
              done = (done_x > 0).item()
  
          return x, done
  
      def reset(self, init_aux=None):
          self.last_aux = None
  
  
  class GeneratorNN(ResettableNN):
  #class NCA(ResettableNN):
      """ A neural cellular automata-type NN to generate levels or wide-representation action distributions."""
  
      def __init__(self, n_in_chans, n_actions, **kwargs):
          super().__init__()
          n_hid_1 = 32
          self.l1 = Conv2d(n_in_chans, n_hid_1, 3, 1, 1, bias=True)
          self.l2 = Conv2d(n_hid_1, n_hid_1, 1, 1, 0, bias=True)
          self.l3 = Conv2d(n_hid_1, n_actions, 1, 1, 0, bias=True)
          self.layers = [self.l1, self.l2, self.l3]
          self.apply(init_weights)
  
      def forward(self, x):
          with th.no_grad():
              x = self.l1(x)
              x = th.nn.functional.relu(x)
              x = self.l2(x)
              x = th.nn.functional.relu(x)
              x = self.l3(x)
              x = th.sigmoid(x)
  
  
          # axis 0 is batch
          # axis 1 is the tile-type (one-hot)
          # axis 0,1 is the x value
          # axis 0,2 is the y value
  
          return x, False
  
  class MixNCA(ResettableNN):
      def __init__(self, *args, **kwargs):
          super(MixNCA, self).__init__()
          self.mix_activ = MixActiv()
  
      def forward(self, x):
          with th.no_grad():
              x = self.l1(x)
              x = self.mix_activ(x)
              x = self.l2(x)
              x = self.mix_activ(x)
              x = self.l3(x)
              x = th.sigmoid(x)
  
  
  class CoordNCA(ResettableNN):
      """ A neural cellular automata-type NN to generate levels or wide-representation action distributions.
      With coordinates as additional input, like a CPPN."""
  
      def __init__(self, n_in_chans, n_actions):
          super().__init__()
          n_hid_1 = 28
          #       n_hid_2 = 16
  
          self.l1 = Conv2d(n_in_chans + 2, n_hid_1, 3, 1, 1, bias=True)
          self.l2 = Conv2d(n_hid_1, n_hid_1, 1, 1, 0, bias=True)
          self.l3 = Conv2d(n_hid_1, n_actions, 1, 1, 0, bias=True)
          self.layers = [self.l1, self.l2, self.l3]
          self.apply(init_weights)
  
      def forward(self, x):
          with th.no_grad():
              coords = get_coord_grid(x, normalize=True)
              x = th.hstack((coords, x))
              x = self.l1(x)
              x = th.nn.functional.relu(x)
              x = self.l2(x)
              x = th.nn.functional.relu(x)
              x = self.l3(x)
              x = th.sigmoid(x)
  
          # axis 0 is batch
          # axis 1 is the tile-type (one-hot)
          # axis 0,1 is the x value
          # axis 0,2 is the y value
  
          return x, False
  
  
  from pytorch_neat.cppn import create_cppn, Leaf
  import neat
  from neat.genome import DefaultGenome
  
  def get_coord_grid(x, normalize=False):
      width = x.shape[-2]
      height = x.shape[-1]
      X = th.arange(width)
      Y = th.arange(height)
      if normalize:
          X = X / width
          Y = Y / height
      else:
          X = X / 1
          Y = Y / 1
      X, Y = th.meshgrid(X, Y)
      x = th.stack((X, Y)).unsqueeze(0)
  
      return x
  
  
  #class ReluCPPN(ResettableNN):
  class FeedForwardCPPN(nn.Module):
      def __init__(self, n_in_chans, n_actions):
          super().__init__()
          n_hid = 64
          self.l1 = Conv2d(2, n_hid, kernel_size=1)
          self.l2 = Conv2d(n_hid, n_hid, kernel_size=1)
          self.l3 = Conv2d(n_hid, n_actions, kernel_size=1)
          self.layers = [self.l1, self.l2, self.l3]
          self.apply(init_weights)
  
      def forward(self, x):
          x = get_coord_grid(x, normalize=True)
          with th.no_grad():
              x = th.relu(self.l1(x))
              x = th.relu(self.l2(x))
              x = th.sigmoid(self.l3(x))
  
          return x, True
  
  
  class GenReluCPPN(ResettableNN):
      def __init__(self, n_in_chans, n_actions):
          super().__init__()
          n_hid = 64
          self.l1 = Conv2d(2+n_in_chans, n_hid, kernel_size=1)
          self.l2 = Conv2d(n_hid, n_hid, kernel_size=1)
          self.l3 = Conv2d(n_hid, n_actions, kernel_size=1)
          self.layers = [self.l1, self.l2, self.l3]
          self.apply(init_weights)
  
      def forward(self, x):
          coord_x = get_coord_grid(x, normalize=True)
          x = th.cat((x, coord_x), axis=1)
          with th.no_grad():
              x = th.relu(self.l1(x))
              x = th.relu(self.l2(x))
              x = th.sigmoid(self.l3(x))
  
          return x, True
  
  
  class SinCPPN(ResettableNN):
      def __init__(self, n_in_chans, n_actions):
          super().__init__()
          n_hid = 64
          self.l1 = Conv2d(2, n_hid, kernel_size=1)
          self.l2 = Conv2d(n_hid, n_hid, kernel_size=1)
          self.l3 = Conv2d(n_hid, n_actions, kernel_size=1)
          self.layers = [self.l1, self.l2, self.l3]
          self.apply(init_weights)
  
      def forward(self, x):
          x = get_coord_grid(x, normalize=True) * 2
          with th.no_grad():
              x = th.sin(self.l1(x))
              x = th.sin(self.l2(x))
              x = th.sigmoid(self.l3(x))
  
          return x, True
  
  
  class GenSinCPPN(ResettableNN):
      def __init__(self, n_in_chans, n_actions):
          super().__init__()
          n_hid = 64
          self.l1 = Conv2d(2+n_in_chans, n_hid, kernel_size=1)
          self.l2 = Conv2d(n_hid, n_hid, kernel_size=1)
          self.l3 = Conv2d(n_hid, n_actions, kernel_size=1)
          self.layers = [self.l1, self.l2, self.l3]
          self.apply(init_weights)
  
      def forward(self, x):
          coord_x = get_coord_grid(x, normalize=True) * 2
          x = th.cat((x, coord_x), axis=1)
          with th.no_grad():
              x = th.sin(self.l1(x))
              x = th.sin(self.l2(x))
              x = th.sigmoid(self.l3(x))
  
          return x, True
  
  
  class MixCPPN(ResettableNN):
      def __init__(self, n_in_chans, n_actions):
          super().__init__()
          n_hid = 64
          self.l1 = Conv2d(2, n_hid, kernel_size=1)
          self.l2 = Conv2d(n_hid, n_hid, kernel_size=1)
          self.l3 = Conv2d(n_hid, n_actions, kernel_size=1)
          self.layers = [self.l1, self.l2, self.l3]
          self.apply(init_weights)
          self.mix_activ = MixActiv()
  
  
      def forward(self, x):
          x = get_coord_grid(x, normalize=True) * 2
          with th.no_grad():
              x = self.mix_activ(self.l1(x))
              x = self.mix_activ(self.l2(x))
              x = th.sigmoid(self.l3(x))
  
          return x, True
  
  
  class GenMixCPPN(ResettableNN):
      def __init__(self, n_in_chans, n_actions):
          super().__init__()
          n_hid = 64
          self.l1 = Conv2d(2+n_in_chans, n_hid, kernel_size=1)
          self.l2 = Conv2d(n_hid, n_hid, kernel_size=1)
          self.l3 = Conv2d(n_hid, n_actions, kernel_size=1)
          self.layers = [self.l1, self.l2, self.l3]
          self.apply(init_weights)
          self.mix_activ = MixActiv()
  
  
      def forward(self, x):
          coord_x = get_coord_grid(x, normalize=True) * 2
          x = th.cat((x, coord_x), axis=1)
          with th.no_grad():
              x = self.mix_activ(self.l1(x))
              x = self.mix_activ(self.l2(x))
              x = th.sigmoid(self.l3(x))
  
          return x, True
  
  
  class FixedGenCPPN(ResettableNN):
      """A fixed-topology CPPN that takes additional channels of noisey input to prompts its output.
      Like a CoordNCA but without the repeated passes and with 1x1 rather than 3x3 kernels."""
      # TODO: Maybe try this with 3x3 conv, just to cover our bases?
      def __init__(self, n_in_chans, n_actions):
          super().__init__()
          n_hid = 64
          self.l1 = Conv2d(2 + n_in_chans, n_hid, kernel_size=1)
          self.l2 = Conv2d(n_hid, n_hid, kernel_size=1)
          self.l3 = Conv2d(n_hid, n_actions, kernel_size=1)
          self.layers = [self.l1, self.l2, self.l3]
          self.apply(init_weights)
  
      def forward(self, x):
          coord_x = get_coord_grid(x, normalize=True) * 2
          x = th.cat((x, coord_x), axis=1)
          with th.no_grad():
              x = th.sin(self.l1(x))
              x = th.sin(self.l2(x))
              x = th.sigmoid(self.l3(x))
  
          return x, True
  
  
  class CPPN(ResettableNN):
      def __init__(self, n_in_chans, n_actions):
          super().__init__()
          neat_config_path = 'config_cppn'
          self.neat_config = neat.config.Config(DefaultGenome, neat.reproduction.DefaultReproduction,
                                                neat.species.DefaultSpeciesSet, neat.stagnation.DefaultStagnation,
                                                neat_config_path)
          self.n_actions = n_actions
          self.neat_config.genome_config.num_outputs = n_actions
          self.neat_config.genome_config.num_hidden = 2
          self.genome = DefaultGenome(0)
          self.genome.configure_new(self.neat_config.genome_config)
          self.input_names = ['x_in', 'y_in']
          self.output_names = ['tile_{}'.format(i) for i in range(n_actions)]
          self.cppn = create_cppn(self.genome, self.neat_config, self.input_names, self.output_names)
  
      def mate(self, ind_1, fit_0, fit_1):
          self.genome.fitness = fit_0
          ind_1.genome.fitness = fit_1
          return self.genome.configure_crossover(self.genome, ind_1.genome, self.neat_config.genome_config)
  
      def mutate(self):
  #       print(self.input_names, self.neat_config.genome_config.input_keys, self.genome.nodes)
          self.genome.mutate(self.neat_config.genome_config)
          self.cppn = create_cppn(self.genome, self.neat_config, self.input_names, self.output_names)
  
      def draw_net(self):
          draw_net(self.neat_config, self.genome,  view=True, filename='cppn')
  
      def forward(self, x):
          X = th.arange(x.shape[-2])
          Y = th.arange(x.shape[-1])
          X, Y = th.meshgrid(X/X.max(), Y/Y.max())
          tile_probs = [self.cppn[i](x_in=X, y_in=Y) for i in range(self.n_actions)]
          multi_hot = th.stack(tile_probs, axis=0)
          multi_hot = multi_hot.unsqueeze(0)
          return multi_hot, True
  
  
  class CPPNCA(ResettableNN):
      def __init__(self, n_in_chans, n_actions):
          super().__init__()
          n_hid_1 = 32
          with th.no_grad():
              self.l1 = Conv2d(n_in_chans, n_hid_1, 3, 1, 1, bias=True)
              self.l2 = Conv2d(n_hid_1, n_hid_1, 1, 1, 0, bias=True)
              self.l3 = Conv2d(n_hid_1, n_actions, 1, 1, 0, bias=True)
          self.layers = [self.l1, self.l2, self.l3]
          self.apply(init_weights)
          n_nca_params = sum(p.numel() for p in self.parameters())
          self.cppn_body = GenCPPN(n_in_chans, n_actions)
          self.normal = th.distributions.multivariate_normal.MultivariateNormal(th.zeros(1), th.eye(1))
  
      def mate(self):
          raise NotImplementedError
  
      def mutate(self):
          self.cppn_body.mutate()
  
          with th.no_grad():
              for layer in self.layers:
                  dw = self.normal.sample(layer.weight.shape)
                  layer.weight = th.nn.Parameter(layer.weight + dw.squeeze(-1))
                  db = self.normal.sample(layer.bias.shape)
                  layer.bias = th.nn.Parameter(layer.bias + db.squeeze(-1))
  
      def forward(self, x):
          with th.no_grad():
              x = self.l1(x)
              x = th.nn.functional.relu(x)
              x = self.l2(x)
              x = th.nn.functional.relu(x)
              x = th.sigmoid(x)
          x, _ = self.cppn_body(x)
          return x, False
  
  
  class GenCPPN(CPPN):
      def __init__(self, n_in_chans, n_actions):
          super().__init__(n_in_chans, n_actions)
          neat_config_path = 'config_cppn'
          self.neat_config = neat.config.Config(DefaultGenome, neat.reproduction.DefaultReproduction,
                                                neat.species.DefaultSpeciesSet, neat.stagnation.DefaultStagnation,
                                                neat_config_path)
          self.n_actions = n_actions
          self.neat_config.genome_config.num_outputs = n_actions
          self.genome = DefaultGenome(0)
          self.input_names = ['x_in', 'y_in'] + ['tile_{}_in'.format(i) for i in range(n_actions)]
          n_inputs = len(self.input_names)
          self.output_names = ['tile_{}_out'.format(i) for i in range(n_actions)]
          self.neat_config.genome_config.input_keys = (-1*np.arange(n_inputs) - 1).tolist()
          self.neat_config.genome_config.num_inputs = n_inputs
          self.neat_config.genome_config.num_hidden = 2
          self.genome.configure_new(self.neat_config.genome_config)
          self.cppn = create_cppn(self.genome, self.neat_config, self.input_names, self.output_names)
  
      def forward(self, x):
          X = th.arange(x.shape[-2])
          Y = th.arange(x.shape[-1])
          X, Y = th.meshgrid(X/X.max(), Y/Y.max())
          inputs = {'x_in': X, 'y_in': Y}
          inputs.update({'tile_{}_in'.format(i): th.Tensor(x[0,i,:,:]) for i in range(self.n_actions)})
          tile_probs = [self.cppn[i](**inputs) for i in range(self.n_actions)]
          multi_hot = th.stack(tile_probs, axis=0)
          multi_hot = multi_hot.unsqueeze(0)
          return multi_hot, True
  
  
  class Individual(qdpy.phenotype.Individual):
      "An individual for mutating with operators. Assuming we're using vanilla MAP-Elites here."
      def __init__(self, model_cls, n_in_chans, n_actions):
          super(Individual, self).__init__()
          self.model = model_cls(n_in_chans, n_actions)
          self.fitness = Fitness([0])
          self.fitness.delValues()
  
      def mutate(self):
          self.model.mutate()
  
      def mate(self, ind_1):
          assert len(self.fitness.values) == 1 == len(ind_1.fitness.values)
          self.model.mate(ind_1.model, fit_0=self.fitness.values[0], fit_1=ind_1.fitness.values[0])
  
      def __eq__(self, ind_1):
          if not hasattr(ind_1, "model"): return False
          return self.model == ind_1.model
  
  
  # FIXME: this guy don't work
  class GeneratorNNDenseSqueeze(ResettableNN):
      """ A neural cellular automata-type NN to generate levels or wide-representation action distributions."""
  
      def __init__(self, n_in_chans, n_actions, observation_shape, n_flat_actions):
          super().__init__()
          n_hid_1 = 16
          # Hack af. Pad the input to make it have root 2? idk, bad
          sq_i = 2
          assert observation_shape[-1] == observation_shape[-2]
  
          #       while sq_i < observation_shape[-1]:
          #           sq_i = sq_i**2
          #       pad_0 = sq_i - observation_shape[-1]
          self.l1 = Conv2d(n_in_chans, n_hid_1, 3, 1, 0, bias=True)
          self.l2 = Conv2d(n_hid_1, n_hid_1, 3, 2, 0, bias=True)
          self.flatten = th.nn.Flatten()
          n_flat = self.flatten(
              self.l2(self.l1(th.zeros(size=observation_shape)))
          ).shape[-1]
          # n_flat = n_hid_1
          self.d1 = Linear(n_flat, n_flat_actions)
          #       self.d2 = Linear(16, n_flat_actions)
          self.layers = [self.l1, self.l2, self.d1]
          self.apply(init_weights)
  
      def forward(self, x):
          with th.no_grad():
              x = self.l1(x)
              x = th.nn.functional.relu(x)
              x = self.l2(x)
              x = th.nn.functional.relu(x)
              #           for i in range(int(np.log2(x.shape[2])) + 1):
              #               x = self.l2(x)
              #               x = th.nn.functional.relu(x)
              x = self.flatten(x)
              x = self.d1(x)
              x = th.sigmoid(x)
              #           x = self.d2(x)
              #           x = th.sigmoid(x)
  
          return x, False
  
  
  class GeneratorNNDense(ResettableNN):
      """ A neural cellular automata-type NN to generate levels or wide-representation action distributions."""
  
      def __init__(self, n_in_chans, n_actions, observation_shape, n_flat_actions):
          super().__init__()
          n_hid_1 = 16
          n_hid_2 = 32
          self.conv1 = Conv2d(n_in_chans, n_hid_1, kernel_size=3, stride=2)
          self.conv2 = Conv2d(n_hid_1, n_hid_2, kernel_size=3, stride=2)
          self.conv3 = Conv2d(n_hid_2, n_hid_2, kernel_size=3, stride=2)
          self.flatten = th.nn.Flatten()
          n_flat = self.flatten(
              self.conv3(self.conv2(self.conv1(th.zeros(size=observation_shape))))
          ).shape[-1]
          #       self.fc1 = Linear(n_flat, n_flat_actions)
          self.fc1 = Linear(n_flat, n_hid_2)
          self.fc2 = Linear(n_hid_2, n_flat_actions)
          self.layers = [self.conv1, self.conv2, self.conv3, self.fc1, self.fc2]
          self.apply(init_weights)
  
      def forward(self, x):
          with th.no_grad():
              x = F.relu(self.conv1(x))
              x = F.relu(self.conv2(x))
              x = F.relu(self.conv3(x))
              x = self.flatten(x)
              x = F.relu(self.fc1(x))
              x = F.softmax(self.fc2(x), dim=1)
  
          return x, False
  
  
  class PlayerNN(ResettableNN):
      def __init__(self, n_tile_types, n_actions=4):
          super().__init__()
          self.n_tile_types = n_tile_types
          assert "zelda" in PROBLEM
          self.l1 = Conv2d(n_tile_types, 16, 3, 1, 0, bias=True)
          self.l2 = Conv2d(16, 16, 3, 2, 1, bias=True)
          self.l3 = Conv2d(16, n_actions, 3, 1, 1, bias=True)
          self.layers = [self.l1, self.l2, self.l3]
          self.apply(init_play_weights)
          self.flatten = th.nn.Flatten()
          self.net_reward = 0
          self.n_episodes = 0
  
      def forward(self, x):
          x = th.Tensor(get_one_hot_map(x, self.n_tile_types))
          x = x.unsqueeze(0)
          with th.no_grad():
              x = th.relu(self.l1(x))
  
              for i in range(int(np.log2(x.shape[2])) + 1):
                  #           for i in range(1):
                  x = th.relu(self.l2(x))
              x = th.relu(self.l3(x))
  
              #           x = x.argmax(1)
              #           x = x[0]
              x = x.flatten()
              x = th.softmax(x, axis=0)
              # x = [x.argmax().item()]
              act_ids = np.arange(x.shape[0])
              probs = x.detach().numpy()
              x = np.random.choice(act_ids, 1, p=probs)
  
          return x
  
      def assign_reward(self, rew):
          self.net_reward += rew
          self.n_episodes += 1
  
      def reset(self):
          self.net_reward = 0
          self.n_episodes = 0
  
      def get_reward(self):
          mean_rew = self.net_reward / self.n_episodes
  
          return mean_rew
  
  
  def init_weights(m):
      if type(m) == th.nn.Linear:
          th.nn.init.xavier_uniform_(m.weight)
          m.bias.data.fill_(0.01)
  
      if type(m) == th.nn.Conv2d:
          th.nn.init.orthogonal_(m.weight)
  
  
  def init_play_weights(m):
      if type(m) == th.nn.Linear:
          th.nn.init.xavier_uniform(m.weight, gain=0)
          m.bias.data.fill_(0.00)
  
      if type(m) == th.nn.Conv2d:
          #       th.nn.init.orthogonal_(m.weight, gain=np.sqrt(2))
          th.nn.init.constant_(m.weight, 0)
  
  
  def set_nograd(nn):
      for param in nn.parameters():
          param.requires_grad = False
  
  
  def get_init_weights(nn):
      """
      Use to get flat vector of weights from PyTorch model
      """
      init_params = []
  
      if isinstance(nn, CPPN):
          for node in nn.cppn:
              if isinstance(node, Leaf):
                  continue
              init_params.append(node.weights)
              init_params.append(node.bias)
      else:
          for lyr in nn.layers:
              init_params.append(lyr.weight.view(-1).numpy())
              init_params.append(lyr.bias.view(-1).numpy())
      init_params = np.hstack(init_params)
      print("number of initial NN parameters: {}".format(init_params.shape))
  
      return init_params
  
  
  def set_weights(nn, weights):
      if ALGO == "ME":
          # then out nn is contained in the individual
          individual = weights  # I'm sorry mama
          return individual.model
      with th.no_grad():
          n_el = 0
  
          if isinstance(nn, CPPN):
              for node in nn.cppn:
                  l_weights = weights[n_el : n_el + len(node.weights)]
                  n_el += len(node.weights)
                  node.weights = l_weights
                  b_weight = weights[n_el: n_el + 1]
                  n_el += 1
                  node.bias = b_weight
          else:
              for layer in nn.layers:
                  l_weights = weights[n_el : n_el + layer.weight.numel()]
                  n_el += layer.weight.numel()
                  l_weights = l_weights.reshape(layer.weight.shape)
                  layer.weight = th.nn.Parameter(th.Tensor(l_weights))
                  layer.weight.requires_grad = False
                  b_weights = weights[n_el : n_el + layer.bias.numel()]
                  n_el += layer.bias.numel()
                  b_weights = b_weights.reshape(layer.bias.shape)
                  layer.bias = th.nn.Parameter(th.Tensor(b_weights))
                  layer.bias.requires_grad = False
  
      return nn
  
  
  def get_one_hot_map(int_map, n_tile_types):
      obs = (np.arange(n_tile_types) == int_map[..., None]).astype(int)
      obs = obs.transpose(2, 0, 1)
  
      return obs
  
  
  """
  Behavior Characteristics Functions
  """
  
  
  def get_entropy(int_map, env):
      """
      Function to calculate entropy of levels represented by integers
      int_map (numpy array of ints): representation of level
      env (gym-pcgrl environment instance): used to get the action space dims
      returns the entropy of the level normalized roughly to a range of 0.0 to 1.0
      """
      # FIXME: make this robust to different action spaces
      n_classes = len(env._prob._prob)
      max_val = -(1 / n_classes) * np.log(1 / n_classes) * n_classes
      total = len(int_map.flatten())
      entropy = 0.0
  
      for tile in range(n_classes):
          p = (tile == int_map.flatten()).astype(int).sum() / total
  
          if p != 0:
              entropy -= p * np.log(p)
  
      return entropy / max_val
  
  
  def get_counts(int_map, env):
      """
      Function to calculate the tile counts for all possible tiles
      int_map (numpy array of ints): representation of level
      env (gym-pcgrl environment instance): used to get the action space dims
      returns a python list with tile counts for each tile normalized to a range of 0.0 to 1.0
      """
      max_val = env._prob._width * env._prob._height  # for example 14*14=196
  
      return [
          np.sum(int_map.flatten() == tile) / max_val
          for tile in range(len(env._prob._prob))
      ]
  
  
  def get_emptiness(int_map, env):
      """
      Function to calculate how empty the level is
      int_map (numpy array of ints): representation of level
      env (gym-pcgrl environment instance): used to get the action space dims
      returns an emptiness value normalized to a range of 0.0 to 1.0
      """
      max_val = env._prob._width * env._prob._height  # for example 14*14=196
  
      return np.sum(int_map.flatten() == 0) / max_val
  
  
  def get_hor_sym(int_map, env):
      """
      Function to get the horizontal symmetry of a level
      int_map (numpy array of ints): representation of level
      env (gym-pcgrl environment instance): used to get the action space dims
      returns a symmetry float value normalized to a range of 0.0 to 1.0
      """
      max_val = env._prob._width * env._prob._height / 2  # for example 14*14/2=98
      m = 0
  
      if int(int_map.shape[0]) % 2 == 0:
          m = np.sum(
              (
                  int_map[: int(int_map.shape[0] / 2)]
                  == np.flip(int_map[int(int_map.shape[0] / 2) :], 0)
              ).astype(int)
          )
          m = m / max_val
      else:
          m = np.sum(
              (
                  int_map[: int(int_map.shape[0] / 2)]
                  == np.flip(int_map[int(int_map.shape[0] / 2) + 1 :], 0)
              ).astype(int)
          )
          m = m / max_val
  
      return m
  
  
  def get_ver_sym(int_map, env):
      """
      Function to get the vertical symmetry of a level
      int_map (numpy array of ints): representation of level
      env (gym-pcgrl environment instance): used to get the action space dims
      returns a symmetry float value normalized to a range of 0.0 to 1.0
      """
      max_val = env._prob._width * env._prob._height / 2  # for example 14*14/2=98
      m = 0
  
      if int(int_map.shape[1]) % 2 == 0:
          m = np.sum(
              (
                  int_map[:, : int(int_map.shape[1] / 2)]
                  == np.flip(int_map[:, int(int_map.shape[1] / 2) :], 1)
              ).astype(int)
          )
          m = m / max_val
      else:
          m = np.sum(
              (
                  int_map[:, : int(int_map.shape[1] / 2)]
                  == np.flip(int_map[:, int(int_map.shape[1] / 2) + 1 :], 1)
              ).astype(int)
          )
          m = m / max_val
  
      return m
  
  
  # SYMMETRY
  
  
  def get_sym(int_map, env):
      """
      Function to get the vertical symmetry of a level
      int_map (numpy array of ints): representation of level
      env (gym-pcgrl environment instance): used to get the action space dims
      returns a symmetry float value normalized to a range of 0.0 to 1.0
      """
      result = (get_ver_sym(int_map, env) + get_hor_sym(int_map, env)) / 2.0
  
      return result
  
  
  # CO-OCCURRANCE
  
  
  def get_co(int_map, env):
      max_val = env._prob._width * env._prob._height * 4
      result = (
          np.sum((np.roll(int_map, 1, axis=0) == int_map).astype(int))
          + np.sum((np.roll(int_map, -1, axis=0) == int_map).astype(int))
          + np.sum((np.roll(int_map, 1, axis=1) == int_map).astype(int))
          + np.sum((np.roll(int_map, -1, axis=1) == int_map).astype(int))
      )
  
      return result / max_val
  
  
  def get_regions(stats):
      return stats["regions"]
  
  
  def get_path_length(stats):
      return stats["path-length"]
  
  
  # TODO: call this once to return the releveant get_bc function, then call this after each eval, so that we don't have to repeatedly compare strings
  
  
  def get_bc(bc_name, int_map, stats, env):
      if bc_name in stats.keys():
          return stats[bc_name]
      elif bc_name == "co-occurance":
          return get_co(int_map, env)
      elif bc_name == "symmetry":
          return get_sym(int_map, env)
      elif bc_name == "symmetry-vertical":
          return get_ver_sym(int_map, env)
      elif bc_name == "symmetry-horizontal":
          return get_hor_sym(int_map, env)
      elif bc_name == "emptiness":
          return get_emptiness(int_map, env)
      elif bc_name == "entropy":
          return get_entropy(int_map, env)
      elif bc_name == "NONE":
          return 0
      else:
          print("The BC {} is not recognized.".format(bc_name))
          raise Exception
  
          return 0.0
  
  
  class PlayerLeft(nn.Module):
      def __init__(self):
          super().__init__()
          self.act_i = 0
  
      def forward(self, obs):
          return [0]
  
  
  class RandomPlayer(nn.Module):
      def __init__(self, action_space):
          super().__init__()
          self.action_space = action_space
          self.act_i = 0
  
      def forward(self, obs):
          return [self.action_space.sample()]
  
  
  class PlayerRight(nn.Module):
      def __init__(self):
          super().__init__()
          self.act_i = 0
  
      def forward(self, obs):
          return [1]
  
  
  def log_archive(archive, name, itr, start_time, level_json=None):
      if ALGO == "ME":
          # Do this inside optimizer ..?
          return
      # TensorBoard Logging.
      df = archive.as_pandas(include_solutions=False)
      elapsed_time = time.time() - start_time
      writer.add_scalar("{} ArchiveSize".format(name), len(df), itr)
      writer.add_scalar("{} score/mean".format(name), df["objective"].mean(), itr)
      writer.add_scalar("{} score/max".format(name), df["objective"].max(), itr)
      writer.add_scalar("{} score/min".format(name), df["objective"].min(), itr)
  
      # Change: log mean, max, and min for all stats
  
      if level_json:
          stats = ["batch_reward", "targets"]
  
          if N_INIT_STATES > 1:
              stats += ["variance", "diversity"]
          # level_json = {'level': final_levels.tolist(),'batch_reward':[batch_reward] * len(final_levels.tolist()), 'variance': [variance_penalty] * len(final_levels.tolist()), 'diversity':[diversity_bonus] * len(final_levels.tolist()),'targets':trg.tolist(), **bc_dict}
  
          for stat in stats:
              writer.add_scalar(
                  "Training {}/min".format(stat), np.min(level_json[stat]), itr
              )
              writer.add_scalar(
                  "Training {}/mean".format(stat), np.mean(level_json[stat]), itr
              )
              writer.add_scalar(
                  "Training {}/max".format(stat), np.max(level_json[stat]), itr
              )
  
      # Logging.
  
      if itr % 1 == 0:
          print(f"> {itr} itrs completed after {elapsed_time:.2f} s")
          print(f"  - {name} Archive Size: {len(df)}")
          print(f"  - {name} Max Score: {df['objective'].max()}")
          print(f"  - {name} Mean Score: {df['objective'].mean()}")
          print(f"  - {name} Min Score: {df['objective'].min()}")
  
  
  N_PLAYER_STEPS = 100
  
  
  def play_level(env, level, player):
      env._rep._old_map = level
      env._rep._random_start = False
      p_obs = env.reset()
  
      if not env.is_playable():
          return 0, None
      # TODO: check if env is playable!
      env.set_active_agent(1)
  
      if RENDER:
          env.render()
      net_p_rew = 0
      action_hist = []
  
      for p_i in range(N_PLAYER_STEPS):
          action = player(p_obs["map"])
  
          if isinstance(action, th.Tensor):
              # TODO: this logic belongs with the model
              player_coords = env._prob.player.coords
              action = np.array(action)[player_coords[0], player_coords[1]]
          elif isinstance(action, list) or isinstance(action, np.ndarray):
              assert len(action) == 1
              action = action[-1]
          else:
              raise Exception
          action_hist.append(action)
          p_obs, p_rew, p_done, p_info = env.step(action)
  
          if RENDER:
              env.render()
          # net_p_rew += p_rew
          net_p_rew = p_rew
  
          if p_done:
              break
  
      #   player.assign_reward(net_p_rew)
      action_freqs = np.bincount(action_hist, minlength=len(env.player_actions))
      action_entropy = scipy.stats.entropy(action_freqs)
      local_action_entropy = np.mean(
          [
              scipy.stats.entropy(
                  np.bincount(action_hist[i : i + 10], minlength=len(env.player_actions))
              )
              for i in np.arange(0, len(action_hist) - 10, 6)
          ]
      )
      local_action_entropy = np.nan_to_num(local_action_entropy)
  
      return net_p_rew, [action_entropy, local_action_entropy]
  
  
  @ray.remote
  def multi_evo(
      env,
      model,
      model_w,
      n_tile_types,
      init_states,
      bc_names,
      static_targets,
      seed,
      player_1,
      player_2,
      proc_id=None,
      init_states_archive=None,
      index=None,
  ):
      if init_states is None:
          init_states = get_init_states(init_states_archive, tuple(index))
  
      if proc_id is not None:
          print("simulating id: {}".format(proc_id))
      model = set_weights(model, model_w)
      result = simulate(
          env=env,
          model=model,
          n_tile_types=n_tile_types,
          init_states=init_states,
          bc_names=bc_names,
          static_targets=static_targets,
          seed=seed,
          player_1=player_1,
          player_2=player_2,
      )
  
      return result
  
  
  @ray.remote
  def multi_play_evo(
      env,
      gen_model,
      player_1_w,
      n_tile_types,
      init_states,
      play_bc_names,
      static_targets,
      seed,
      player_1,
      player_2,
      playable_levels,
      proc_id=None,
  ):
  
      if proc_id is not None:
          print("simulating id: {}".format(proc_id))
      player_1 = set_weights(player_1, player_1_w)
      obj, bcs = player_simulate(
          env=env,
          n_tile_types=n_tile_types,
          play_bc_names=play_bc_names,
          seed=seed,
          player_1=player_1,
          playable_levels=playable_levels,
      )
  
      return obj, bcs
  
  
  def gen_playable_levels(env, gen_model, init_states, n_tile_types):
      """ To get only the playable levels of a given generator, so that we can run player evaluations on them more quickly."""
      final_levels = []
  
      for int_map in init_states:
          obs = get_one_hot_map(int_map, n_tile_types)
  
          if RENDER:
              env.render()
          done = False
          n_step = 0
          last_int_map = None
  
          while not done:
              int_tensor = th.unsqueeze(th.Tensor(obs), 0)
              action, done = gen_model(int_tensor)[0].numpy()
              obs = action
              int_map = done or action.argmax(axis=0)
              env._rep._map = int_map
              done = done or (int_map == last_int_map).all() or n_step >= N_STEPS
  
              #           if INFER and not EVALUATE:
              #               time.sleep(1 / 30)
  
              if done:
                  gen_model.reset()
                  env._rep._old_map = int_map
                  env._rep._random_start = False
                  _ = env.reset()
  
                  if env.is_playable():
                      final_levels.append(int_map)
              n_step += 1
  
      return final_levels
  
  
  def player_simulate(
      env, n_tile_types, play_bc_names, player_1, playable_levels, seed=None
  ):
      n_evals = 10
      net_reward = 0
      bcs = []
  
      for int_map in playable_levels * n_evals:
          if INFER:
              #           env.render()
              input("ready player 1")
          p_1_rew, p_bcs = play_level(env, int_map, player_1)
          bcs.append(p_bcs)
  
          if INFER:
              print("p_1 reward: ", p_1_rew)
          net_reward += p_1_rew
  
      reward = net_reward / len(playable_levels * n_evals)
      bcs = [np.mean([bcs[j][i] for j in range(len(bcs))]) for i in range(len(bcs[0]))]
  
      return reward, bcs
  
  def plot_score_heatmap(scores, score_name, bc_names, cmap_str="magma", bcs_in_filename=True,
                         lower_bounds=None, upper_bounds=None,
                         x_bounds=None, y_bounds=None):
      scores = scores.T
      ax = plt.gca()
      ax.set_xlim(lower_bounds[0], upper_bounds[0])
      ax.set_ylim(lower_bounds[1], upper_bounds[1])
      label_fontdict = {
          'fontsize': 16,
      }
      ax.set_xlabel(bc_names[0], fontdict=label_fontdict)
      ax.set_ylabel(bc_names[1], fontdict=label_fontdict)
      vmin = np.nanmin(scores)
      vmax = np.nanmax(scores)
      t = ax.pcolormesh(
          x_bounds,
          y_bounds,
          scores,
          cmap=matplotlib.cm.get_cmap(cmap_str),
          vmin=vmin,
          vmax=vmax,
      )
      ax.figure.colorbar(t, ax=ax, pad=0.1)
  
      if SHOW_VIS:
          plt.show()
      if bcs_in_filename:
          f_name = score_name + "_" + "-".join(bc_names)
      else:
          f_name = score_name
  
      if not RANDOM_INIT_LEVELS:
          f_name = f_name + "_fixLvls"
      f_name += ".png"
      plt.title(score_name, fontdict={'fontsize': 24})
      plt.tight_layout()
      plt.savefig(os.path.join(SAVE_PATH, f_name))
      plt.close()
  
  def simulate(
          env,
          model,
          n_tile_types,
          init_states,
          bc_names,
          static_targets,
          seed=None,
          player_1=None,
          player_2=None,
          render_levels=False
  ):
      """
      Function to run a single trajectory and return results.
  
      Args:
          env (gym.Env): A copy of the binary-wide-v0 environment.
          model (np.ndarray): The array of weights for the policy.
          seed (int): The seed for the environment.
          player_sim (bool): Are we collecting obj and bcs for the player, rather than the generator?
      Returns:
          total_reward (float): The reward accrued by the lander throughout its
              trajectory.
          path_length (float): The path length of the final solution.
          regions (float): The number of distinct regions of the final solution.
      """
      global N_INIT_STATES
  
      if seed is not None:
          env.seed(seed)
  
      if PLAY_LEVEL:
          assert player_1 is not None
          assert player_2 is not None
  
      if CMAES:
          bc_names = ["NONE", "NONE"]
      # Allow us to manually set the level-map on reset (using the "_old_map" attribute)
      # Actually we have found a more efficient workaround for now.
  
      #   env._rep._random_start = False
      #   if n_episode == 0 and False:
      #       env._rep._old_map = init_state
      #       obs = env.reset()
      #       int_map = obs['map']
      n_init_states = init_states.shape[0]
      width = init_states.shape[1]
      height = init_states.shape[2]
      bcs = np.empty(shape=(len(bc_names), n_init_states))
      # if SAVE_LEVELS:
      trg = np.empty(shape=(n_init_states))
      final_levels = np.empty(shape=init_states.shape, dtype=np.uint8)
      batch_reward = 0
      batch_time_penalty = 0
      batch_targets_penalty = 0
      batch_play_bonus = 0
  
      if render_levels:
          level_frames = []
  
      for (n_episode, init_state) in enumerate(init_states):
          # NOTE: Sneaky hack. We don't need initial stats. Never even reset. Heh. Be careful!!
          # Set the representation to begin in the upper left corner
          env._rep._map = init_state.copy()
          env._prob.path_coords = []
          env._prob.path_length = None
          # Only applies to narrow and turtle. Better than using reset, but ugly, and not optimal
          # TODO: wrap the env instead
          env._rep._x = env._rep._y = 0
          # env._rep._x = np.random.randint(env._prob._width)
          # env._rep._y = np.random.randint(env._prob._height)
          int_map = init_state
          obs = get_one_hot_map(int_map, n_tile_types)
  
          if RENDER:
              env.render()
  
              if INFER:
                  #               time.sleep(10/30)
                  #               input()
                  pass
          done = False
  
          n_step = 0
  
          while not done:
              if render_levels:
                  level_frames.append(env.render(mode="rgb_array"))
              #           in_tensor = th.unsqueeze(
              #               th.unsqueeze(th.tensor(np.float32(obs['map'])), 0), 0)
              in_tensor = th.unsqueeze(th.Tensor(obs), 0)
              action, done = model(in_tensor)
              action = action[0].numpy()
              # There is probably a better way to do this, so we are not passing unnecessary kwargs, depending on representation
              action, skip = preprocess_action(
                  action,
                  int_map=env._rep._map,
                  x=env._rep._x,
                  y=env._rep._y,
                  n_dirs=N_DIRS,
                  n_tiles=n_tile_types,
              )
              change, x, y = env._rep.update(action)
              int_map = env._rep._map
              obs = get_one_hot_map(env._rep.get_observation()["map"], n_tile_types)
              preprocess_observation(obs, x=env._rep._x, y=env._rep._y)
              #           int_map = action.argmax(axis=0)
              #           obs = get_one_hot_map(int_map, n_tile_types)
              #           env._rep._map = int_map
              done = done or not (change or skip) or n_step >= N_STEPS
              # done = n_step >= N_STEPS
  
              #           if INFER and not EVALUATE:
              #               time.sleep(1 / 30)
  
              if done:
                  model.reset()
                  if render_levels:
                      # get final level state
                      level_frames.append(env.render(mode="rgb_array"))
                  # we'll need this to compute Hamming diversity
                  final_levels[n_episode] = int_map
                  stats = env._prob.get_stats(
                      get_string_map(int_map, env._prob.get_tile_types()),
                      # lenient_paths = True,
                  )
  
                  # get BCs
                  # Resume here. Use new BC function.
  
                  for i in range(len(bc_names)):
                      bc_name = bc_names[i]
                      bcs[i, n_episode] = get_bc(bc_name, int_map, stats, env)
  
                  # TODO: reward calculation should depend on self.reward_names
                  # ad hoc reward: shorter episodes are better?
                  time_penalty = n_step
                  batch_time_penalty -= time_penalty
  
                  # we want to hit each of our static targets exactly, penalize for anything else.
                  # for ranges, we take our least distance to any element in the range
                  targets_penalty = 0
  
                  for k in static_targets:
                      if k in bc_names:
                          continue
  
                      if isinstance(static_targets[k], tuple):
                          # take the smallest distance from current value to any point in range
                          # NOTE: we're assuming this metric is integer-valued
                          targets_penalty += abs(
                              np.arange(static_targets[k][0], static_targets[k][1]) - stats[k]
                          ).min()
                      else:
                          targets_penalty += abs(static_targets[k] - stats[k])
                  #                   targets_penalty = np.sum([abs(static_targets[k] - stats[k]) if not isinstance(static_targets[k], tuple) else abs(np.arange(*static_targets[k]) - stats[k]).min() for k in static_targets])
                  batch_targets_penalty -= targets_penalty
                  # if SAVE_LEVELS:
                  trg[n_episode] = -targets_penalty
  
                  if PLAY_LEVEL:
                      if INFER:
                          env.render()
                          input("ready player 1")
                      p_1_rew, p_bcs = play_level(env, int_map, player_1)
  
                      if INFER:
                          print("p_1 reward: ", p_1_rew)
                          input("ready player 2")
                      p_2_rew, p_bcs = play_level(env, int_map, player_2)
  
                      if INFER:
                          print("p_2 reward: ", p_2_rew)
  
                      max_regret = env._prob.max_reward - env._prob.min_reward
                      # add this in case we get worst possible regret (don't want to punish a playable map)
                      batch_play_bonus += max_regret + p_1_rew - p_2_rew
  
              if RENDER:
                  if INFER:
                      stats = env._prob.get_stats(
                          get_string_map(int_map, env._prob.get_tile_types()),
                          # lenient_paths=True,
                      )
                  env.render()
  
  
              if done and INFER:  # and not (EVALUATE and THREADS):
                  if not EVALUATE:
                      #                   time.sleep(5 / 30)
                      print(
                          "stats: {}\n\ntime_penalty: {}\n targets_penalty: {}".format(
                              stats, time_penalty, targets_penalty
                          )
                      )
              last_int_map = int_map
              n_step += 1
      final_bcs = [bcs[i].mean() for i in range(bcs.shape[0])]
      batch_targets_penalty = TARGETS_PENALTY_WEIGHT * batch_targets_penalty / max(N_INIT_STATES, 1)
      # batch_targets_penalty = batch_targets_penalty / N_INIT_STATES
      batch_reward += batch_targets_penalty
  
      if PLAY_LEVEL:
          batch_reward += batch_play_bonus / max(N_INIT_STATES, 1)
          time_penalty, targets_penalty, variance_penalty, diversity_bonus = (
              None,
              None,
              None,
              None,
          )
      else:
          #       batch_time_penalty = batch_time_penalty / N_INIT_STATES
  
          N_INIT_STATES = n_init_states
  
          if N_INIT_STATES > 1 and (batch_targets_penalty == 0 or not CASCADE_REWARD):
              # Calculate stats that depend on having generated multiple levels. If using gated reward, only calculate these additional components of reward if level is
              # perfectly valid.
              # Variance penalty is the negative average (per-BC) standard deviation from the mean BC vector.
              variance_penalty = (
                  -np.sum([bcs[i].std() for i in range(bcs.shape[0])]) / bcs.shape[0]
              )
              # Diversity bonus. We want minimal variance along BCS *and* diversity in terms of the map.
              # Sum pairwise hamming distances between all generated maps.
              diversity_bonus = np.sum(
                  [
                      np.sum(final_levels[j] != final_levels[k]) if j != k else 0
                      for k in range(N_INIT_STATES)
                      for j in range(N_INIT_STATES)
                  ]
              ) / (N_INIT_STATES * N_INIT_STATES - 1)
              # ad hoc scaling :/
              diversity_bonus = 10 * diversity_bonus / (width * height)
              batch_reward = batch_reward + max(0, variance_penalty + diversity_bonus)
          else:
              variance_penalty = None
              diversity_bonus = None
  
      if SAVE_LEVELS:
          bc_dict = {}
  
          for i in range(len(bc_names)):
              bc_name = bc_names[i]
              bc_dict[bc_name] = bcs[i, :].tolist()
          level_json = {
              "level": final_levels.tolist(),
              "batch_reward": [batch_reward] * len(final_levels.tolist()),
              "variance": [variance_penalty] * len(final_levels.tolist()),
              "diversity": [diversity_bonus] * len(final_levels.tolist()),
              "targets": trg.tolist(),
              **bc_dict,
          }
      else:
          level_json = {
              "level": final_levels.tolist(),
              "batch_reward": [batch_reward] * len(final_levels.tolist()),
              "variance": [variance_penalty] * len(final_levels.tolist()),
              "diversity": [diversity_bonus] * len(final_levels.tolist()),
              "targets": trg.tolist(),
          }
  
      if render_levels:
          return level_frames
  
      if not INFER:
          return level_json, batch_reward, final_bcs
      else:
          return (
              level_json,
              batch_reward,
              final_bcs,
              (
                  batch_time_penalty,
                  batch_targets_penalty,
                  variance_penalty,
                  diversity_bonus,
              ),
          )
  
  
  class EvoPCGRL:
      def __init__(self):
          self.init_env()
          assert self.env.observation_space["map"].low[0, 0] == 0
          # get number of tile types from environment's observation space
          # here we assume that all (x, y) locations in the observation space have the same upper/lower bound
          self.n_tile_types = self.env.observation_space["map"].high[0, 0] + 1
          self.width = self.env._prob._width
          self.height = self.env._prob._height
  
          # FIXME why not?
          # self.width = self.env._prob._width
  
          # TODO: make reward a command line argument?
          # TODO: multi-objective compatibility?
          self.bc_names = BCS
  
          # calculate the bounds of our behavioral characteristics
          # NOTE: We assume a square map for some of these (not ideal).
          # regions and path-length are applicable to all PCGRL problems
          self.bc_bounds = self.env._prob.cond_bounds
          self.bc_bounds.update(
              {
                  "co-occurance": (0.0, 1.0),
                  "symmetry": (0.0, 1.0),
                  "symmetry-vertical": (0.0, 1.0),
                  "symmetry-horizontal": (0.0, 1.0),
                  "emptiness": (0.0, 1.0),
                  "entropy": (0.0, 1.0),
              }
          )
  
          self.static_targets = self.env._prob.static_trgs
  
          if REEVALUATE_ELITES or (RANDOM_INIT_LEVELS and args.n_init_states != 0):
              init_level_archive_args = (N_INIT_STATES, self.height, self.width)
          else:
              init_level_archive_args = ()
  
          if ALGO == "ME":
              if RANDOM_INIT_LEVELS and args.n_init_states != 0:
                  gen_archive_cls = MEInitStatesArchive
              else:
                  gen_archive_cls = MEGrid
  
          elif REEVALUATE_ELITES:
              # If we are constantly providing new random seeds to generators, we may want to regularly re-evaluate
              # elites
              gen_archive_cls = FlexArchive
          elif RANDOM_INIT_LEVELS and not args.n_init_states == 0:
              # If we have random seeds each generation but are not re-evaluating elites, then we want to hang onto these
              # random seeds.
              gen_archive_cls = InitStatesArchive
          #           gen_archive_cls = GridArchive
          else:
              gen_archive_cls = GridArchive
              init_level_archive_args = ()
  
          if PLAY_LEVEL:
              self.play_bc_names = ["action_entropy", "local_action_entropy"]
              self.play_bc_bounds = {
                  "action_entropy": (0, 4),
                  "local_action_entropy": (0, 4),
              }
              self.gen_archive = gen_archive_cls(
                  [100 for _ in self.bc_names],
                  # [1],
                  # [(-1, 1)],
                  [self.bc_bounds[bc_name] for bc_name in self.bc_names],
              )
              self.play_archive = FlexArchive(
                  # minimum of: 100 for each behavioral characteristic, or as many different values as the BC can take on, if it is less
                  # [min(100, int(np.ceil(self.bc_bounds[bc_name][1] - self.bc_bounds[bc_name][0]))) for bc_name in self.bc_names],
                  [100 for _ in self.play_bc_names],
                  # min/max for each BC
                  [self.play_bc_bounds[bc_name] for bc_name in self.play_bc_names],
              )
          else:
              if CMAES:
                  # Restrict the archive to 1 cell so that we are effectively doing CMAES. BCs should be ignored.
                  self.gen_archive = gen_archive_cls(
                      [1, 1], [(0, 1), (0, 1)], *init_level_archive_args
                  )
              else:
                  self.gen_archive = gen_archive_cls(
                      # minimum of 100 for each behavioral characteristic, or as many different values as the BC can take on, if it is less
                      # [min(100, int(np.ceil(self.bc_bounds[bc_name][1] - self.bc_bounds[bc_name][0]))) for bc_name in self.bc_names],
                      [100 for _ in self.bc_names],
                      # min/max for each BC
                      [self.bc_bounds[bc_name] for bc_name in self.bc_names],
                      *init_level_archive_args,
                  )
  
          reps_to_out_chans = {
              "cellular": self.n_tile_types,
              "wide": self.n_tile_types,
              "narrow": self.n_tile_types + 1,
              "turtle": self.n_tile_types + N_DIRS,
          }
  
          reps_to_in_chans = {
              "cellular": self.n_tile_types,
              "wide": self.n_tile_types,
              "narrow": self.n_tile_types + 1,
              "turtle": self.n_tile_types + 1,
          }
  
          n_out_chans = reps_to_out_chans[REPRESENTATION]
          n_in_chans = reps_to_in_chans[REPRESENTATION]
  
          if MODEL == "CNN":
              # Adding n_tile_types as a dimension here. Why would this not be in the env's observation space though? Should be one-hot by default?
              observation_shape = (
                  1,
                  self.n_tile_types,
                  *self.env.observation_space["map"].shape,
              )
  
              if isinstance(self.env.action_space, gym.spaces.Box):
                  action_shape = self.env.action_space.shape
                  assert len(action_shape) == 3
                  n_flat_actions = action_shape[0] * action_shape[1] * action_shape[2]
              elif isinstance(self.env.action_space, gym.spaces.MultiDiscrete):
                  nvec = self.env.action_space.nvec
                  assert len(nvec) == 3
                  n_flat_actions = nvec[0] + nvec[1] + nvec[2]
              elif isinstance(self.env.action_space, gym.spaces.Discrete):
                  n_flat_actions = self.env.action_space.n
              else:
                  raise NotImplementedError(
                      "I don't know how to handle this action space: {}".format(
                          type(self.env.action_space)
                      )
                  )
              self.gen_model = GeneratorNNDense(
                  n_in_chans=self.n_tile_types,
                  n_actions=n_out_chans,
                  observation_shape=observation_shape,
                  n_flat_actions=n_flat_actions,
              )
          # TODO: remove this, just call model "NCA"
          elif MODEL == "NCA":
              self.gen_model = globals()["GeneratorNN"](
                  n_in_chans=self.n_tile_types, n_actions=n_out_chans
              )
          else:
              self.gen_model = globals()[MODEL](
                  n_in_chans=self.n_tile_types, n_actions=n_out_chans
              )
          set_nograd(self.gen_model)
          initial_w = get_init_weights(self.gen_model)
          assert len(initial_w.shape) == 1
          self.n_generator_weights = initial_w.shape[0]
          self.n_player_weights = 0
          # TODO: different initial weights per emitter as in pyribs lunar lander relanded example?
  
          if MODEL == "NCA":
              init_step_size = 1
          elif MODEL == "CNN":
              init_step_size = 1
          else:
              init_step_size = 1
  
          if CMAES:
              # The optimizing emitter will prioritize fitness over exploration of behavior space
              emitter_type = OptimizingEmitter
          else:
              emitter_type = ImprovementEmitter
  
          batch_size = 30
          n_emitters = 5
          if ALGO == "ME":
              pass
  
          elif args.mega:
              gen_emitters = [
                  GradientImprovementEmitter(
                      self.gen_archive,
                      initial_w.flatten(),
                      # TODO: play with initial step size?
                      sigma_g=10.0,
                      stepsize=0.002,  # Initial step size.
                      gradient_optimizer="adam",
                      selection_rule="mu",
                      batch_size=batch_size,
                  )
                  for _ in range(n_emitters)  # Create 5 separate emitters.
              ]
          else:
              gen_emitters = [
                  #           ImprovementEmitter(
                  emitter_type(
                      self.gen_archive,
                      initial_w.flatten(),
                      # TODO: play with initial step size?
                      init_step_size,  # Initial step size.
                      batch_size=batch_size,
                  )
                  for _ in range(n_emitters)  # Create 5 separate emitters.
              ]
  
          if PLAY_LEVEL:
              # Concatenate designer and player weights
              self.play_model = PlayerNN(
                  self.n_tile_types, n_actions=len(self.env.player_actions)
              )
              set_nograd(self.play_model)
              initial_play_w = get_init_weights(self.play_model)
              assert len(initial_play_w.shape) == 1
              self.n_player_weights = initial_play_w.shape[0]
              play_emitters = [
                  OptimizingEmitter(
                      self.play_archive,
                      initial_play_w.flatten(),
                      # NOTE: Big step size, no good otherwise
                      1,  # Initial step size.
                      batch_size=batch_size,
                  )
                  for _ in range(n_emitters)  # Create 5 separate emitters.
              ]
              self.play_optimizer = Optimizer(self.play_archive, play_emitters)
          if ALGO == "ME":
              ind_cls_args = {
                      'model_cls': globals()[MODEL],
                     'n_in_chans': self.n_tile_types,
                     'n_actions': self.n_tile_types,
              }
  
              self.gen_optimizer = MEOptimizer(self.gen_archive,
                                               ind_cls=Individual,
                                               batch_size=n_emitters*batch_size,
                                               ind_cls_args=ind_cls_args,
                                               )
          else:
              self.gen_optimizer = Optimizer(self.gen_archive, gen_emitters)
  
          # These are the initial maps which will act as seeds to our NCA models
  
          if args.n_init_states == 0:
              # special square patch
              self.init_states = np.zeros(shape=(1, self.height, self.width))
              self.init_states[0, 5:-5, 5:-5] = 1
          else:
              #           self.init_states = np.random.randint(
              #               0, self.n_tile_types, (N_INIT_STATES, self.width, self.height)
              #           )
              self.init_states = gen_random_levels(N_INIT_STATES, self.env)
  
          self.start_time = time.time()
          self.total_itrs = N_GENERATIONS
          self.n_itr = 1
  
          if PLAY_LEVEL:
              self.player_1 = PlayerNN(self.n_tile_types)
              self.player_2 = RandomPlayer(self.env.player_action_space)
          else:
              self.player_1 = None
              self.player_2 = None
          # This directory might already exist if a previous experiment failed before the first proper checkpoint/save
  
          if not os.path.isdir(SAVE_PATH):
              os.mkdir(SAVE_PATH)
          # Save the command line arguments with which we launched
          with open(os.path.join(SAVE_PATH, "settings.json"), "w", encoding="utf-8") as f:
              json.dump(arg_dict, f, ensure_ascii=False, indent=4)
  
      def evolve(self):
  
          net_p_itr = 0
  
          for itr in tqdm(range(self.n_itr, self.total_itrs + 1)):
              # Request models from the optimizer.
  
              if args.mega:
                  gen_sols = self.gen_optimizer.ask(grad_estimate=True)
              else:
                  # if algo is ME, these are "individual" objects
                  gen_sols = self.gen_optimizer.ask()
  
              # Evaluate the models and record the objectives and BCs.
              objs, bcs = [], []
              stats = ["batch_reward", "variance", "diversity", "targets"]
              stat_json = {
                  "batch_reward": [],
                  "variance": [],
                  "diversity": [],
                  "targets": [],
              }
  
              if RANDOM_INIT_LEVELS and args.n_init_states != 0:
                  init_states = gen_random_levels(N_INIT_STATES, self.env)
              else:
                  init_states = self.init_states
  
              if THREADS:
                  n_sols = len(gen_sols)
                  if N_PROC is not None:
                      n_proc = N_PROC
                  else:
                      n_proc = n_sols
                  n_launches = np.ceil(n_sols / n_proc)
                  results = []
                  for n_launch in range(int(n_launches)):
                      futures = [
                          multi_evo.remote(
                              self.env,
                              self.gen_model,
                              model_w,
                              self.n_tile_types,
                              init_states,
                              self.bc_names,
                              self.static_targets,
                              seed,
                              player_1=self.player_1,
                              player_2=self.player_2,
                          )
                          for model_w in gen_sols
                      ]
                      results += ray.get(futures)
                      del futures
                      auto_garbage_collect()
  
                  for result in results:
                      level_json, m_obj, m_bcs = result
  
                      if SAVE_LEVELS:
                          df = pd.DataFrame(level_json)
                          df = df[df["targets"] == 0]
  
                          if len(df) > 0:
                              df.to_csv(
                                  os.path.join(SAVE_PATH, "levels.csv"),
                                  mode="a",
                                  header=False,
                                  index=False,
                              )
                      objs.append(m_obj)
                      bcs.append([*m_bcs])
                      [stat_json[stat].extend(level_json[stat]) for stat in stats]
                  del results
                  auto_garbage_collect()
              else:
                  for model_w in gen_sols:
                      gen_model = set_weights(self.gen_model, model_w)
                      level_json, m_obj, m_bcs = simulate(
                          env=self.env,
                          model=gen_model,
                          n_tile_types=self.n_tile_types,
                          init_states=init_states,
                          bc_names=self.bc_names,
                          static_targets=self.static_targets,
                          seed=seed,
                          player_1=self.player_1,
                          player_2=self.player_2,
                      )
  
                      if SAVE_LEVELS:
                          # Save levels to disc
                          df = pd.DataFrame(level_json)
                          df = df[df["targets"] == 0]
  
                          if len(df) > 0:
                              df.to_csv(
                                  os.path.join(SAVE_PATH, "levels.csv"),
                                  mode="a",
                                  header=False,
                                  index=False,
                              )
                      objs.append(m_obj)
                      bcs.append(m_bcs)
                      [stat_json[stat].extend(level_json[stat]) for stat in stats]
  
              if RANDOM_INIT_LEVELS:
                  # Tell the archive what the initial states are, so that we can record them in case an individual is
                  # added.
                  self.gen_archive.set_init_states(init_states)
              # Send the results back to the optimizer.
              if args.mega:
                  # TODO: Here we need the jacobian
                  jacobian = None
                  self.gen_optimizer.tell(objs, bcs, jacobian=jacobian)
              else:
                  self.gen_optimizer.tell(objs, bcs)
  #               for emitter in self.gen_optimizer.emitters:
  #
  
              # Re-evaluate elite generators. If doing CMAES, re-evaluate every iteration. Otherwise, try to let the archive grow.
  
              if REEVALUATE_ELITES and (CMAES or self.n_itr % 1 == 0):
                  df = self.gen_archive.as_pandas()
                  #               curr_archive_size = len(df)
                  high_performing = df.sample(frac=1)
                  elite_models = np.array(high_performing.loc[:, "solution_0":])
                  elite_bcs = np.array(high_performing.loc[:, "behavior_0":"behavior_1"])
  
                  if THREADS:
                      futures = [
                          multi_evo.remote(
                              self.env,
                              self.gen_model,
                              elite_models[i],
                              self.n_tile_types,
                              init_states,
                              self.bc_names,
                              self.static_targets,
                              seed,
                              player_1=self.player_1,
                              player_2=self.player_2,
                          )
                          for i in range(min(max(len(elite_models) // 2, 1), 150 // 2))
                      ]
                      results = ray.get(futures)
  
                      for (el_i, result) in enumerate(results):
                          old_el_bcs = elite_bcs[el_i]
                          level_json, el_obj, el_bcs = result
  
                          if SAVE_LEVELS:
                              # Save levels to disk
                              df = pd.DataFrame(level_json)
                              df = df[df["targets"] == 0]
  
                              if len(df) > 0:
                                  df.to_csv(
                                      os.path.join(SAVE_PATH, "levels.csv"),
                                      mode="a",
                                      header=False,
                                      index=False,
                                  )
                          #                       mean_obj, mean_bcs, obj_hist, bc_hist = self.gen_archive.pop_elite(el_obj, el_bcs, old_el_bcs)
                          results[el_i] = self.gen_archive.pop_elite(
                              el_obj, el_bcs, old_el_bcs
                          )
                          [stat_json[stat].extend(level_json[stat]) for stat in stats]
  
                      for (el_i, result) in enumerate(results):
                          self.gen_archive.update_elite(*result)
                      del results
                      auto_garbage_collect()
  
                  else:
                      # 150 to match number of new-model evaluations
  
                      for elite_i in range(min(max(len(elite_models) // 2, 1), 150 // 2)):
                          # print(elite_i)
                          # pprint.pprint(self.gen_archive.obj_hist, width=1)
                          # pprint.pprint(self.gen_archive.bc_hist, width=1)
                          old_el_bcs = elite_bcs[elite_i]
                          gen_model_weights = elite_models[elite_i]
                          gen_model = set_weights(self.gen_model, gen_model_weights)
  
                          level_json, el_obj, el_bcs = simulate(
                              env=self.env,
                              model=gen_model,
                              n_tile_types=self.n_tile_types,
                              init_states=init_states,
                              bc_names=self.bc_names,
                              static_targets=self.static_targets,
                              seed=seed,
                              player_1=self.player_1,
                              player_2=self.player_2,
                          )
                          idx = self.gen_archive.get_index(old_el_bcs)
                          [stat_json[stat].extend(level_json[stat]) for stat in stats]
                          self.gen_archive.update_elite(
                              *self.gen_archive.pop_elite(el_obj, el_bcs, old_el_bcs)
                          )
  
              #               last_archive_size = len(self.gen_archive.as_pandas(include_solutions=False))
  
              log_archive(self.gen_archive, "Generator", itr, self.start_time, stat_json)
  
              # FIXME: implement these
              #           self.play_bc_names = ['action_entropy', 'action_entropy_local']
  
              if PLAY_LEVEL:
                  # elite_model_w = self.gen_archive.get_random_elite()[0]
                  df = self.gen_archive.as_pandas()
                  high_performing = df.sort_values("objective", ascending=False)
                  models = np.array(high_performing.loc[:, "solution_0":])
                  np.random.shuffle(models)
                  playable_levels = []
  
                  for m_i in range(len(models)):
                      elite_model_w = models[m_i]
                      gen_model = set_weights(self.gen_model, elite_model_w)
                      playable_levels += gen_playable_levels(
                          self.env, self.gen_model, self.init_states, self.n_tile_types
                      )
  
                      if len(playable_levels) >= 50:
                          break
  
                  if len(playable_levels) >= 10:
                      play_start_time = time.time()
                      self.playable_levels = playable_levels
  
                      for p_itr in tqdm(range(1, 2)):
                          net_p_itr += 1
                          play_sols = self.play_optimizer.ask()
                          objs, bcs = [], []
  
                          if THREADS:
                              futures = [
                                  multi_play_evo.remote(
                                      self.env,
                                      gen_model,
                                      player_w,
                                      self.n_tile_types,
                                      init_states,
                                      self.play_bc_names,
                                      self.static_targets,
                                      seed,
                                      player_1=self.player_1,
                                      player_2=self.player_2,
                                      playable_levels=playable_levels,
                                  )
                                  for player_w in play_sols
                              ]
                              results = ray.get(futures)
  
                              for result in results:
                                  m_obj, m_bcs = result
                                  objs.append(m_obj)
                                  bcs.append([*m_bcs])
                              del results
                              auto_garbage_collect()
                          else:
                              play_i = 0
  
                              for play_w in play_sols:
                                  play_i += 1
                                  play_model = set_weights(self.play_model, play_w)
                                  m_obj, m_bcs = player_simulate(
                                      env=self.env,
                                      n_tile_types=self.n_tile_types,
                                      play_bc_names=self.play_bc_names,
                                      seed=seed,
                                      player_1=self.player_1,
                                      playable_levels=playable_levels,
                                  )
                                  objs.append(m_obj)
                                  bcs.append(m_bcs)
                          self.play_optimizer.tell(objs, bcs)
  
                          # TODO: parallelize me
                          df = self.play_archive.as_pandas()
                          high_performing = df.sort_values("objective", ascending=False)
                          elite_models = np.array(high_performing.loc[:, "solution_0":])
  
                          for elite_i in range(10):
                              play_model_weights = elite_models[elite_i]
                              init_nn = set_weights(self.play_model, play_model_weights)
  
                              obj, bcs = player_simulate(
                                  self.env,
                                  self.n_tile_types,
                                  self.play_bc_names,
                                  init_nn,
                                  playable_levels=playable_levels,
                              )
  
                              self.play_archive.update_elite(obj, bcs)
  
                          #    m_objs.append(obj)
                          # bc_a = get_bcs(init_nn)
                          # obj = np.mean(m_objs)
                          # objs.append(obj)
                          # bcs.append([bc_a])
                          log_archive(self.play_archive, "Player", p_itr, play_start_time)
  
                          if net_p_itr > 0 and net_p_itr % SAVE_INTERVAL == 0:
                              # Save checkpoint during player evolution loop
                              self.save()
  
                          df = self.play_archive.as_pandas()
                          high_performing = df.sort_values("objective", ascending=False)
                          elite_scores = np.array(high_performing.loc[:, "objective"])
  
                          if np.array(elite_scores).max() >= self.env._prob.max_reward:
                              break
  
                      # TODO: assuming an archive of one here! Make it more general, like above for generators
                      play_model = set_weights(
                          self.play_model, self.play_archive.get_random_elite()[0]
                      )
  
              if itr % SAVE_INTERVAL == 0 or itr == 1:
                  # Save checkpoint during generator evolution loop
                  self.save()
  
  #           if itr % VIS_INTERVAL == 0 or itr == 1:
  #               ckp_dir = os.path.join(SAVE_PATH, "checkpoint_{}".format(itr))
  
  #               if not os.path.isdir(ckp_dir):
  #                   os.mkdir(ckp_dir)
  
  #               if not CMAES:
  #                   # Otherwise the heatmap would just be a single cell
  #                   self.visualize(itr=itr)
  #               archive_objs = np.array(
  #                   self.gen_archive.as_pandas(include_solutions=False).loc[
  #                       :, "objective"
  #                   ]
  #               )
  #               save_train_stats(archive_objs, itr=itr)
  
              self.n_itr += 1
  
      def save(self):
          global ENV
          ENV = self.env
          self.env = None
          evo_path = os.path.join(SAVE_PATH, "evolver.pkl")
  
          os.system(
              'mv "{}" "{}"'.format(evo_path, os.path.join(SAVE_PATH, "last_evolver.pkl"))
          )
          pickle.dump(
              self, open(os.path.join(SAVE_PATH, "evolver.pkl"), "wb"), protocol=4
          )
          self.env = ENV
  
      def init_env(self):
          """Initialize the PCGRL level-generation RL environment and extract any useful info from it."""
  
          env_name = "{}-{}-v0".format(PROBLEM, REPRESENTATION)
          self.env = gym.make(env_name)
          self.env.adjust_param(render=RENDER)
  
          if CMAES:
              # Give a little wiggle room from targets, to allow for some diversity
  
              if "binary" in PROBLEM:
                  path_trg = self.env._prob.static_trgs["path-length"]
                  self.env._prob.static_trgs.update(
                      {"path-length": (path_trg - 20, path_trg)}
                  )
              elif "zelda" in PROBLEM:
                  path_trg = self.env._prob.static_trgs["path-length"]
                  self.env._prob.static_trgs.update(
                      {"path-length": (path_trg - 40, path_trg)}
                  )
              elif "sokoban" in PROBLEM:
                  sol_trg = self.env._prob.static_trgs["sol-length"]
                  self.env._prob.static_trgs.update(
                      {"sol-length": (sol_trg - 10, sol_trg)}
                  )
              elif "smb" in PROBLEM:
                  pass
              else:
                  raise NotImplemented
  
          global N_DIRS
  
          if hasattr(self.env._rep, "_dirs"):
              N_DIRS = len(self.env._rep._dirs)
          else:
              N_DIRS = 0
  
          global N_STEPS
  
          #       if N_STEPS is None:
          #       if REPRESENTATION != "cellular":
          max_ca_steps = args.n_steps
          max_changes = self.env._prob._height * self.env._prob._width
          reps_to_steps = {
              "cellular": max_ca_steps,
              "wide": max_changes,
              #           "narrow": max_changes,
              "narrow": max_changes,
              #           "turtle": max_changes * 2,
              "turtle": 2 * max_changes,
              # So that it can move around to each tile I guess
          }
          N_STEPS = reps_to_steps[REPRESENTATION]
  
      def visualize(self, itr=None):
          archive = self.gen_archive
          # # Visualize Result
          #       grid_archive_heatmap(archive, vmin=self.reward_bounds[self.reward_names[0]][0], vmax=self.reward_bounds[self.reward_names[0]][1])
          #       if PROBLEM == 'binary':
          #           vmin = -20
          #           vmax = 20
          #       elif PROBLEM == 'zelda':
          #           vmin = -20
          #           vmax = 20
          #       grid_archive_heatmap(archive, vmin=vmin, vmax=vmax)
          if ALGO == "ME":
              obj_min, obj_max = archive.fitness_extrema[0]
              qdpy.plots.plotGridSubplots(archive.quality_array[..., 0], os.path.join(SAVE_PATH, 'fitness.pdf'),
                                          plt.get_cmap("inferno_r"), archive.features_domain,
                                          archive.fitness_domain[0], nbTicks=None)
          else:
              plt.figure(figsize=(8, 6))
              df_obj = archive.as_pandas()["objective"]
              obj_min = df_obj.min()
              obj_max = df_obj.max()
              vmin = np.floor(obj_min)
              vmax = np.ceil(obj_max)
              grid_archive_heatmap(archive, vmin=vmin, vmax=vmax)
              label_fontdict = {
                  'fontsize': 16,
              }
              if not CMAES:
                  plt.xlabel(self.bc_names[0], fontdict=label_fontdict)
                  plt.ylabel(self.bc_names[1], fontdict=label_fontdict)
              if itr is not None:
                  save_path = os.path.join(SAVE_PATH, "checkpoint_{}".format(itr))
              else:
                  save_path = SAVE_PATH
              plt.title('fitness', fontdict={'fontsize': 24})
              plt.tight_layout()
              plt.savefig(os.path.join(save_path, "fitness.png"))
              #       plt.gca().invert_yaxis()  # Makes more sense if larger BC_1's are on top.
  
  
              if SHOW_VIS:
                  plt.show()
              plt.close()
  
          # Print table of results
  #       df = archive.as_pandas()
          # high_performing = df[df["objective"] > 200].sort_values("objective", ascending=False)
  
      #       print(df)
  
      def infer(self, concat_gifs=True):
          assert INFER
          self.init_env()
          archive = self.gen_archive
          if args.algo == "ME":
              nonempty_idxs = np.stack(np.where(
                  np.isnan(archive.quality_array) == False), axis=1)
              # Assume 2nd BC is a measure of complexity
              # Sort according to 2nd BC
              idxs = nonempty_idxs.tolist()
              idxs.sort(key=lambda x: x[1])
              idxs_T = tuple(np.array(idxs).T)
              objs = archive.quality_array[idxs_T]
              # Get list of individuals in same order. First get list of features belonging to individuals in bin,
              # then get individual by bin-coordinate
              bcs = [archive.features[tuple(idx[:-1])][idx[-1]].values for idx in idxs]
              models = [archive.solutions[tuple(idx[:-1])][idx[-1]] for idx in idxs]
              # Get rid of bin coordinate for our purposes
              # TODO: for more flexibility, instead adapt the below to get this bin coordinate
              idxs = [idx[:-1] for idx in idxs]
          else:
              df = archive.as_pandas()
              rows = df.sort_values("behavior_1", ascending=False)
              models = np.array(rows.loc[:, "solution_0":])
              bcs_0 = np.array(rows.loc[:, "behavior_0"])
              bcs_1 = np.array(rows.loc[:, "behavior_1"])
              objs = np.array(rows.loc[:, "objective"])
              # FIXME: don't need these
              idxs = np.array(rows.loc[:, "index_0":"index_1"])
          global N_INIT_STATES
          global N_EVAL_STATES
          global RENDER
          global RANDOM_INIT_LEVELS
  
          if RENDER_LEVELS:
              RENDER = False
              #           N_INIT_STATES = 1
  
              if "smb" in PROBLEM:
                  d = 4
                  figw, figh = 32, 4
              elif "zelda" in PROBLEM:
                  d = 3
                  figw, figh = self.env._prob._width, self.env._prob._height
              else:
                  d = 6  # number of rows and columns
                  figw, figh = self.env._prob._width, self.env._prob._height
  
              if CMAES:
                  n_rows = 2
                  n_cols = 5
                  n_figs = n_rows * d
                  fig, axs = plt.subplots(
                      ncols=d,
                      nrows=n_rows,
                      figsize=(figw * n_cols / d, figh * n_rows / d),
                  )
                  df_g = df.sort_values(by=["objective"], ascending=False)
                  grid_models = np.array(df_g.loc[:, "solution_0":])
                  level_frames = []
  
                  for (i, model) in enumerate(grid_models):
                      for j in range(n_figs):
                          n_row = j // d
                          n_col = j % d
                          axs[n_row, n_col].set_axis_off()
                          # TODO: select for diversity?
                          # parallelization would be kind of pointelss here
                          init_nn = set_weights(self.gen_model, model)
                          # run simulation, but only on a single level-seed
  #                       init_state = gen_random_levels(1, self.env)
                          #                       init_state = np.random.randint(
                          #                           0, self.n_tile_types, size=(1, *self.init_states.shape[1:])
                          #                       )
                          #                       _, _, _, (
                          #                           time_penalty,
                          #                           targets_penalty,
                          #                           variance_penalty,
                          #                           diversity_bonus,
                          #                       ) = simulate(
  
                          raise NotImplementedError
                          TT()  # don't have a way of rendering CMAES yet??
                          level_frames_i = simulate(
                              self.env,
                              init_nn,
                              self.n_tile_types,
                              self.init_states[0:1],
                              self.bc_names,
                              self.static_targets,
                              seed=None,
                              render_levels=True,
                          )
                          if not concat_gifs:
                              save_level_frames(level_frames_i, i)
                          else:
                              level_frames += level_frames_i
                          # Get image
  #                       img = self.env.render(mode="rgb_array")
                          img = level_frames[-1]
                          axs[n_row, n_col].imshow(img, aspect=1)
                  if concat_gifs:
                      save_level_frames(level_frames, 'concat')
  
              else:
                  fig, axs = plt.subplots(ncols=d, nrows=d, figsize=(figw, figh))
                  if ALGO == "ME":
                      pass
                  else:
                      df_g = df.sort_values(by=["behavior_0", "behavior_1"], ascending=False)
  
                  df_g["row"] = np.floor(
                      np.linspace(0, d, len(df_g), endpoint=False)
                  ).astype(int)
                  level_frames = []
  
                  for row_num in range(d):
                      row = df_g[df_g["row"] == row_num]
                      row = row.sort_values(by=["behavior_1"], ascending=True)
                      row["col"] = np.arange(0, len(row), dtype=int)
                      idx = np.floor(np.linspace(0, len(row) - 1, d)).astype(int)
                      row = row[row["col"].isin(idx)]
                      row = row.drop(["row", "col"], axis=1)
                      grid_models = np.array(row.loc[:, "solution_0":])
  
                      for col_num in range(len(row)):
                          model = grid_models[col_num]
  #                       axs[row_num, col_num].set_axis_off()
                          axs[-col_num-1, -row_num-1].set_axis_off()
  
                          # initialize weights
                          gen_model = set_weights(self.gen_model, model)
  
                          # run simulation, but only on the first level-seed
                          #                       _, _, _, (
                          #                           time_penalty,
                          #                           targets_penalty,
                          #                           variance_penalty,
                          #                           diversity_bonus,
                          #                       ) = simulate(
                          level_frames_i = simulate(
                              self.env,
                              gen_model,
                              self.n_tile_types,
                              self.init_states[0:1],
                              self.bc_names,
                              self.static_targets,
                              seed=None,
                              render_levels=True,
                          )
                          if not concat_gifs:
                              save_level_frames(level_frames_i, '{}_{}'.format(row_num, col_num))
                          level_frames += level_frames_i
                          # Get image
                          #                       img = self.env.render(mode="rgb_array")
                          img = level_frames[-1]
  #                       axs[row_num, col_num].imshow(img, aspect="auto")
                          axs[-col_num-1, -row_num-1].imshow(img, aspect="auto")
                  if concat_gifs:
                      save_level_frames(level_frames, 'concat')
              fig.subplots_adjust(hspace=0.01, wspace=0.01)
              plt.tight_layout()
              fig.savefig(
                  os.path.join(SAVE_PATH, "levelGrid_{}-bin.png".format(d)), dpi=300
              )
              plt.close()
  
          if PLAY_LEVEL:
              player_simulate(
                  self.env,
                  self.n_tile_types,
                  self.play_bc_names,
                  self.play_model,
                  playable_levels=self.playable_levels,
                  seed=None,
              )
          i = 0
  
          if EVALUATE:
              # First, visualize and aggregate the scores of the elites as they currently stand in the grid
  
              if not VISUALIZE:
                  # visualize if we haven't already
                  self.visualize()
              # aggregate scores of individuals currently in the grid
              save_train_stats(objs, archive, self.env, self.bc_names)
  
              # The level spaces which we will attempt to map to
              problem_eval_bc_names = {
                  "binary": [
  #                   ("regions", "path-length")
                        ],
                  "zelda": [
  #                   ("nearest-enemy", "path-length"),
  #                   ("symmetry", "path-length"),
  #                   ("emptiness", "path-length"),
                  ],
                  "sokoban": [
  #                   ("crate", "sol-length")
                  ],
                  "smb": [
  #                   ("emptiness", "jumps")
                  ],
              }
  
  #           for k in problem_eval_bc_names.keys():
  #               problem_eval_bc_names[k] += [
  #                   # ("NONE"),
  #                   ("emptiness", "symmetry")
  #               ]
  
              for (k, v) in problem_eval_bc_names.items():
                  if k in PROBLEM:
                      eval_bc_names = v
  
                      break
              # toss our elites into an archive with different BCs. For fun!
              eval_bc_names = list(set([tuple(self.bc_names)] + eval_bc_names))
  
              if not CMAES:
                  if ALGO == "ME":
                      eval_archives = [
                          MEGrid(
                              [N_BINS for _ in eval_bcs],
                              [self.bc_bounds[bc_name] for bc_name in eval_bcs],
                          )
                          for eval_bcs in eval_bc_names
                      ]
                  else:
                      eval_archives = [
                          GridArchive(
                              # minimum of 100 for each behavioral characteristic, or as many different values as the BC can take on, if it is less
                              # [min(100, int(np.ceil(self.bc_bounds[bc_name][1] - self.bc_bounds[bc_name][0]))) for bc_name in self.bc_names],
                              [N_BINS for _ in eval_bcs],
                              # min/max for each BC
                              [self.bc_bounds[bc_name] for bc_name in eval_bcs],
                          )
                          for eval_bcs in eval_bc_names
                      ]
                      [
                          eval_archive.initialize(solution_dim=len(models[0]))
                          for eval_archive in eval_archives
                      ]
  
              RENDER = False
              # Iterate through our archive of trained elites, evaluating them and storing stats about them.
              # Borrowing logic from grid_archive_heatmap from pyribs.
  
              # Retrieve data from archive
              if ALGO == 'ME':
                  lower_bounds = [archive.features_domain[i][0] for i in range(len(archive.features_domain))]
                  upper_bounds = [archive.features_domain[i][1] for i in range(len(archive.features_domain))]
                  x_dim, y_dim = archive.shape
              else:
                  lower_bounds = archive.lower_bounds
                  upper_bounds = archive.upper_bounds
                  x_dim, y_dim = archive.dims
              x_bounds = np.linspace(lower_bounds[0], upper_bounds[0], x_dim + 1)
              y_bounds = np.linspace(lower_bounds[1], upper_bounds[1], y_dim + 1)
  
              # Color for each cell in the heatmap
              fitness_scores = np.full((y_dim, x_dim), np.nan)
              playability_scores = np.full((y_dim, x_dim), np.nan)
              diversity_scores = np.full((y_dim, x_dim), np.nan)
              reliability_scores = np.full((y_dim, x_dim), np.nan)
              eval_fitness_scores = []
              eval_playability_scores = []
              eval_diversity_scores = []
              eval_reliability_scores = []
  
              if not CMAES:
                  for j in range(len(eval_archives)):
                      eval_fitness_scores.append(np.full((y_dim, x_dim), np.nan))
                      eval_playability_scores.append(np.full((y_dim, x_dim), np.nan))
                      eval_reliability_scores.append(np.full((y_dim, x_dim), np.nan))
                      eval_diversity_scores.append(np.full((y_dim, x_dim), np.nan))
  
              def record_scores(
                  id_0,
                  id_1,
                  batch_reward,
                  targets_penalty,
                  diversity_bonus,
                  variance_penalty,
                  fitness_scores,
                  playability_scores,
                  diversity_scores,
                  reliability_scores,
              ):
                  fitness_scores[id_0, id_1] = batch_reward
                  playability_scores[id_0, id_1] = targets_penalty
  
                  if diversity_bonus is not None:
                      diversity_scores[id_0, id_1] = diversity_bonus
  
                  if variance_penalty is not None:
                      reliability_scores[id_0, id_1] = variance_penalty
  
              def save_levels(level_json, overwrite=False, headers=False):
                  df = pd.DataFrame.from_dict(level_json
                                    )
                  #               df = df[df['targets'] == 0]
  
                  if overwrite:
                      write_mode = "w"
                  else:
                      write_mode = "a"
  
                  if len(df) > 0:
                      csv_name = "eval_levels"
  
                      if not RANDOM_INIT_LEVELS:
                          csv_name += "_fixLvls"
                      csv_name += ".csv"
                      if headers:
                          header = df.columns
                      else:
                          header = None
                      df.to_csv(
                          os.path.join(SAVE_PATH, csv_name),
                          mode=write_mode,
                          header=header,
                          index=False,
                      )
  
              init_states_archive = None
  
              if RANDOM_INIT_LEVELS:
                  # Effectively doing inference on a (presumed) held-out set of levels
  
                  if CMAES:
                      N_EVAL_STATES = N_INIT_STATES = 100
                  else:
                      N_EVAL_STATES = N_INIT_STATES = 20  #= 100  # e.g. 10
  
                  init_states = gen_random_levels(N_INIT_STATES, self.env)
              #               init_states = np.random.randint(
              #                   0,
              #                   self.n_tile_types,
              #                   size=(N_EVAL_STATES, *self.init_states.shape[1:]),
              #               )
              elif args.fix_level_seeds or args.n_init_states == 0:
                  # If level seeds were fixed throughout training, use those
                  init_states = self.init_states
                  N_EVAL_STATES = N_INIT_STATES = init_states.shape[0]
              else:
                  init_states_archive = self.gen_archive.init_states_archive
                  init_states = None
                  # Otherwise, use the init level seeds that were entered into the archive with each elite
  
              n_train_bcs = len(self.bc_names)
  
              if THREADS:
                  futures = [
                      multi_evo.remote(
                          self.env,
                          self.gen_model,
                          model_w,
                          self.n_tile_types,
                          init_states,
                          [bc for bc_names in eval_bc_names for bc in bc_names],
                          self.static_targets,
                          seed,
                          player_1=self.player_1,
                          player_2=self.player_2,
                          proc_id=i,
                          init_states_archive=init_states_archive,
                          index=tuple(idxs[i]),
                      )
                      for (i, model_w) in enumerate(models)
                  ]
                  results = ray.get(futures)
                  i = 0
  
                  for result in results:
  
                      level_json, batch_reward, final_bcs, (
                          time_penalty,
                          batch_targets_penalty,
                          variance_penalty,
                          diversity_bonus,
                      ) = result
                      # id_0 = idxs_0[i]
                      # id_1 = idxs_1[i]
                      grid_bcs = final_bcs[:n_train_bcs]
                      # TODO: remove this (it's for backward compatibility) since we've implemented get_index for qdpy
                      #   grid
                      if ALGO == "ME":
                          id_0, id_1 = archive.index_grid(tuple(grid_bcs))
                      else:
                          id_0, id_1 = archive.get_index(np.array(grid_bcs))
  
                      if SAVE_LEVELS:
                          save_levels(level_json, overwrite=i == 0, headers=i==0)
                      # Record directly from evolved archive since we are guaranteed to have only one elite per cell
                      record_scores(
                          id_0,
                          id_1,
                          batch_reward,
                          batch_targets_penalty,
                          diversity_bonus,
                          variance_penalty,
                          fitness_scores,
                          playability_scores,
                          diversity_scores,
                          reliability_scores,
                      )
  
                      if not CMAES:
                          for j, eval_archive in enumerate(eval_archives):
                              # Record componentes of the fitness for each cell in each evaluation archive
                              # NOTE: assume 2 BCs per eval archive
                              eval_bcs = np.array(
  #                               final_bcs[n_train_bcs + 2 * j : n_train_bcs + 2 * j + 2]
                                  final_bcs[2 * j: 2 * (j + 1)]
                              )
                              if ALGO == "ME":
                                  id_0, id_1 = archive.index_grid(tuple(eval_bcs))
                                  # Dummy individual
                                  individual = Individual(type(self.gen_model), self.n_tile_types, self.n_tile_types)
                                  individual.fitness = Fitness([batch_reward])
                                  individual.features = Features(final_bcs)
                                  idx = eval_archive.add(individual)
                                  ind_added = idx is not None
                              else:
                                  id_0, id_1 = eval_archive.get_index(eval_bcs)
                                  # Add dummy solution weights for now
                                  status, _ = eval_archive.add(
                                      np.zeros(eval_archive.solution_dim),
                                      batch_reward,
                                      eval_bcs,
                                  )
                                  ind_added = status != AddStatus.NOT_ADDED
  
                              if ind_added:
                                  # For eval archive, only record new best individuals in each filled cell
                                  record_scores(
                                      id_0,
                                      id_1,
                                      batch_reward,
                                      batch_targets_penalty,
                                      diversity_bonus,
                                      variance_penalty,
                                      eval_fitness_scores[j],
                                      eval_playability_scores[j],
                                      eval_diversity_scores[j],
                                      eval_reliability_scores[j],
                                  )
                      i += 1
                  auto_garbage_collect()
  
              else:
                  # NOTE: Note maintaining this single-threaded code at the moment, can refactor and bring it up to date later
  
                  while i < len(models):
                      # iterate through all models and record stats, on either training seeds or new ones (to test evaluation)
                      model = models[i]
                      id_0, id_1 = idxs[i]
  
                      if init_states is None:
                          init_states_archive = archive.init_states_archive
                      else:
                          init_states_archive = None
  
                      if init_states is None:
                          init_states = get_init_states(
                              init_states_archive, tuple(idxs[i])
                          )
  
                      gen_model = set_weights(self.gen_model, model)
                      level_json, batch_reward, final_bcs, (
                          time_penalty,
                          targets_penalty,
                          variance_penalty,
                          diversity_bonus,
                      ) = simulate(
                          env=self.env,
                          model=gen_model,
                          n_tile_types=self.n_tile_types,
                          init_states=init_states,
                          bc_names=self.bc_names,
                          static_targets=self.static_targets,
                          seed=None,
                          player_1=self.player_1,
                          player_2=self.player_2,
                      )
  
                      if SAVE_LEVELS:
                          save_levels(level_json)
                      record_scores(
                          id_0,
                          id_1,
                          batch_reward,
                          targets_penalty,
                          diversity_bonus,
                          variance_penalty,
                          fitness_scores,
                          playability_scores,
                          diversity_scores,
                          reliability_scores,
                      )
  
              if ALGO == "ME":
                  n_filled_bins = eval_archive.filled_bins
                  assert len(models) == archive.filled_bins
                  n_total_bins = archive.size
              else:
                  n_filled_bins = len(eval_archive._occupied_indices)
                  assert len(models) == len(archive._occupied_indices)
                  n_total_bins = archive.bins
              qd_score = get_qd_score(archive, self.env, self.bc_names)
              eval_qd_score = get_qd_score(eval_archive, self.env, self.bc_names)
              stats = {
                  "generations completed": self.n_itr,
                  "% train archive full": len(models) / n_total_bins,
                  "archive size": n_filled_bins,
                  "QD score": qd_score,
                  "eval QD score": eval_qd_score,
                  "% eval archives full": {},
                  "eval archive sizes": {},
                  "eval QD scores": {},
              }
  
              if not CMAES:
                  plot_args = {
                      'lower_bounds': lower_bounds,
                      'upper_bounds': upper_bounds,
                      'x_bounds': x_bounds,
                      'y_bounds': y_bounds,
                  }
                  plot_score_heatmap(playability_scores, "playability", self.bc_names, **plot_args,
                                     bcs_in_filename=False)
                  plot_score_heatmap(diversity_scores / 10, "diversity", self.bc_names, **plot_args, bcs_in_filename=False)
                  plot_score_heatmap(reliability_scores, "reliability", self.bc_names, **plot_args, bcs_in_filename=False)
                  plot_score_heatmap(fitness_scores, "fitness_eval", self.bc_names, **plot_args, bcs_in_filename=False)
  
                  for j, eval_archive in enumerate(eval_archives):
                      bc_names = eval_bc_names[j]
  
                      if bc_names != ("NONE") and bc_names != tuple(self.bc_names):
                          plot_score_heatmap(
                              eval_playability_scores[j], "playability", bc_names, **plot_args,
                          )
                          plot_score_heatmap(
                              eval_diversity_scores[j] / 10, "diversity", bc_names, **plot_args,
                          )
                          plot_score_heatmap(
                              eval_reliability_scores[j], "reliability", bc_names, **plot_args,
                          )
                          plot_score_heatmap(
                              eval_fitness_scores[j], "fitness_eval", bc_names, **plot_args,
                          )
  
                      if bc_names == tuple(self.bc_names):
                          # in case a bug appears here, where performance differs from training to inference,
                          # include this redundant data to try and pinpoint it. Note that this is only redundant in
                          # stats_fixLvls, though, because otherwise, we are doing evaluation in the same BC space.
                          pct_archive_full = (
                              n_filled_bins / n_total_bins
                          )
  
                          if not RANDOM_INIT_LEVELS:
                              # then this will be the same as the
                              #                           if not len(eval_archive._occupied_indices) / eval_archive.bins == stats["% train archive full"]:
                              #                           continue
                              pass
                          else:
                              pass
                          stats["% elites maintained"] = (
                              pct_archive_full / stats["% train archive full"]
                          )
                          stats["% QD score maintained"] = stats["eval QD score"] / stats["QD score"]
  
                          stats["% fresh train archive full"] = pct_archive_full
                          stats["% fresh train archive full"] = pct_archive_full
                      n_occupied = n_filled_bins
  #                   assert n_occupied == len(eval_archive._occupied_indices)
                      bcs_key = "-".join(bc_names)
                      stats["% eval archives full"].update(
                          {
                              bcs_key: n_occupied / n_total_bins,
                      })
                      stats["eval archive sizes"].update({
                          bcs_key: n_occupied,
                      })
                      stats["eval QD scores"].update({
                          bcs_key: get_qd_score(eval_archive, self.env, bc_names)
                      })
  
              stats.update(
                  {
                      "playability": get_stats(playability_scores),
                      "diversity": get_stats(diversity_scores / 10),
                      "reliability": get_stats(reliability_scores),
                  }
              )
              f_name = "stats"
  
              if not RANDOM_INIT_LEVELS:
                  f_name = f_name + "fixLvls"
              f_name += ".json"
              with open(os.path.join(SAVE_PATH, f_name), "w", encoding="utf-8") as f:
                  json.dump(stats, f, ensure_ascii=False, indent=4)
  
              return
  
          while i < len(models):
              #           model = self.archive.get_random_elite()[0]
              #           model = models[np.random.randint(len(models))]
  
              model = models[i]
              gen_model = set_weights(self.gen_model, model)
  
              #           RANDOM_INIT_LEVELS = not opts.fix_level_seeds
  
              if RANDOM_INIT_LEVELS and args.n_init_states != 0:
                  init_states = gen_random_levels(N_INIT_STATES, self.env)
              elif not args.fix_level_seeds and args.n_init_states != 0:
                  init_states_archive = archive.init_states_archive
                  init_states = get_init_states(init_states_archive, tuple(idxs[i]))
              else:
                  init_states = self.init_states
              _, _, _, (
                  time_penalty,
                  targets_penalty,
                  variance_penalty,
                  diversity_bonus,
              ) = simulate(
                  self.env,
                  gen_model,
                  self.n_tile_types,
                  init_states,
                  self.bc_names,
                  self.static_targets,
                  seed=None,
                  player_1=self.player_1,
                  player_2=self.player_2,
              )
              #           input("Mean behavior characteristics:\n\t{}: {}\n\t{}: {}\nMean reward:\n\tTotal: {}\n\ttime: {}\n\ttargets: {}\n\tvariance: {}\n\tdiversity: {}\nPress any key for next generator...".format(
              #               self.bc_names[0], bcs_0[i], self.bc_names[1], bcs_1[i], objs[i], time_penalty, targets_penalty, variance_penalty, diversity_bonus))
              i += 1
  
  
  #           if i == len(models):
  #               i=0
  
  
  def gen_random_levels(n_init_states, env):
      init_states = np.random.randint(
          0, len(env._prob.get_tile_types()), (N_INIT_STATES, env._prob._height, env._prob._width)
      )
      return init_states
  
  #   init_states = np.zeros(shape=(n_init_states, env._prob._height, env._prob._width))
  #   init_state_maps = []
  
  #   for i in range(N_INIT_STATES):
  #       env._rep.reset(
  #           env._prob._width,
  #           env._prob._height,
  #           get_int_prob(env._prob._prob, env._prob.get_tile_types()),
  #       )
  #       #                   init_state_maps.append(np.expand_dims(get_one_hot_map(self.env._rep._map, self.n_tile_types), axis=0))
  #       init_state_maps.append(np.expand_dims(env._rep._map, axis=0))
  
  #   init_states[:] = np.vstack(init_state_maps)
  #   # init_states = np.zeros(
  #   #    0, self.n_tile_types, size=self.init_states.shape
  #   # )
  
  #   return init_states
  
  
  if __name__ == "__main__":
      """
      Set Parameters
      """
      N_BINS = 100
      CA_ACTION = True
      args, arg_dict = get_args()
      global INFER
      global EVO_DIR
      global CUDA
      global RENDER
      global PROBLEM
      global SHOW_VIS
      global VISUALIZE
      global N_STEPS
      global N_GENERATIONS
      global N_INIT_STATES
      global N_INFER_STEPS
      global BCS
      global RENDER_LEVELS
      global THREADS
      global PLAY_LEVEL
      global CMAES
      global EVALUATE
      global SAVE_LEVELS
      global RANDOM_INIT_LEVELS
      global CASCADE_REWARD
      global REPRESENTATION
      global MODEL
      global REEVALUATE_ELITES
      global preprocess_action
      global N_PROC
      global ALGO
      global seed
      CONCAT_GIFS = False
      if arg_dict["exp_name"] == '5':
          seed = 420
      else:
          try:
              seed = int(arg_dict["exp_name"])
          except Exception:
              print("Assigning random seed")
              seed = np.random.randint()
      print("Random number seed is: {}".format(seed))
      N_PROC = arg_dict["n_cpu"]
      MODEL = arg_dict["model"]
      ALGO = arg_dict["algo"]
      if ALGO == "ME":
          # TODO: implement wrapper around other models generically
          assert MODEL in ["CPPN", "GenCPPN", "CPPNCA"]
      else:
          assert ALGO == "CMAME"
      REPRESENTATION = arg_dict["representation"]
      CASCADE_REWARD = arg_dict["cascade_reward"]
      REEVALUATE_ELITES = not arg_dict["fix_elites"] and arg_dict["n_init_states"] != 0
      RANDOM_INIT_LEVELS = (
          not arg_dict["fix_level_seeds"]
          and arg_dict["n_init_states"] != 0
          or REEVALUATE_ELITES
      )
  
      if REEVALUATE_ELITES:
          # Otherwise there is no point in re-evaluating them
          assert RANDOM_INIT_LEVELS
      CMAES = arg_dict["behavior_characteristics"] == ["NONE"]
      EVALUATE = arg_dict["evaluate"]
      PLAY_LEVEL = arg_dict["play_level"]
      BCS = arg_dict["behavior_characteristics"]
      N_GENERATIONS = arg_dict["n_generations"]
      N_INIT_STATES = arg_dict["n_init_states"]
      N_STEPS = arg_dict["n_steps"]
  
      SHOW_VIS = arg_dict["show_vis"]
      PROBLEM = arg_dict["problem"]
      CUDA = False
      VISUALIZE = arg_dict["visualize"]
      INFER = arg_dict["infer"] or EVALUATE
      N_INFER_STEPS = N_STEPS
      #   N_INFER_STEPS = 100
      RENDER_LEVELS = arg_dict["render_levels"]
      THREADS = arg_dict["multi_thread"]  # or EVALUATE
      SAVE_INTERVAL = arg_dict["save_interval"]
      VIS_INTERVAL = 50
      if "CPPN" in MODEL:
          if MODEL != "CPPNCA" and "Gen" not in MODEL:
              assert N_INIT_STATES == 0 and not RANDOM_INIT_LEVELS and not REEVALUATE_ELITES
          if MODEL != "CPPNCA":
              assert N_STEPS == 1
  
      SAVE_LEVELS = arg_dict["save_levels"] or EVALUATE
  
      #   exp_name = 'EvoPCGRL_{}-{}_{}_{}-batch_{}-step_{}'.format(PROBLEM, REPRESENTATION, BCS, N_INIT_STATES, N_STEPS, arg_dict['exp_name'])
      #   exp_name = "EvoPCGRL_{}-{}_{}_{}_{}-batch".format(
      #       PROBLEM, REPRESENTATION, MODEL, BCS, N_INIT_STATES
      #   )
      exp_name = "EvoPCGRL_"
      if ALGO == "ME":
          exp_name += "ME_"
      exp_name += "{}-{}_{}_{}_{}-batch_{}-pass".format(
          PROBLEM, REPRESENTATION, MODEL, BCS, N_INIT_STATES, N_STEPS
      )
  
      if CASCADE_REWARD:
          exp_name += "_cascRew"
  
      if not RANDOM_INIT_LEVELS:
          exp_name += "_fixLvls"
  
      if not REEVALUATE_ELITES:
          exp_name += "_fixElites"
  
      if args.mega:
          exp_name += "_MEGA"
      exp_name += "_" + arg_dict["exp_name"]
      SAVE_PATH = os.path.join("evo_runs", exp_name)
      if MODEL not in preprocess_action_funcs:
          if "CPPN" in MODEL:
              preprocess_action = preprocess_action_funcs['CPPN'][REPRESENTATION]
          else:
              preprocess_action = preprocess_action_funcs['NCA'][REPRESENTATION]
      else:
          preprocess_action = preprocess_action_funcs[MODEL][REPRESENTATION]
      if MODEL not in preprocess_observation_funcs:
          preprocess_observation = preprocess_observation_funcs['NCA'][REPRESENTATION]
      else:
          preprocess_observation = preprocess_observation_funcs[MODEL][REPRESENTATION]
  
      def init_tensorboard():
          assert not INFER
          # Create TensorBoard Log Directory if does not exist
          LOG_NAME = "./runs/" + datetime.now().strftime("%Y%m%d-%H%M%S") + "-" + exp_name
          writer = SummaryWriter(LOG_NAME)
  
          return writer
  
  
      if THREADS:
          ray.init()
  
      try:
          try:
              evolver = pickle.load(open(os.path.join(SAVE_PATH, "evolver.pkl"), "rb"))
          except:
              evolver = pickle.load(
                  open(os.path.join(SAVE_PATH, "last_evolver.pkl"), "rb")
              )
          print("Loaded save file at {}".format(SAVE_PATH))
  
          if VISUALIZE:
              evolver.visualize()
  
          if INFER:
              global RENDER
              RENDER = True
              N_STEPS = N_INFER_STEPS
  
              #           if not RANDOM_INIT_LEVELS:
              # evaluate on initial level seeds that each generator has seen before
              RANDOM_INIT_LEVELS = False
              evolver.infer(concat_gifs=CONCAT_GIFS)
              save_grid(csv_name="eval_levels_fixLvls")
              # evaluate on random initial level seeds
              RANDOM_INIT_LEVELS = True
              evolver.infer(concat_gifs=CONCAT_GIFS)
              save_grid(csv_name="eval_levels")
          #           save_grid(csv_name="levels")
  
          if not (INFER or VISUALIZE):
              writer = init_tensorboard()
              # then we train
              RENDER = arg_dict["render"]
              evolver.init_env()
              evolver.total_itrs = arg_dict["n_generations"]
              evolver.evolve()
      except FileNotFoundError as e:
          if not INFER:
              RENDER = arg_dict["render"]
              print(
                  "Failed loading from an existing save-file. Evolving from scratch. The error was: {}".format(
                      e
                  )
              )
              writer = init_tensorboard()
              evolver = EvoPCGRL()
              evolver.evolve()
          else:
              print(
                  "Loading from an existing save-file failed. Cannot run inference. The error was: {}".format(
                      e
                  )
              )
  