import rp2
from rp2 import PIO
from utime import sleep
from machine import Pin

@rp2.asm_pio(
    out_init=[PIO.OUT_LOW]*8,      # Initialize 8 output pins (segments A-G + DP) to LOW
    sideset_init=[PIO.OUT_LOW]*4   # Initialize 4 side-set pins (digit select lines) to LOW
)
def sevseg():
    wrap_target()                  # Mark the start of a loopable PIO block
    label("0")                     # Label this instruction as "0" for looping back later

    pull(noblock)           .side(0)
    """ Pull 32-bit data from FIFO into OSR (non-blocking); 
        while setting all digit-select lines LOW (avoid ghosting) """

    mov(x, osr)             .side(0)
    """ Move pulled data from OSR into register X; 
        keep digit lines LOW during transfer """

    out(pins, 8)            .side(1)
    """ Output 8 bits to segment pins (digit 1 data);
        Activate digit 1 by setting GPIO10 HIGH (side=1) """

    out(pins, 8)            .side(2)
    """ Output 8 bits for digit 2;
        Activate digit 2 by setting GPIO11 HIGH (side=2) """

    out(pins, 8)            .side(4)
    """ Output 8 bits for digit 3;
        Activate digit 3 by setting GPIO12 HIGH (side=4) """

    out(pins, 8)            .side(8)
    """ Output 8 bits for digit 4;
        Activate digit 4 by setting GPIO13 HIGH (side=8) """

    jmp("0")                .side(0)
    """ Loop back to label "0";
        Reset all digit select lines LOW before restarting """

    wrap()                         # End of the loopable PIO program

# --- Start State Machine for PIO --- #
sm = rp2.StateMachine(
    0,                # Use state machine 0
    sevseg,           # Load the sevseg PIO program
    freq=2000,        # Set clock frequency to 2kHz (controls refresh rate)
    out_base=Pin(2),  # Segment output pins start from GPIO2 (covers GPIO2–9)
    sideset_base=Pin(10)  # Digit select pins start from GPIO10 (GPIO10–13)
)
sm.active(1)          # Start the state machine

# 7-segment digit encodings (common anode)
digits = [
  0b11000000,  # 0
  0b11111001,  # 1
  0b10100100,  # 2 
  0b10110000,  # 3
  0b10011001,  # 4
  0b10010010,  # 5
  0b10000010,  # 6
  0b11111000,  # 7
  0b10000000,  # 8
  0b10011000,  # 9
]

# Convert a number (0–9999) to 4-digit 7-segment display format
def segmentize(num):
  return (
    digits[num % 10]                             # Units digit
    | digits[num // 10 % 10] << 8                # Tens digit
    | digits[num // 100 % 10] << 16              # Hundreds digit
    | digits[num // 1000 % 10] << 24             # Thousands digit
  )

counter = 0                   # Initialize counter

while True:
  sm.put(segmentize(counter))  # Send 32-bit encoded 7-segment data to PIO
  counter += 1                 # Increment counter
  sleep(0.1)                   # Delay 100ms before next update