(norganizer) web status filter
[experiments.git] / evolution / scene.py
1
2 import logging
3 from pygame import Rect
4
5 class QuadTree(object):
6 """An implementation of a quad-tree.
7
8 This QuadTree started life as a version of [1] but found a life of its own
9 when I realised it wasn't doing what I needed. It is intended for static
10 geometry, ie, items such as the landscape that don't move.
11
12 This implementation inserts items at the current level if they overlap all
13 4 sub-quadrants, otherwise it inserts them recursively into the one or two
14 sub-quadrants that they overlap.
15
16 Items being stored in the tree must be a pygame.Rect or have have a
17 .rect (pygame.Rect) attribute that is a pygame.Rect
18 ...and they must be hashable.
19
20 Acknowledgements:
21 [1] http://mu.arete.cc/pcr/syntax/quadtree/1/quadtree.py
22 """
23
24 class QuadTreeNode:
25 def __init__(self, item):
26 self.item = item
27 self.rect = Rect(item.pos[0], item.pos[1], 1, 1)
28
29 def __init__(self, items, depth=8, bounding_rect=None):
30 """Creates a quad-tree.
31
32 @param items:
33 A sequence of items to store in the quad-tree. Note that these
34 items must be a pygame.Rect or have a .rect attribute.
35
36 @param depth:
37 The maximum recursion depth.
38
39 @param bounding_rect:
40 The bounding rectangle of all of the items in the quad-tree. For
41 internal use only.
42 """
43
44 # The sub-quadrants are empty to start with.
45 self.nw = self.ne = self.se = self.sw = None
46
47 # If we've reached the maximum depth then insert all items into this
48 # quadrant.
49 depth -= 1
50 if depth == 0 or not items:
51 self.items = items
52 return
53
54 # Find this quadrant's centre.
55 if bounding_rect:
56 bounding_rect = Rect( bounding_rect )
57 else:
58 # If there isn't a bounding rect, then calculate it from the items.
59 bounding_rect = Rect( items[0] )
60 for item in items[1:]:
61 bounding_rect.union_ip( item )
62 cx = self.cx = bounding_rect.centerx
63 cy = self.cy = bounding_rect.centery
64
65 self.items = []
66 nw_items = []
67 ne_items = []
68 se_items = []
69 sw_items = []
70
71 for item in items:
72 # Which of the sub-quadrants does the item overlap?
73 in_nw = item.rect.left <= cx and item.rect.top <= cy
74 in_sw = item.rect.left <= cx and item.rect.bottom >= cy
75 in_ne = item.rect.right >= cx and item.rect.top <= cy
76 in_se = item.rect.right >= cx and item.rect.bottom >= cy
77
78 # If it overlaps all 4 quadrants then insert it at the current
79 # depth, otherwise append it to a list to be inserted under every
80 # quadrant that it overlaps.
81 if in_nw and in_ne and in_se and in_sw:
82 self.items.append(item)
83 else:
84 if in_nw: nw_items.append(item)
85 if in_ne: ne_items.append(item)
86 if in_se: se_items.append(item)
87 if in_sw: sw_items.append(item)
88
89 # Create the sub-quadrants, recursively.
90 if nw_items:
91 self.nw = QuadTree(nw_items, depth, (bounding_rect.left, bounding_rect.top, cx, cy))
92 if ne_items:
93 self.ne = QuadTree(ne_items, depth, (cx, bounding_rect.top, bounding_rect.right, cy))
94 if se_items:
95 self.se = QuadTree(se_items, depth, (cx, cy, bounding_rect.right, bounding_rect.bottom))
96 if sw_items:
97 self.sw = QuadTree(sw_items, depth, (bounding_rect.left, cy, cx, bounding_rect.bottom))
98
99
100 def hit(self, x, y, w, h):
101 """Returns the items that overlap a bounding rectangle.
102 Returns the set of all items in the quad-tree that overlap with a
103 bounding rectangle.
104 """
105
106 rect = Rect(x, y, w, h)
107
108 # Find the hits at the current level.
109 hits = set( [ self.items[n] for n in rect.collidelistall( self.items ) ] )
110
111 # Recursively check the lower quadrants.
112 if self.nw and rect.left <= self.cx and rect.top <= self.cy:
113 hits |= self.nw.hit(x, y, w, h)
114 if self.sw and rect.left <= self.cx and rect.bottom >= self.cy:
115 hits |= self.sw.hit(x, y, w, h)
116 if self.ne and rect.right >= self.cx and rect.top <= self.cy:
117 hits |= self.ne.hit(x, y, w, h)
118 if self.se and rect.right >= self.cx and rect.bottom >= self.cy:
119 hits |= self.se.hit(x, y, w, h)
120
121 return hits
122
123
124 class Scene:
125 class Types:
126 DRAWABLE = "Drawable"
127 PHYSICAL = "Physical"
128 ALIVE = "Alive"
129 GOODS = "Goods"
130
131 class Render:
132 Good = 0
133 Bot = 1
134
135 class GoodTypes:
136 FOOD = "Food"
137
138 def __init__(self, name, w, h):
139 self.name = name
140 self.w = w
141 self.h = h
142 self.all = set()
143 self.byType = {}
144 self.byEntity = {}
145 self.logger = logging.getLogger(self.name)
146
147 def add(self, entity):
148 self.logger.info("Scene: add %s", entity)
149 self.all.add(entity)
150 for type in entity.TYPES:
151 self.byType.setdefault(type, []).append(entity)
152
153 def remove(self, entity):
154 self.logger.info("Scene: remove %s", entity)
155 self.all.remove(entity)
156 for type in entity.TYPES:
157 entities = self.byType.get(type, [])
158 if entity in entities: entities.remove(entity)
159
160 def getByType(self, type):
161 return self.byType.get(type, [])
162
163 def getAt(self, entity, type):
164 return [e for e in self.getByType(type) if (e.pos[0] == entity.pos[0] and e.pos[1] == entity.pos[1])]
165
166 def getAround(self, entity, distance, type):
167 hits = self.tree.hit(
168 entity.pos[0] - distance / 2,
169 entity.pos[1] - distance / 2,
170 distance, distance
171 )
172 results = [hit.item for hit in hits if type in hit.item.TYPES]
173 return results
174
175 def spatialHash(self):
176 bodies = [QuadTree.QuadTreeNode(e) for e in self.all if Scene.Types.PHYSICAL in e.TYPES]
177 self.tree = QuadTree(bodies)
178
179
180