from machine import Pin
import time

# Pin definitions
LED_PIN = 2          # Standard LED
RGB_RED_PIN = 3      # RGB LED Red pin
RGB_GREEN_PIN = 4    # RGB LED Green pin
RGB_BLUE_PIN = 5     # RGB LED Blue pin
SHCP_PIN = 8         # Shift Register clock pin
STCP_PIN = 7         # Shift Register latch pin
DS_PIN = 6           # Shift Register data pin

# Initialize pins
led = Pin(LED_PIN, Pin.OUT)
rgb_red = Pin(RGB_RED_PIN, Pin.OUT)
rgb_green = Pin(RGB_GREEN_PIN, Pin.OUT)
rgb_blue = Pin(RGB_BLUE_PIN, Pin.OUT)
shcp = Pin(SHCP_PIN, Pin.OUT)
stcp = Pin(STCP_PIN, Pin.OUT)
ds = Pin(DS_PIN, Pin.OUT)

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

# Global variables for state tracking
current_number = 0
last_led_update = 0
last_rgb_update = 0
last_display_update = 0
rgb_state = 0

# Function to shift out data to the 74HC595 shift registers
def shift_out(data_pin, clock_pin, value, bit_order=0):
    # Loop through each bit
    for i in range(8):
        if bit_order == 0:  # MSB first
            bit = (value >> (7 - i)) & 1
        else:  # LSB first
            bit = (value >> i) & 1
            
        # Set data pin
        data_pin.value(bit)
        
        # Pulse clock pin
        clock_pin.value(1)
        clock_pin.value(0)

# Function to set RGB LED color
def set_rgb_color(state):
    if state == 0:  # Red
        rgb_red.value(1)
        rgb_green.value(0)
        rgb_blue.value(0)
    elif state == 1:  # Green
        rgb_red.value(0)
        rgb_green.value(1)
        rgb_blue.value(0)
    elif state == 2:  # Blue
        rgb_red.value(0)
        rgb_green.value(0)
        rgb_blue.value(1)
    elif state == 3:  # Yellow (Red + Green)
        rgb_red.value(1)
        rgb_green.value(1)
        rgb_blue.value(0)
    elif state == 4:  # Cyan (Green + Blue)
        rgb_red.value(0)
        rgb_green.value(1)
        rgb_blue.value(1)
    elif state == 5:  # Magenta (Red + Blue)
        rgb_red.value(1)
        rgb_green.value(0)
        rgb_blue.value(1)
    elif state == 6:  # White (Red + Green + Blue)
        rgb_red.value(1)
        rgb_green.value(1)
        rgb_blue.value(1)

# Function to display a number on the 7-segment displays
def display_number(number):
    tens = number // 10
    ones = number % 10
    
    tens_code = segment_codes[tens]
    ones_code = segment_codes[ones]
    
    # Ensure latch is low before shifting data
    stcp.value(0)
    
    # Shift out data for ones digit (goes to sr2 via sr1)
    shift_out(ds, shcp, ones_code)
    
    # Shift out data for tens digit (goes to sr1)
    shift_out(ds, shcp, tens_code)
    
    # Latch the data after both digits are shifted
    stcp.value(1)
    stcp.value(0)

# Initialize displays and LEDs
display_number(0)
set_rgb_color(rgb_state)
rgb_state = (rgb_state + 1) % 7
last_rgb_update = time.ticks_ms()

# Main loop
while True:
    current_time = time.ticks_ms()
    
    # Update LED state every 2 seconds
    if time.ticks_diff(current_time, last_led_update) >= 2000:
        led.value(not led.value())
        last_led_update = current_time
    
    # Update RGB LED color every 2 seconds
    if time.ticks_diff(current_time, last_rgb_update) >= 2000:
        set_rgb_color(rgb_state)
        rgb_state = (rgb_state + 1) % 7
        last_rgb_update = current_time
    
    # Update display every 2 seconds
    if time.ticks_diff(current_time, last_display_update) >= 2000:
        current_number = (current_number + 1) % 100
        display_number(current_number)
        last_display_update = current_time