light_maymays/guitarhero/main.py

import math

import pygame
import serial

def clamp(x, ab):
    (a, b) = ab
    return max(a, min(b, x))

def rescale(x, from_limits, to_limits):
    (a, b) = from_limits
    x_0_1 = (x - a) / (b - a)

    (c, d) = to_limits
    return c + (d - c) * x_0_1

class MovingHead:
    def __init__(self, start_addr):
        self.start_addr = start_addr

        self.pan = 0  # -pi/2 to pi/2
        self.tilt = 0  # -3pi/2 to 3pi/2
        self.speed = 0
        self.dimmer = 0  # 0 to 1
        self.rgbw = (0, 0, 0, 0)

    def __str__(self):
        return f"MovingHead({self.start_addr}): pan={self.pan!r}, tilt={self.tilt!r}, speed={self.speed!r}, dimmer={self.dimmer!r}, rgbw={self.rgbw!r}"
    
    def render(self, dst):
        
        pan = rescale(self.pan, (-1.5 * math.pi, 1.5 * math.pi), (255, 0))
        pan = clamp(int(pan), (0, 255))
        pan_fine = 0

        tilt = rescale(self.tilt, (-0.5 * math.pi, 0.5 * math.pi), (0, 255))
        tilt = clamp(int(tilt), (0, 255))
        tilt_fine = 0

        dimmer = clamp(7 + int(127 * self.dimmer), (7, 134))

        (r, g, b, w) = self.rgbw

        channels = [
            pan,
            pan_fine,
            tilt,
            tilt_fine,
            self.speed,
            dimmer,
            r,
            g,
            b,
            w,
            0,  # color mode
            0,  # auto jump speed
            0,  # control mode
            0,  # reset
        ]

        offset = self.start_addr - 1

        dst[offset : offset + len(channels)] = channels

if __name__ == "__main__":

    movingheads = [
        MovingHead(1),
        MovingHead(15),
        MovingHead(29),
        MovingHead(43),
    ]

    movingheads[0].rgbw = ( 59, 130, 246, 0)  # blue
    movingheads[1].rgbw = (245, 158,  11, 0)  # yellow
    movingheads[2].rgbw = (239,  68,  68, 0)  # red
    movingheads[3].rgbw = ( 16, 185, 129, 0)  # green

    for head in movingheads:
        head.pan = -0.5 * math.pi

    dmx_data = bytearray(512)
    
    # pygame initialization

    pygame.init()

    screen = pygame.display.set_mode((420, 69))
    pygame.display.set_caption("meh")

    pygame.joystick.init()
    gamepad = pygame.joystick.Joystick(0)
    gamepad.init()

    clock = pygame.time.Clock()

    # main loop

    running = True

    with serial.Serial("/dev/ttyUSB0", 500_000) as ser:

        def sync():
            # wait for sync
            while True:
                b = ser.readline()
                if b.strip() == b"Sync.":
                    return
                
        
        sync()

        while running:
            
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                    
            # axes:
            pitch = gamepad.get_axis(4)  # 0 horizontal, -1 up, 1 down
            whammy = gamepad.get_axis(3)  # -1 default, 1 if max. pressed
            
            # buttons:
            button_state = [
                gamepad.get_button(i)
                for i in [
                    0,  # green
                    1,  # red
                    3,  # yellow
                    2,  # blue
                    4,  # orange
                ]
            ]
            
            # strumm bar
            (_, bar) = gamepad.get_hat(0)

            print(f"pitch: {pitch:+0.2f}, whammy: {whammy:+0.2f}, buttons: {button_state}, bar: {bar: 1}")
            
            # render

            for head_id, button_id in enumerate([3, 2, 1, 0]):
                movingheads[head_id].dimmer = 1 if button_state[button_id] else 0.2

            pitch = max(0, min(1, -1 * pitch))  # 0 horizontal, 1 up
            tilt = rescale(pitch, (0, 1), (-0.5 * math.pi, 0))
            for head in movingheads:
                head.tilt = tilt
                
            whammy_0_1 = (whammy + 1) / 2
            
            ANGLE_HOME  = -0.5 * math.pi
            ANGLE_OUTER = 0.5235988
            ANGLE_INNER = 0.2617994
            
            movingheads[0].pan = rescale(whammy_0_1, (0, 1), (ANGLE_HOME, ANGLE_HOME + ANGLE_OUTER))
            movingheads[1].pan = rescale(whammy_0_1, (0, 1), (ANGLE_HOME, ANGLE_HOME + ANGLE_INNER))
            movingheads[2].pan = rescale(whammy_0_1, (0, 1), (ANGLE_HOME, ANGLE_HOME - ANGLE_INNER))
            movingheads[3].pan = rescale(whammy_0_1, (0, 1), (ANGLE_HOME, ANGLE_HOME - ANGLE_OUTER))

            for head in movingheads:
                # print(head)
                head.render(dmx_data)
            
            ser.write(dmx_data)
            ser.flush()
            response = ser.readline()
            if response.strip() != b"Ack.":
                print(f"received bad response: {response!r}")
                sync()

            try:
                clock.tick(50)
            except KeyboardInterrupt:
                running = False
    
    pygame.quit()