+import pymunk
+import pymunk.pygame_util
+import pygame.display
+import pygame.draw
+
+from .constants import (
+ SCREEN_SIZE, LIGHT_CATEGORY, FITTINGS_CATEGORY)
+from .utils import debug_timer
+
+LIGHT_FILTER = pymunk.ShapeFilter(
+ mask=pymunk.ShapeFilter.ALL_MASKS ^ (
+ LIGHT_CATEGORY | FITTINGS_CATEGORY),
+ categories=LIGHT_CATEGORY)
+
+FITTINGS_FILTER = pymunk.ShapeFilter(
+ mask=pymunk.ShapeFilter.ALL_MASKS ^ (
+ LIGHT_CATEGORY | FITTINGS_CATEGORY),
+ categories=FITTINGS_CATEGORY)
+
+# Just match lights, nothing else
+LIT_BY_FILTER = pymunk.ShapeFilter(mask=LIGHT_CATEGORY)
+
+
+def screen_rays(pos):
+ """ An iterable that returns ordered rays from pos to the edge of the
+ screen, starting with the edge point (0, 0) and continuing clockwise
+ in pymunk coordinates.
+ """
+ w, h = SCREEN_SIZE
+ left, right, bottom, top = 0, w, 0, h
+ step = 1
+ for y in range(0, h, step):
+ yield pymunk.Vec2d(left, y)
+ for x in range(0, w, step):
+ yield pymunk.Vec2d(x, top)
+ for y in range(top, -1, -step):
+ yield pymunk.Vec2d(right, y)
+ for x in range(right, -1, -step):
+ yield pymunk.Vec2d(x, bottom)
+
+
+@debug_timer("lights.calculate_ray_polys")
+def calculate_ray_polys(space, body, position):
+ position = pymunk.Vec2d(position)
+ vertices = [position]
+ ray_polys = []
+ for ray in screen_rays(position):
+ info = space.segment_query_first(position, ray, 1, LIGHT_FILTER)
+ point = ray if info is None else info.point
+ vertices.append(point)
+ if len(vertices) > 3:
+ trial_poly = pymunk.Poly(None, vertices)
+ trial_poly.update(pymunk.Transform.identity())
+ query_prev = trial_poly.point_query(vertices[-2])
+ query_pos = trial_poly.point_query(position)
+ if query_prev.distance < -0.01 or query_pos.distance < -0.01:
+ new_poly = pymunk.Poly(body, vertices[:-1])
+ vertices = [position, vertices[-1]]
+ ray_polys.append(new_poly)
+ else:
+ vertices = trial_poly.get_vertices() + [point]
+ if len(vertices) > 2:
+ ray_polys.append(pymunk.Poly(body, vertices))
+ return ray_polys
+
+
+class LightManager(object):
+ """ Manages a set of lights. """
+
+ def __init__(self, space, gamestate):
+ self._space = space
+ self._lights = [
+ BaseLight.load(cfg) for cfg in gamestate.station["lights"]]
+ for light in self._lights:
+ light.add(self._space)
+
+ def toggle_nearest(self, *args, **kw):
+ light = self.nearest(*args, **kw)
+ if light:
+ light.toggle()
+
+ def nearest(self, pos, surfpos=False, max_distance=1.0):
+ if surfpos:
+ surface = pygame.display.get_surface()
+ pos = pymunk.pygame_util.from_pygame(pos, surface)
+ point_info = self._space.point_query_nearest(
+ pos, max_distance, pymunk.ShapeFilter(mask=FITTINGS_CATEGORY))
+ if point_info is not None:
+ return point_info.shape.body.light
+ return None
+
+ def lit_by(self, pos, surfpos=False, max_distance=0.0):
+ if surfpos:
+ surface = pygame.display.get_surface()
+ pos = pymunk.pygame_util.from_pygame(pos, surface)
+ point_info_list = self._space.point_query(
+ pos, max_distance, pymunk.ShapeFilter(mask=LIGHT_CATEGORY))
+ lights = [p.shape.body.light for p in point_info_list]
+ return [light for light in lights if light.on]
+
+ def light_query(self, shape):
+ """Query the lights by shape"""
+ old_filter = shape.filter
+ # We need to restrict matches to only the lights
+ shape.filter = LIT_BY_FILTER
+ shape_info_list = self._space.shape_query(shape)
+ shape.filter = old_filter
+ lights = [p.shape.body.light for p in shape_info_list]
+ return [light for light in lights if light.on]
+
+ def render_light(self, surface):
+ for light in self._lights:
+ light.render_light(surface)
+
+ def render_fittings(self, surface):
+ for light in self._lights:
+ light.render_fitting(surface)
+