Add function that returns whether a position is lit.
[tabakrolletjie.git] / tabakrolletjie / lights.py
index 92b66e0a8d4b91907fdec632724e00b0ffa9fba8..3c9795212fdcc00b8e4b1b782c7de62991f1f593 100644 (file)
 """
 
 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)
+
+
+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 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)
 
 
 class BaseLight(object):
     """ Common light functionality. """
 
+    COLOURS = {
+        "red": (255, 0, 0),
+        "green": (0, 255, 0),
+        "blue": (0, 255, 255),
+        "yellow": (255, 255, 0),
+        "white": (255, 255, 255),
+    }
+
     def __init__(self, colour, position):
-        self.body = pymunk.Body(0, 0, pymunk.body.Body.STATIC)
         self.colour = colour
-        self.position = position
-
-    def add_to_space(self, space):
-        space.remove(self.body, **self.body.shapes())
-        shapes = self.determine_ray_polys(space)
-        space.add(self.body, **shapes)
-
-    def determine_ray_polys(self, space):
-        raise NotImplementedError(
-            "Lights should implement .determine_ray_polys.")
+        self.position = pymunk.Vec2d(position)
+        self.on = True
+        self.body = pymunk.Body(0, 0, pymunk.body.Body.STATIC)
+        self.fitting = pymunk.Circle(self.body, 10.0, self.position)
+        self.body.light = self
 
     @classmethod
     def load(cls, config):
@@ -30,6 +135,46 @@ class BaseLight(object):
             if c.__name__.lower() == light_type]
         return light_class(**kw)
 
+    def add(self, space):
+        if self.body.space is not None:
+            space.remove(self.body, *self.body.shapes)
+        shapes = self.shapes_for_ray_polys(
+            calculate_ray_polys(space, self.body, self.position))
+        for shape in shapes:
+            shape.filter = LIGHT_FILTER
+        self.fitting.filter = FITTINGS_FILTER
+        space.add(self.body, self.fitting, *shapes)
+
+    def shapes_for_ray_polys(self, ray_polys):
+        return ray_polys
+
+    def toggle(self):
+        self.on = not self.on
+
+    def render_light(self, surface):
+        if not self.on:
+            return
+        subsurface = surface.copy()
+        light_colour = self.COLOURS[self.colour]
+        for shape in self.body.shapes:
+            if shape is self.fitting:
+                continue
+            pygame_poly = [
+                pymunk.pygame_util.to_pygame(v, surface) for v in
+                shape.get_vertices()]
+            pygame.draw.polygon(
+                subsurface, light_colour, pygame_poly, 0)
+            pygame.draw.aalines(
+                subsurface, light_colour, True, pygame_poly, 1)
+        subsurface.set_alpha(50)
+        surface.blit(subsurface, (0, 0), None)
+
+    def render_fitting(self, surface):
+        pygame.draw.circle(
+            surface, (255, 255, 0),
+            pymunk.pygame_util.to_pygame(self.fitting.offset, surface),
+            int(self.fitting.radius))
+
 
 class SpotLight(BaseLight):
     def __init__(
@@ -37,6 +182,7 @@ class SpotLight(BaseLight):
         super(SpotLight, self).__init__(colour, position)
         self.direction = direction
         self.spread = spread
+        self.i = 0
 
 
 class Lamp(BaseLight):