lighting.py 3.06 KB
Newer Older
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
1 2
#!/usr/bin/env python3

3
from time import sleep
4
from random import randint
5
import configparser
6
from threading import Thread
7

8 9 10 11
import zmq
import spidev
import ws2812

12 13 14 15 16 17

class Color:
    green: int
    red: int
    blue: int

18 19 20 21 22
    def __init__(self, green, red, blue):
        self.green = green
        self.red = red
        self.blue = blue

23 24 25 26 27 28
    def as_tuple(self):
        return self.green, self.red, self.blue


class ColorManager(Thread):
    current_rand = 1
29
    unscaled_blue = 100
30 31 32
    current_color = Color(0, 0, 0)
    activeFlag = True
    spi = spidev.SpiDev()
33
    running = True
34 35 36 37

    def __init__(self):
        super().__init__()
        self.spi.open(1, 0)
38 39 40 41
        self.lights_off()

    def lights_off(self):
        ws2812.write2812(self.spi, [(0, 0, 0) * 4])
42 43 44 45

    def get_color(self):
        if self.current_color.green + self.current_rand > 80:
            self.current_rand = -randint(1, 8)
46 47
            self.current_color = Color(80, 255, self.current_color.blue)
        elif self.current_color.green + self.current_rand < 0:
48
            self.current_rand = randint(1, 8)
49
            self.current_color = Color(0, 255, self.current_color.blue)
Federico Meloda's avatar
Federico Meloda committed
50
        else:
51
            self.current_color.green += self.current_rand
52
            self.current_color.red = 255
53
            self.current_color.blue = self.unscaled_blue * self.current_color.green//255
54
        if self.unscaled_blue < 100 :
55
            self.unscaled_blue += 1
56

57 58 59 60 61
    def fade(self):
        if self.current_color.red > 0:
            self.current_color.red -= 1
        if self.current_color.green > 0:
            self.current_color.green -= 1
62

63 64 65
    def stop(self):
        self.running = False

66 67
    def start_fade(self):
        self.activeFlag = False
68

69
    def burn(self):
70
        self.unscaled_blue = 0
71

72
    def run(self):
73
        while self.running:
74 75 76 77
            if self.activeFlag:
                self.get_color()
            else:
                self.fade()
Federico Meloda's avatar
Federico Meloda committed
78

79 80
            ws2812.write2812(self.spi, [self.current_color.as_tuple() * 4])
            sleep(0.001)
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
81 82


83
def main():
84 85 86
    config = configparser.ConfigParser({'host': '10.10.10.1', 'port': '5556'})
    config.read('config.cfg')

Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
87 88
    port = config["DEFAULT"]['port']
    host = config["DEFAULT"]['host']
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
89

90 91 92
    # Socket to talk to server
    context = zmq.Context()
    socket = context.socket(zmq.SUB)
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
93

94
    print("Connecting to coordinator...")
95
    socket.connect(f"tcp://{host}:{port}")
96

97
    topicfilter = "events"
98
    socket.setsockopt_string(zmq.SUBSCRIBE, topicfilter)
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
99

100 101
    cm = ColorManager()

102 103 104 105
    while True:
        string = socket.recv()
        topic, messagedata = string.split()
        print(topic, messagedata)
106

107
        if messagedata == b"liftoff":
108 109
            cm = ColorManager()
            cm.start()
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
110
            cm.burn()
111
        elif messagedata == b"stage_1_separation":
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
112
            cm.burn()
Federico Meloda's avatar
Federico Meloda committed
113
        elif messagedata == b"stage_2_separation":
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
114
            cm.burn()
115
            cm.start_fade()
116 117
        elif messagedata == b"abort":
            cm.lights_off()
118
            cm.stop()
119 120
        else:
            print("Unhandled message")
121 122 123 124


if __name__ == "__main__":
    main()