import os

import random

from time import time

import matplotlib.pyplot as plt

import numpy as np

import pandas as pd

import torch

from gymnasium import spaces

# Mappings for changing direction if reflected.

# Cannot cross a row boundary moving right or left.

ROW_CROSSING = {1: 2, 3: -2, 5: 1, 6: -1, 7: 1, 8: -1}

# Cannot cross a column boundary moving up or down.

COL_CROSSING = {2: 2, 4: -2, 5: 3, 6: 1, 7: -1, 8: -3}

class Dot:

def __init__(self, r: int = 0, c: int = 0, t: int = 1) -> None:

# Initialize current point and tail to initial point.

self.row = np.ones(t, dtype="int32") * r

self.col = np.ones(t, dtype="int32") * c

def move(self, r: int, c: int):

"""

Cycle path history and set new current coordinates.

:param r: row

:param c: column

"""

# Cycle the path history.

for t in reversed(range(1, len(self.row))):

self.row[t] = self.row[t - 1]

self.col[t] = self.col[t - 1]

# Set new current point

self.row[0] = r

self.col[0] = c

class DotSimulator:

"""

This simulator lets us generate dots and make them move.

It's especially useful in keeping entitled cats occupied,

but instead of feline neurons, we use this for fake ones.

Specifically, this generates a grid for each timestep, where a specified

number of points have values of 1 with fading tails ("decay"), designating

the current positions and movements of their corresponding dots. All other

points are set to 0. From timestep to timestep, the dots either remain

where they are or move one space.

The 2D observation of the current state is provided every step, as well as

the reward, completion flag, and sucessful interception flag. It may be

helpful to amplify the grid values when encoding them as spike trains.

:param t: int: number of timesteps/samples of grids with dot movements

:param height: int: height dimension of the grid (rows)

:param width: int: width dimension of the grid (columns)

:param decay: int: length of decaying tail behind a dot (its path history)

:param dots: int: number of target dots

:param herrs: number of distraction dots (red herrings)

:param pandas: Bool: print as pandas dataframe versus graphical plots

:param write: Bool: write grids to file to be plotted later.

:param mute: Bool: mute graphical rendering (can write to file or print pandas)

:param speed: int: set movement speed of dots.

:param randr: float: set the randomization rate of target movements.

:param allow_stay: bool: allow a dot to remain in place as a movement choice.

:param seed: int: optional seed for RNG in movement generation.

:param fname: string: optional filename for saving grids to file

:param fpath: string: optional file path for saving grids to file

:param diag: Bool: allow diagonal movement.

:param bound_hand: str: bounds handling when a dot reaches the world's end.

'stay': dots will simply be prevented from crossing the edges.

'bounce': dot positions and directions will be reflected.

'trans': dot positions will be mirrored to the opposite edge.

:param fit_func: str: Fitness function.

'euc': Single Euclidean (Pythagorean) distance value

'disp': Tuple of x,y displacement values

'rng' : Range rings--the closer the ring, the lower the number

'dir' : directional--+1 if moving in the right direction

-1 if moving in the wrong direction

0 if neither.

:param ring_size: int: set range ring size for range ring fitness function.

:param bullseye: int: set reward for successful intercept; default = 10.0

:param teleport: Bool: teleport network dot after intercept; default = true

"""

def __init__(self, t: int, **kwargs) -> None:

self.timesteps = t # total timesteps

self.ts = 0 # initialize current timestep to 0

""" Keyword arguments """

self.h = kwargs.get("height", 28) # height dimension

self.w = kwargs.get("width", 28) # width dimension

self.decay = kwargs.get("decay", 1) # length of a dot's tail (path history)

self.ndots = kwargs.get("dots", 1) # Number of dots

self.herrs = kwargs.get("herrs", 0) # Red herrings (distractions)

self.pandas = kwargs.get("pandas", False) # print as pandas DF

self.write2F = kwargs.get("write", False) # write grids to file

self.mute = kwargs.get("mute", False) # mute displayed rendering

self.speed = kwargs.get("speed", 1) # dot movement speed

self.randr = kwargs.get("randr", 1.0) # rate of random movement

self.minch = int(not kwargs.get("allow_stay", True)) # allow stay choice.

# Grab system time and trim off extra large parts of the number.

sysTime = time()

sysTime = int(1e10 * (sysTime - 1e6 * (sysTime // 1e6)))

# Save off RNG seed.

self.seed = kwargs.get("seed", 0)

if self.seed == 0:

self.seed = sysTime

# Create filename if one isn't provided.

path = kwargs.get("fpath", "out")

self.filename = kwargs.get("fname", "grids_s" + str(self.seed) + ".csv")

self.fileCnt = 0

self.filename = (

path

+ "/"

+ self.filename[:-5]

+ "_"

+ str(self.fileCnt)

+ self.filename[-4:]

)

if not os.path.exists(path):

os.makedirs(path)

# Expand movement options if diagonal is allowed.

if kwargs.get("diag", False):

self.choices = 9

else:

self.choices = 5

# Enumerated bounds handling when a dot traverses the region's edge.

bh = kwargs.get("bound_hand", "stay")

if bh == "stay":

self.b_handling = 0 # Don't move if directed past the edge.

elif bh == "bounce":

self.b_handling = 1 # Bounce off edge (reflect coordinates).

elif bh == "trans":

self.b_handling = 2 # Translate to opposite side of the region.

else:

assert False, "Unsupported bounds handling"

# Enumerated fitness (reward) function.

ff = kwargs.get("fit_func", "euc")

if ff == "euc":

self.fit_func = 0 # Single Euclidean (Pythagorean) distance value

elif ff == "disp":

self.fit_func = 1 # Tuple of x,y displacement values

elif ff == "rng":

self.fit_func = 2 # Range rings--the closer the ring, the lower the number

elif ff == "dir":

self.fit_func = 3 # direction--moving closer or farther away?

else:

assert False, "Unsupported fitness function"

self.ring_size = kwargs.get("ring_size", 2) # Range ring size

self.bullseye = kwargs.get("bullseye", 10.0) # Intercept reward

self.teleport = kwargs.get("teleport", True) # Teleport after intercept

# Initialize empty lists of relevant and distraction dots.

self.netDot = Dot(0, 0, self.decay)

self.dots = []

self.herrings = []

self.obs = np.zeros((self.h, self.w))

self.action_space = spaces.Discrete(self.choices)

self.newPlot = True # One-time flag

def step(self, action):

"""

Generates a grid for the current timestep.

See above for full description.

:param action: network's prediction of the dot movement

:return obs: observation of grid matrix of shape (h,w)

:return reward: precision of network's prediction in Euclidean distance

:return done: indicates termination of simulation

:return intercept: indicates a successful intercept this step

"""

# Increment timestep

self.ts += 1 # Increment timestep

# If the random rate is high enough, update movement direction.

if random.uniform(0, 1) <= self.randr:

self.dotDir = random.randint(

self.minch, self.choices - 1

) # five possible options

# Initialize empty grid and populate as we update dots.

self.obs = np.zeros((self.h, self.w))

# Update network dot according to the network's action.

self.prevRow = self.netDot.row[0]

self.prevCol = self.netDot.col[0]

if action is not None:

self.movePoint(self.netDot, action)

# self.obs = self.obs/(self.ndots + self.herrs) # normalize

reward, intercept = self.compute_reward()

# Teleport network dot if intercept is successful.

if intercept and self.teleport:

bh1, bh2 = self.h // 5, 4 * self.h // 5

bw1, bw2 = self.w // 5, 4 * self.w // 5

r, c = random.randint(bh1, bh2), random.randint(bw1, bw2)

self.netDot = Dot(r, c, self.decay)

# Redo grid observation.

self.obs = np.zeros((self.h, self.w))

self.obs[r, c] = 1

# Move all relevant dots in the same direction.

for d in self.dots:

self.movePoint(d)

# Move distraction dots with individually randomized motion.

for h in self.herrings:

self.movePoint(h, random.randint(self.minch, self.choices - 1))

reward = torch.Tensor(np.array(reward))

done = self.timesteps <= self.ts

return self.obs, reward, done, intercept

def reset(self):

"""

Reset dots to initial positions, and reset RNG seed.

"""

self.ts = 0 # reset timesteps

# Reset RNG

random.seed(self.seed)

# provide default observation

self.obs = np.zeros((self.h, self.w))

# Set boundaries (so we don't spawn points on the edge.

# Not that there's a real problem with it, but it's boring.

bh1, bh2 = self.h // 5, 4 * self.h // 5

bw1, bw2 = self.w // 5, 4 * self.w // 5

# Start dots in the middle.

# midr = self.h//2 <= we randomize this now.

# midc = self.w//2

# We know that the sum from n=0 to n=N of 1 - n/N = (N + 1)/2

# Thus, computing the initial grid space for a dot, given the

# length of its tail N would be (self.decay + 1)/2.

# But... we also have to cap it at 1. So, who cares?

# Reinitalize network dot placement.

r, c = random.randint(bh1, bh2), random.randint(bw1, bw2)

self.netDot = Dot(r, c, self.decay)

self.obs[r, c] = 1

# Reinitalize target dot placement with initial movement direction.

self.dots = []

self.dotDir = random.randint(self.minch, self.choices - 1)

for d in range(self.ndots):

r, c = random.randint(bh1, bh2), random.randint(bw1, bw2)

self.dots.append(Dot(r, c, self.decay))

self.obs[r, c] = 1

# Reinitalize red herring placement.

self.herrings = []

for h in range(self.herrs):

r, c = random.randint(bh1, bh2), random.randint(bw1, bw2)

self.herrings.append(Dot(r, c, self.decay))

self.obs[r, c] = 1

return self.obs

def movePoint(self, d: Dot, dotDir: int = -1):

"""

Apply clockwise directional enumeration.

:param dotDir: enumerated movement as described above.

:param/return r: current row => next row

:param/return c: current column => next column

"""

# If not provided, use the known current direction.

targetDir = False # flag if we're using the target's direction.

if dotDir < 0:

dotDir = self.dotDir

targetDir = True

r, c = d.row[0], d.col[0]

""" Apply clockwise directional enumeration. """

# 0 means stay, though we also won't go past the edge.

if dotDir == 1: # up

r += self.speed

elif dotDir == 2: # right

c += self.speed

elif dotDir == 3: # down

r -= self.speed

elif dotDir == 4: # left

c -= self.speed

elif dotDir == 5: # up and right

r += self.speed

c += self.speed

elif dotDir == 6: # down and right

r -= self.speed

c += self.speed

elif dotDir == 7: # down and left

r -= self.speed

c -= self.speed

elif dotDir == 8: # up and left

r += self.speed

c -= self.speed

elif dotDir != 0: # Woops

assert False, "Unsupported dot direction"

""" When a dot attempts to move past an edge... """

# Stay put.

if self.b_handling == 0:

r = max(min(r, self.h - 1), 0)

c = max(min(c, self.w - 1), 0)

# direction stays the same.

# Bounce: reflect its coordinates back into the region.

elif self.b_handling == 1:

if r < 0 or self.h <= r:

r = self.h - 1 - r % self.h # reflect row

if targetDir:

self.dotDir += ROW_CROSSING[dotDir]

if c < 0 or self.w <= c:

c = self.w - 1 - c % self.w # reflect column

if targetDir:

self.dotDir += COL_CROSSING[dotDir]

# Translate: the dot will continue in the same direction

# from the opposite side of the region.

elif self.b_handling == 2:

r = r % self.h # Mirror row

c = c % self.w # Mirror column

# direction stays the same.

# Woops

else:

assert False, "Unsupported bounds handling"

# Update the saved point in the Dot class.

# This also cycles the path history.

d.move(r, c)

# Update the grid with this point and its decaying trail.

for t in range(self.decay):

self.obs[d.row[t], d.col[t]] = min(

self.obs[d.row[t], d.col[t]] + 1 - t / self.decay, 1

)

def compute_reward(self):

"""

Computes reward according to the chosen fitness function.

Returns reward and flag indicating a successful intercept.

"""

# Add bull's eye reward (if we're using it)

if (

self.bullseye != 0

and self.dots[0].row[0] == self.netDot.row[0]

and self.dots[0].col[0] == self.netDot.col[0]

):

return self.bullseye, True

reward = 0.0

# Euclidean distance

if self.fit_func == 0:

reward = -np.hypot(

self.dots[0].row[0] - self.netDot.row[0],

self.dots[0].col[0] - self.netDot.col[0],

)

# Displacement tensor

elif self.fit_func == 1:

reward = torch.Tensor(

[

self.dots[0].row[0] - self.netDot.row[0],

self.dots[0].col[0] - self.netDot.col[0],

]

)

# Range rings; default range ring size = 2

elif self.fit_func == 2:

reward = (

-np.hypot(

self.dots[0].row[0] - self.netDot.row[0],

self.dots[0].col[0] - self.netDot.col[0],

)

// self.ring_size

)

# Directional

elif self.fit_func == 3:

rd1 = abs(self.dots[0].row[0] - self.prevRow)

rd2 = abs(self.dots[0].row[0] - self.netDot.row[0])

cd1 = abs(self.dots[0].col[0] - self.prevCol)

cd2 = abs(self.dots[0].col[0] - self.netDot.col[0])

if rd2 < rd1:

reward += 1.0 # right row movement

elif rd1 < rd2:

reward -= 1.0 # wrong row movement

if cd2 < cd1:

reward += 1.0 # right col movement

elif cd1 < cd2:

reward -= 1.0 # wrong col movement

# Woops

else:

assert False, "Unsupported fitness function"

return reward, False

def render(self):

"""

Display current state, either in ASCII or graphic plots.

"""

# Double value of network dot only for visual aid in rendering.

temp = self.obs

temp[self.netDot.row, self.netDot.col] *= 2

# Write to file if requested.

if self.write2F:

with open(self.filename, "ab") as f:

np.savetxt(f, temp, delimiter=",")

# Print as pandas dataframe if requested.

if self.pandas:

print("Timestep:", self.ts)

print(pd.DataFrame(temp, dtype="uint32"))

# Provide graphical rendering if requested.

if not self.mute:

# Otherwise, we'll render it as... I don't know yet.

# print('Timestep:', self.ts)

# get current figure, clear it, and replot.

if self.newPlot:

self.fig = plt.gcf()

plt.figure(self.fig.number)

plt.clf()

plt.imshow(temp, cmap="hot", interpolation="nearest")

# Only display colorbar once.

if self.newPlot:

self.newPlot = False

# plt.ion()

plt.colorbar()

# Pause so that that GUI can do its thing.

plt.pause(1e-8)

def cycleOutFiles(self, newInt=-1):

"""

Increments numbered suffix on output file to start a new one.

"""

underScore = self.filename.rfind("_")

if 0 <= newInt:

self.fileCnt = newInt

else:

self.fileCnt += 1

self.filename = self.filename[: underScore + 1] + str(self.fileCnt) + ".csv"

def addFileSuffix(self, suffix):

"""

Adds suffix to output file (like "train" or "test").

"""

self.filename = self.filename[:-5] + suffix + "_" + self.filename[-5:]

def changeFileSuffix(self, sFrom, sTo):

"""

Adds suffix to output file (like "train" or "test").

"""

self.filename = self.filename.replace(sFrom, sTo)

self.cycleOutFiles(newInt=0) # reset file count.

def driver():

steps = 200

dotSim = DotSimulator(200)

dotSim.reset()

grids = np.zeros((steps, 28, 28))

directions = np.zeros(steps)

done = False

for t in range(steps):

grids[t], reward, done, info = dotSim.step(0)

directions[t] = info["direction"]

vals, cnts = np.unique(directions, return_counts=True)

print(vals, cnts / steps)

if __name__ == "__main__":

driver()