-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathreplay_handler.py
108 lines (100 loc) · 4.36 KB
/
replay_handler.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation as animation
import matplotlib.text as text
import time
import random
import pickle
rad_labels = 'abcdefghijklmnopqrstuvwxyz' #for replay annotations
dir_labels = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' #for replay annotations
class ReplayHandler:
# Constructor
def __init__(self,
rad_creeps_history,
dir_creeps_history,
start_index,
end_index,
):
self.rad_creeps_history = rad_creeps_history
self.rad_creeps_history.sort(key=lambda c: c.birth_count)
self.dir_creeps_history = dir_creeps_history
self.dir_creeps_history.sort(key=lambda c: c.birth_count)
self.index = start_index
self.end_index = end_index
self.fig, self.ax = plt.subplots()
self.sc_r = self.ax.scatter([],[],c=(0.0,1.0,0.0),s=200,edgecolors='face')
self.sc_d = self.ax.scatter([],[],c=(0.65,0.4,0.3),s=200,edgecolors='face')
self.annotations = []
# set limits to screen width/height
self.ax.set_xlim(0,1920)
self.ax.set_ylim(1080,0)
# Print the creeps
def update(self,_):
xs_r, ys_r = ([],[])
xs_d, ys_d = ([],[])
self.clear_annotations()
if self.index > self.end_index:
self.index = self.end_index # keep animating last step
for creep_id, creep in enumerate(self.rad_creeps_history):
was_born = (self.index >= creep.birth_count)
has_died = (self.index > len(creep.health_history)+creep.birth_count)
if was_born and not has_died:
i = self.index-creep.birth_count
if i == len(creep.x_history):
x_coord = creep.x
y_coord = creep.y
health = creep.health
else:
x_coord = creep.x_history[i]
y_coord = creep.y_history[i]
health = creep.health_history[i]
xs_r.append(x_coord)
ys_r.append(y_coord)
creep_label = rad_labels[creep_id%len(rad_labels)]
annotation = text.Annotation(
"%.2f %s" % (health, creep_label),
xy=(x_coord, y_coord),
xytext=(-35,-5),
textcoords = 'offset points')
self.add_annotation(annotation)
for creep_id, creep in enumerate(self.dir_creeps_history):
was_born = (self.index >= creep.birth_count)
has_died = (self.index > len(creep.health_history)+creep.birth_count)
if was_born and not has_died:
i = self.index-creep.birth_count
if i == len(creep.x_history):
x_coord = creep.x
y_coord = creep.y
health = creep.health
else:
x_coord = creep.x_history[i]
y_coord = creep.y_history[i]
health = creep.health_history[i]
xs_d.append(x_coord)
ys_d.append(y_coord)
creep_label = dir_labels[creep_id%len(dir_labels)]
annotation = text.Annotation(
"%s %.2f" % (creep_label, health),
xy=(x_coord, y_coord),
xytext=(10, -5),
textcoords = 'offset points')
self.add_annotation(annotation)
self.ax.set_title("Frame %i" % self.index)
self.sc_r.set_offsets(np.column_stack((xs_r, ys_r)))
self.sc_d.set_offsets(np.column_stack((xs_d, ys_d)))
self.index += 1
return self.sc_r, #self.sc_d
def add_annotation(self, annotation):
self.annotations.append(annotation)
self.ax.add_artist(annotation)
def clear_annotations(self):
for annotation in self.annotations:
annotation.remove()
self.annotations = []
def replay_main(rad_creeps_history,dir_creeps_history, start_index=0, end_index=-1):
R = ReplayHandler(rad_creeps_history, dir_creeps_history, start_index, end_index)
ani = animation.FuncAnimation(R.fig, R.update, interval=50,blit=False)
plt.show()
if __name__ == '__main__':
state = pickle.load( open( "save.p", "rb" ) )
replay_main(state['rad_creeps_history'], state['dire_creeps_history'], 0, state['counter'])