import rp2
from machine import Pin
import time

#----------------------------------------------------------------
# Gamepad Pmod hardware simulator (using the RP2040 PIO peripheral)
# Copyright (C) 2025, Tiny Tapeout LTD
# SPDX-License-Identifier: Apache-2.0
# Author: Uri Shaked
#
# This program uses:
#   - an OUT pin (data) that is driven by the OUT instruction
#   - two sideset pins: the first will be used for the clock (clk),
#     the second for the latch (latch).
#
# (At 1 MHz SM frequency, the bit‐transmission takes 1 + 24×2 + 2 = 51 cycles.
# With delay count = 1600, the full cycle is about (51+1600)=1651 µs ≈ 600 Hz.)
#----------------------------------------------------------------

@rp2.asm_pio(
    out_init=(rp2.PIO.OUT_LOW, ),           # data pin starts low
    in_shiftdir=rp2.PIO.SHIFT_RIGHT,
    out_shiftdir=rp2.PIO.SHIFT_RIGHT,
    sideset_init=(rp2.PIO.OUT_LOW, ) * 2,   # default: clk=0, latch=0
)
def gamepad_tx():
    pull()                              # Block until a word is in FIFO; load OSR.
    mov(x, osr)                         # Save OSR to X (so we have a “cached” value).
    set(y, 25)                           
    in_(y, 26)                          # ISR <- 1600 (0x640)
    wrap_target()
    pull(noblock)                       # If a new word is waiting, update OSR.
    mov(x, osr)                         # Save OSR to X (so we have a “cached” value).
    # --- Bit shifting loop (24 bits) ---
    set(y, 23)             .side(0b00)  # Initialize Y counter; normal state.
    label("bit_loop")
    out(pins, 1)           .side(0b00)  # Shift one bit (data on OUT pin) with clock low.
    jmp(y_dec, "bit_loop") .side(0b10)  # Data valid on clock's rising edge.
    # --- Latch pulse ---
    set(pins, 0)           .side(0b01) 
    nop()                  .side(0b00)
    # --- Delay loop ---
    mov(y, isr)                         # ISR is 1600, giving us ~1/600 period at 1 MHz.
    label("delay_loop")
    jmp(y_dec, "delay_loop")
    wrap()



class Gamepad:
    def __init__(self, data_pin, latch_pin, clk_pin, sm_id=0, freq=1000000):
        """
        Initialize the Gamepad interface.
        
        Arguments:
          data_pin : the GPIO number for the data output.
          clk_pin  : the GPIO number for the clock output (first sideset pin).
          latch_pin: the GPIO number for the latch output (second sideset pin).
          sm_id    : which PIO state machine to use (default 0).
          freq     : state machine clock frequency (default 1 MHz, giving a ~500 kHz bit clock).
        """
        assert clk_pin == latch_pin + 1

        # Initialize the pins.
        self.data_pin = Pin(data_pin, Pin.OUT)
        self.clk_pin = Pin(clk_pin, Pin.OUT)
        self.latch_pin = Pin(latch_pin, Pin.OUT)
        
        # 0xffffff means both controllers are disconnected
        self._last_value = 0xffffff

        # Create the state machine.
        # The state machine is configured with:
        #  - out_base set to the data pin.
        #  - sideset_base set to the clock pin (latch will be the next pin).
        self.sm = rp2.StateMachine(sm_id, gamepad_tx,
                                   freq=freq,
                                   out_base=self.data_pin,
                                   sideset_base=self.latch_pin)
        self.running = False

    def start(self):
        """
        Start the PIO state machine.
        """
        if not self.running:
            self.sm.active(1)
            self.sm.put(self._last_value)
            self.running = True

    def stop(self):
        """
        Disable the PIO machine.
        """
        self.sm.active(0)
        self.running = False

    def send(self, value: int):
        """
        Set a new 24-bit value to transmit.
        """
        self.sm.put(value)
        self._last_value = value

#----------------------------------------------------------------
# Example usage
#----------------------------------------------------------------
if __name__ == "__main__":
    # Create a Gamepad instance using (for example) GPIO0 for data,
    # GPIO1 for clock, and GPIO2 for latch.
    gp = Gamepad(latch_pin=17, clk_pin=18, data_pin=19, sm_id=0, freq=1000000)
    
    # Set an initial 24-bit value (here, 0xABCDEF).
    print("Sending:", hex(0xabcdef))
    gp.send(0xABCDEF)
    
    # Start the state machine.
    gp.start()
    
    try:
        i = 0
        while True:
            time.sleep(0.2)
            print("Sending:", hex(1 << i))
            gp.send(1 << i)
            i += 1
            if i == 24:
                i = 0
    except KeyboardInterrupt:
        gp.stop()
        print("Gamepad stopped.")