Nerf Lazertag - Lasertag

#!/usr/bin/env python3
"""
nerf_gun_controller_nec_v2.py

Adds:
- IR_ACTIVE_HIGH option for transmitter modules that are active-low.
- "--test" mode to transmit a known packet once per second (no mouse needed).
- Increased default repeats for easier bring-up.

NEC Extended (16/16) payload:
  address16 = (team_id << 8) | player_id
  command16 = (weapon_id << 8) | damage

Hardware notes:
- For stable 38 kHz, use a hardware-capable GPIO pin:
  BCM 12 (pin 32), BCM 13 (pin 33), BCM 18 (pin 12), BCM 19 (pin 35)
- Many "IR transmitter" mini boards are under-driven unless powered from 5V.
"""

import time
import threading
import argparse
import serial
import RPi.GPIO as GPIO
import pygame
import pigpio
from evdev import InputDevice, ecodes, list_devices


# --------------------------
# MUZZLE LED (PWM FLASH)
# --------------------------
class MuzzleLed:
    """Simple PWM muzzle LED flasher using RPi.GPIO.

    Notes:
    - Use a *different* GPIO than the IR carrier pin.
    - Requires an appropriate transistor/MOSFET if driving anything beyond a small LED.
    """
    def __init__(self, gpio: int, pwm_hz: int = 1000, duty_8bit: int = 255, pulse_ms: int = 30, active_high: bool = True):
        import threading

        self.gpio = int(gpio)
        self.pwm_hz = int(pwm_hz)
        self.duty_8bit = max(0, min(255, int(duty_8bit)))
        self.pulse_ms = max(1, int(pulse_ms))
        self.active_high = bool(active_high)

        self._lock = threading.Lock()
        self._timer = None

        GPIO.setup(self.gpio, GPIO.OUT, initial=GPIO.LOW if self.active_high else GPIO.HIGH)
        self._pwm = GPIO.PWM(self.gpio, self.pwm_hz)
        self._pwm.start(0.0)  # start off

    def _set_on(self):
        # RPi.GPIO PWM is always "active-high" at the pin level; if you wired active-low, invert duty.
        if self.active_high:
            duty = (self.duty_8bit / 255.0) * 100.0
        else:
            duty = 100.0 - ((self.duty_8bit / 255.0) * 100.0)
        self._pwm.ChangeDutyCycle(max(0.0, min(100.0, duty)))

    def _set_off(self):
        if self.active_high:
            self._pwm.ChangeDutyCycle(0.0)
        else:
            self._pwm.ChangeDutyCycle(100.0)

    def pulse(self):
        """Flash the muzzle LED briefly (non-blocking)."""
        import threading

        with self._lock:
            # Cancel any pending off-timer and restart pulse window
            if self._timer is not None:
                try:
                    self._timer.cancel()
                except Exception:
                    pass
                self._timer = None

            self._set_on()
            self._timer = threading.Timer(self.pulse_ms / 1000.0, self._set_off)
            self._timer.daemon = True
            self._timer.start()

    def shutdown(self):
        with self._lock:
            if self._timer is not None:
                try:
                    self._timer.cancel()
                except Exception:
                    pass
                self._timer = None
        try:
            self._set_off()
        except Exception:
            pass
        try:
            self._pwm.stop()
        except Exception:
            pass
        try:
            GPIO.output(self.gpio, GPIO.LOW if self.active_high else GPIO.HIGH)
        except Exception:
            pass

# --------------------------
# SOUND FILES
# --------------------------
SND_AMMO_EMPTY = "/home/yomama/Music/AmmoEmpty.wav"
SND_FIRE_LOOP  = "/home/yomama/Music/FireLoop.wav"
SND_FIRE_TAIL  = "/home/yomama/Music/FireTail.wav"
SND_RELOAD     = "/home/yomama/Music/Reload.wav"


# --------------------------
# SERIAL HUD (Arduino over USB)
# --------------------------
SERIAL_PORT = "/dev/ttyACM0"
SERIAL_BAUD = 115200

# --------------------------
# MUZZLE LED (FLASH) CONFIG
# --------------------------
#
# You requested PWM on GPIO 12 for the muzzle flash LED.
# IMPORTANT: A single GPIO pin cannot drive both the 38 kHz IR carrier and an
# LED flash at the same time. To honor your request, this file assigns:
#   - GPIO 12 -> MUZZLE LED (PWM flash)
#   - GPIO 13 -> IR transmitter (hardware-capable PWM pin)
# If your wiring differs, change MUZZLE_LED_GPIO and/or IR_TX_GPIO.

MUZZLE_LED_GPIO = 12
MUZZLE_LED_ACTIVE_HIGH = True
MUZZLE_LED_PWM_HZ = 1000          # LED PWM frequency (not the IR carrier)
MUZZLE_LED_DUTY_8BIT = 255        # 0..255
MUZZLE_PULSE_MS = 30              # flash duration per shot

# --------------------------
# IR TRANSMIT CONFIG
# --------------------------
IR_TX_GPIO = 13
IR_CARRIER_HZ = 38000
IR_DUTY = 1/3

# True=active-high, False=active-low
IR_ACTIVE_HIGH = True

IR_REPEATS = 4
IR_REPEAT_GAP_MS = 12

# --------------------------
# LAZERTAG PAYLOAD
# --------------------------
TEAM_ID = 2
PLAYER_ID = 5
WEAPON_ID = 1
DAMAGE_PER_BULLET = 40

# --------------------------
# SIMPLE IR HIT MODE (bring-up)
# --------------------------
# When enabled, the gun always transmits ONE fixed NEC-Extended frame per shot.
# The vest can simply match (address16, command16) and treat it as a hit.
SIMPLE_IR_MODE = True
SIMPLE_ADDRESS16 = 0x00FF
SIMPLE_COMMAND16 = 0xA55A
SIMPLE_DAMAGE = 25  # vest-side damage to apply per valid hit

# --------------------------
# WEAPON CONFIG
# --------------------------
MAG_CAPACITY = 40
RELOAD_SECONDS = 3.0
RPM = 800
SECONDS_PER_SHOT = 60.0 / RPM
WEAPON_NAME = "SMG IR AUTO"

MOUSE_EVENT_PATH = None

# --------------------------
# AUDIO
# --------------------------
class Audio:
    def __init__(self):
        # Increase channel count so multiple sounds can overlap without blocking.
        pygame.mixer.init(frequency=44100, size=-16, channels=2, buffer=512)
        pygame.mixer.set_num_channels(8)

        self._require_file(SND_AMMO_EMPTY)
        self._require_file(SND_FIRE_LOOP)
        self._require_file(SND_FIRE_TAIL)
        self._require_file(SND_RELOAD)

        self.ammo_empty = pygame.mixer.Sound(SND_AMMO_EMPTY)
        self.fire_loop  = pygame.mixer.Sound(SND_FIRE_LOOP)
        self.fire_tail  = pygame.mixer.Sound(SND_FIRE_TAIL)
        self.reload     = pygame.mixer.Sound(SND_RELOAD)

        # Dedicated channels (prevents "sound won't play" when another is active).
        self.ch_fire_loop  = pygame.mixer.Channel(0)
        self.ch_fire_tail  = pygame.mixer.Channel(1)
        self.ch_reload     = pygame.mixer.Channel(2)
        self.ch_ammo_empty = pygame.mixer.Channel(3)

        self.lock = threading.Lock()

    @staticmethod
    def _require_file(path: str):
        with open(path, "rb"):
            pass

    def play_ammo_empty(self):
        with self.lock:
            # Don't let repeated empty clicks stack forever.
            self.ch_ammo_empty.stop()
            self.ch_ammo_empty.play(self.ammo_empty)

    def start_fire_loop(self):
        with self.lock:
            if not self.ch_fire_loop.get_busy():
                self.ch_fire_loop.play(self.fire_loop, loops=-1)

    def stop_fire_loop(self):
        with self.lock:
            if self.ch_fire_loop.get_busy():
                self.ch_fire_loop.stop()

    def play_fire_tail(self):
        with self.lock:
            # Tail should be audible even if reload/empty is happening.
            self.ch_fire_tail.stop()
            self.ch_fire_tail.play(self.fire_tail)

    def play_reload(self):
        with self.lock:
            self.ch_reload.stop()
            self.ch_reload.play(self.reload)


# --------------------------
# HUD (Arduino over USB)
# --------------------------
class HudSerial:
    def __init__(self, port: str, baud: int):
        # Note: The Arduino typically resets when the serial port opens.
        self.ser = serial.Serial(port, baudrate=baud, timeout=0.1)
        time.sleep(2.0)
        self.send_weapon(WEAPON_NAME)

    def _send_line(self, line: str):
        try:
            self.ser.write((line.strip() + "\n").encode("utf-8"))
        except Exception:
            pass

    def send_weapon(self, name: str):
        self._send_line(f"W:{name}")

    def send_ammo(self, cur: int, maxv: int):
        self._send_line(f"A:{cur},{maxv}")

    def send_reload_ms(self, ms_remaining: int):
        self._send_line(f"R:{ms_remaining}")

    def send_reload_done(self):
        self._send_line("RDONE")


# --------------------------
# HUD (disabled / no Arduino)
# --------------------------
class HudNull:
    """No-op HUD.

    This variant is intended for builds where there is no Arduino/TFT HUD.
    """
    def __init__(self, *_, **__):
        pass

    def send_weapon(self, *_):
        pass

    def send_ammo(self, *_):
        pass

    def send_reload_ms(self, *_):
        pass

    def send_reload_done(self, *_):
        pass



# --------------------------
# IR TX: NEC Extended (16/16)
# --------------------------
class NecIrTx:
    def __init__(self, gpio: int, carrier_hz: int = 38000, duty: float = 1/3, active_high: bool = True):
        self.gpio = gpio
        self.carrier_hz = carrier_hz
        self.duty = duty
        self.active_high = active_high

        self.pi = pigpio.pi()
        if not self.pi.connected:
            raise RuntimeError("pigpio not connected. Start: sudo systemctl enable --now pigpiod")

        self.pi.set_mode(self.gpio, pigpio.OUTPUT)
        self.pi.write(self.gpio, 0 if self.active_high else 1)

        self.period_us = int(round(1_000_000 / self.carrier_hz))
        self.on_us = max(1, int(round(self.period_us * self.duty)))
        self.off_us = max(1, self.period_us - self.on_us)

        self._wave_cache = {}

    def _pulse_on(self, us: int):
        if self.active_high:
            return pigpio.pulse(1 << self.gpio, 0, us)
        return pigpio.pulse(0, 1 << self.gpio, us)

    def _pulse_off(self, us: int):
        if self.active_high:
            return pigpio.pulse(0, 1 << self.gpio, us)
        return pigpio.pulse(1 << self.gpio, 0, us)

    def _carrier_mark(self, micros: int):
        pulses = []
        cycles = micros // self.period_us
        rem = micros % self.period_us

        for _ in range(int(cycles)):
            pulses.append(self._pulse_on(self.on_us))
            pulses.append(self._pulse_off(self.off_us))

        if rem > 0:
            on = min(self.on_us, rem)
            off = max(0, rem - on)
            pulses.append(self._pulse_on(on))
            if off > 0:
                pulses.append(self._pulse_off(off))
        return pulses

    def _space(self, micros: int):
        return [self._pulse_off(micros)]

    def _build_wave(self, address16: int, command16: int):
        b0 = address16 & 0xFF
        b1 = (address16 >> 8) & 0xFF
        b2 = command16 & 0xFF
        b3 = (command16 >> 8) & 0xFF

        pulses = []
        pulses += self._carrier_mark(9000)
        pulses += self._space(4500)

        for byte in (b0, b1, b2, b3):
            for bit in range(8):
                bitval = (byte >> bit) & 1
                pulses += self._carrier_mark(560)
                pulses += self._space(1690 if bitval else 560)

        pulses += self._carrier_mark(560)
        return pulses

    def send_raw(self, address16: int, command16: int, repeats: int = 1):
        """Transmit a fixed NEC-Extended 16/16 frame (bring-up mode)."""
        address16 &= 0xFFFF
        command16 &= 0xFFFF
        key = (address16 << 16) | command16

        wave_id = self._wave_cache.get(key)
        if wave_id is None:
            self.pi.wave_add_new()
            self.pi.wave_add_generic(self._build_wave(address16, command16))
            wave_id = self.pi.wave_create()
            if wave_id < 0:
                raise RuntimeError("Failed to create pigpio wave (out of resources).")
            self._wave_cache[key] = wave_id

        for _ in range(max(1, int(repeats))):
            self.pi.wave_send_once(wave_id)
            while self.pi.wave_tx_busy():
                time.sleep(0.001)
            if repeats > 1:
                time.sleep(IR_REPEAT_GAP_MS / 1000.0)

    def send_lazertag(self, team_id: int, player_id: int, weapon_id: int, damage: int, repeats: int = 1):
        address16 = ((team_id & 0xFF) << 8) | (player_id & 0xFF)
        command16 = ((weapon_id & 0xFF) << 8) | (damage & 0xFF)
        key = (address16 << 16) | command16

        wave_id = self._wave_cache.get(key)
        if wave_id is None:
            self.pi.wave_add_new()
            self.pi.wave_add_generic(self._build_wave(address16, command16))
            wave_id = self.pi.wave_create()
            if wave_id < 0:
                raise RuntimeError("Failed to create pigpio wave (out of resources).")
            self._wave_cache[key] = wave_id

        for _ in range(max(1, repeats)):
            self.pi.wave_send_once(wave_id)
            while self.pi.wave_tx_busy():
                time.sleep(0.001)
            if repeats > 1:
                time.sleep(IR_REPEAT_GAP_MS / 1000.0)

    def shutdown(self):
        try:
            self.pi.wave_tx_stop()
        except Exception:
            pass
        try:
            for wid in list(self._wave_cache.values()):
                try:
                    self.pi.wave_delete(wid)
                except Exception:
                    pass
            self._wave_cache.clear()
        except Exception:
            pass
        try:
            self.pi.write(self.gpio, 0 if self.active_high else 1)
        except Exception:
            pass
        try:
            self.pi.stop()
        except Exception:
            pass

# --------------------------
# INPUT: Mouse auto-detection
# --------------------------
def find_mouse_event_device() -> str:
    try:
        import glob
        candidates = sorted(glob.glob("/dev/input/by-id/*-event-mouse"))
        if candidates:
            return candidates[0]
    except Exception:
        pass

    for path in list_devices():
        try:
            dev = InputDevice(path)
            caps = dev.capabilities()
            has_buttons = (ecodes.EV_KEY in caps) and (ecodes.BTN_LEFT in caps[ecodes.EV_KEY])
            has_rel = (ecodes.EV_REL in caps) and (ecodes.REL_X in caps[ecodes.EV_REL]) and (ecodes.REL_Y in caps[ecodes.EV_REL])
            if has_buttons and has_rel:
                return path
        except Exception:
            continue

    raise FileNotFoundError("No mouse-like input device found under /dev/input.")

# --------------------------
# MAIN CONTROLLER
# --------------------------
class NerfGunController:
    def __init__(self):
        GPIO.setwarnings(False)
        GPIO.setmode(GPIO.BCM)

        # Muzzle LED flash (separate from IR TX)
        self.muzzle_led = MuzzleLed(
            MUZZLE_LED_GPIO,
            pwm_hz=MUZZLE_LED_PWM_HZ,
            duty_8bit=MUZZLE_LED_DUTY_8BIT,
            pulse_ms=MUZZLE_PULSE_MS,
            active_high=MUZZLE_LED_ACTIVE_HIGH,
        )

        self.ir_tx = NecIrTx(IR_TX_GPIO, IR_CARRIER_HZ, IR_DUTY, active_high=IR_ACTIVE_HIGH)

        self.audio = Audio()
        self.hud = HudSerial(SERIAL_PORT, SERIAL_BAUD)
        self.hud.send_weapon(WEAPON_NAME)

        self.mag_max = MAG_CAPACITY
        self.mag_cur = MAG_CAPACITY
        self.is_firing = False
        self.is_reloading = False
        self.last_shot_time = 0.0

        self.hud.send_ammo(self.mag_cur, self.mag_max)

        self.fire_thread = threading.Thread(target=self._fire_loop_worker, daemon=True)
        self.fire_thread.start()

        mouse_path = MOUSE_EVENT_PATH or find_mouse_event_device()
        print(f"[INPUT] Using mouse device: {mouse_path}")
        self.mouse = InputDevice(mouse_path)

    def _fire_loop_worker(self):
        while True:
            if not self.is_firing:
                time.sleep(0.005)
                continue

            if self.is_reloading:
                time.sleep(0.01)
                continue

            now = time.time()
            if (now - self.last_shot_time) < SECONDS_PER_SHOT:
                time.sleep(0.001)
                continue

            if self.mag_cur <= 0:
                self.audio.stop_fire_loop()
                self.audio.play_ammo_empty()
                self.is_firing = False
                continue

            self.last_shot_time = now
            self.mag_cur -= 1

            # Visual feedback: brief muzzle flash
            try:
                self.muzzle_led.pulse()
            except Exception:
                pass

            self.audio.start_fire_loop()

            if SIMPLE_IR_MODE:
                self.ir_tx.send_raw(SIMPLE_ADDRESS16, SIMPLE_COMMAND16, repeats=1)
            else:
                self.ir_tx.send_lazertag(TEAM_ID, PLAYER_ID, WEAPON_ID, DAMAGE_PER_BULLET, repeats=IR_REPEATS)

            self.hud.send_ammo(self.mag_cur, self.mag_max)

            if self.mag_cur == 0:
                self.audio.stop_fire_loop()
                self.audio.play_fire_tail()
                self.is_firing = False

    def start_firing(self):
        if self.is_reloading:
            return
        if self.mag_cur <= 0:
            self.audio.play_ammo_empty()
            self.is_firing = False
            return
        self.is_firing = True

    def stop_firing(self):
        if self.is_firing:
            self.is_firing = False
            if self.mag_cur > 0:
                self.audio.stop_fire_loop()
                self.audio.play_fire_tail()
            else:
                self.audio.stop_fire_loop()

    def reload(self):
        if self.is_reloading:
            return
        self.stop_firing()

        self.is_reloading = True
        self.audio.play_reload()

        end_time = time.time() + RELOAD_SECONDS
        while True:
            remaining = end_time - time.time()
            if remaining <= 0:
                break
            self.hud.send_reload_ms(int(remaining * 1000))
            time.sleep(0.1)

        self.mag_cur = self.mag_max
        self.hud.send_reload_done()
        self.hud.send_ammo(self.mag_cur, self.mag_max)
        self.is_reloading = False

    def run_mouse(self):
        LEFT = ecodes.BTN_LEFT
        RIGHT = ecodes.BTN_RIGHT

        print("LAZERNERF Controller (mouse mode)")
        print("Left click = FIRE (hold for auto), Right click = RELOAD")
        if SIMPLE_IR_MODE:
            print(f"NEC payload (SIMPLE): address16=0x{SIMPLE_ADDRESS16:04X} command16=0x{SIMPLE_COMMAND16:04X}")
        else:
            print(f"NEC payload: Team={TEAM_ID} Player={PLAYER_ID} Weapon={WEAPON_ID} Damage={DAMAGE_PER_BULLET}")
        print(f"IR: GPIO={IR_TX_GPIO} carrier={IR_CARRIER_HZ}Hz duty={IR_DUTY:.2f} repeats={IR_REPEATS} active_high={IR_ACTIVE_HIGH}")
        print("Ctrl+C to exit")

        for event in self.mouse.read_loop():
            if event.type != ecodes.EV_KEY:
                continue

            if event.code == LEFT:
                if event.value == 1:
                    self.start_firing()
                elif event.value == 0:
                    self.stop_firing()

            elif event.code == RIGHT:
                if event.value == 1:
                    threading.Thread(target=self.reload, daemon=True).start()

    def shutdown(self):
        try:
            self.ir_tx.shutdown()
        except Exception:
            pass
        try:
            self.muzzle_led.shutdown()
        except Exception:
            pass
        GPIO.cleanup()

def run_test_mode():
    tx = None
    try:
        tx = NecIrTx(IR_TX_GPIO, IR_CARRIER_HZ, IR_DUTY, active_high=IR_ACTIVE_HIGH)
        print("[TEST] Sending one NEC packet per second. Ctrl+C to stop.")
        print(f"[TEST] GPIO={IR_TX_GPIO} active_high={IR_ACTIVE_HIGH} repeats={IR_REPEATS}")
        if SIMPLE_IR_MODE:
            print(f"[TEST] Payload (SIMPLE): address16=0x{SIMPLE_ADDRESS16:04X} command16=0x{SIMPLE_COMMAND16:04X}")
        else:
            print(f"[TEST] Payload: Team={TEAM_ID} Player={PLAYER_ID} Weapon={WEAPON_ID} Damage={DAMAGE_PER_BULLET}")

        while True:
            if SIMPLE_IR_MODE:
                tx.send_raw(SIMPLE_ADDRESS16, SIMPLE_COMMAND16, repeats=1)
            else:
                tx.send_lazertag(TEAM_ID, PLAYER_ID, WEAPON_ID, DAMAGE_PER_BULLET, repeats=IR_REPEATS)
            time.sleep(1.0)
    except KeyboardInterrupt:
        pass
    finally:
        if tx:
            tx.shutdown()

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--test", action="store_true", help="Transmit a known packet once per second (no input needed).")
    ap.add_argument("--active-low", action="store_true", help="Use if your IR transmitter module turns on when SIG is LOW.")
    args = ap.parse_args()

    global IR_ACTIVE_HIGH
    if args.active_low:
        IR_ACTIVE_HIGH = False

    if args.test:
        run_test_mode()
        return

    ctl = None
    try:
        ctl = NerfGunController()
        ctl.run_mouse()
    finally:
        if ctl is not None:
            ctl.shutdown()

if __name__ == "__main__":
    main()