lighting.py 2.84 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 33 34 35 36
    current_color = Color(0, 0, 0)
    activeFlag = True
    spi = spidev.SpiDev()

    def __init__(self):
        super().__init__()
        self.spi.open(1, 0)
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
37
        ws2812.write2812(self.spi, [(0,0,0) * 4])
38 39 40 41

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

53 54 55 56 57
    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
58

59 60
    def start_fade(self):
        self.activeFlag = False
61

62
    def burn(self):
63
        self.unscaled_blue = 0
64

65 66 67 68 69 70
    def run(self):
        while self.is_alive():
            if self.activeFlag:
                self.get_color()
            else:
                self.fade()
Federico Meloda's avatar
Federico Meloda committed
71

72 73
            ws2812.write2812(self.spi, [self.current_color.as_tuple() * 4])
            sleep(0.001)
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
74 75


76
def main():
77 78 79
    config = configparser.ConfigParser({'host': '10.10.10.1', 'port': '5556'})
    config.read('config.cfg')

Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
80 81
    port = config["DEFAULT"]['port']
    host = config["DEFAULT"]['host']
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
82

83 84 85
    # Socket to talk to server
    context = zmq.Context()
    socket = context.socket(zmq.SUB)
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
86

87
    print("Connecting to coordinator...")
88
    socket.connect(f"tcp://{host}:{port}")
89

90
    topicfilter = "events"
91
    socket.setsockopt_string(zmq.SUBSCRIBE, topicfilter)
Aljaž Srebrnič's avatar
Aljaž Srebrnič committed
92

93 94
    cm = ColorManager()

95 96 97 98
    while True:
        string = socket.recv()
        topic, messagedata = string.split()
        print(topic, messagedata)
99

100
        if messagedata == b"liftoff":
101 102
            cm = ColorManager()
            cm.start()
103
            burn()
104
        elif messagedata == b"stage_1_separation":
105
            burn()
Federico Meloda's avatar
Federico Meloda committed
106
        elif messagedata == b"stage_2_separation":
107
            burn()
108
            cm.start_fade()
109 110
        else:
            print("Unhandled message")
111 112 113 114


if __name__ == "__main__":
    main()