from machine import Pin, PWM
import time

# Pin definitions
SHCP_PIN = Pin(2, Pin.OUT)  # Shift register clock (was A2)
STCP_PIN = Pin(1, Pin.OUT)  # Storage register clock (was A1)
DS_PIN = Pin(0, Pin.OUT)    # Serial data input (was A0)
LED1_PIN = Pin(3, Pin.OUT)  # LED 1 (was 2)
LED2_PIN = Pin(4, Pin.OUT)  # LED 2 (was 3)
LED3_PIN = Pin(5, Pin.OUT)  # LED 3 (was 4)

# RGB LED PWM setup
RGB1_R_PIN = PWM(Pin(6))    # RGB1 Red (was 5)
RGB1_G_PIN = PWM(Pin(7))    # RGB1 Green (was 6)
RGB1_B_PIN = PWM(Pin(8))    # RGB1 Blue (was 7)
RGB2_R_PIN = PWM(Pin(9))    # RGB2 Red (was 8)
RGB2_G_PIN = PWM(Pin(10))   # RGB2 Green (was 9)
RGB2_B_PIN = PWM(Pin(11))   # RGB2 Blue (was 10)

# Initialize PWM with frequency of 1000Hz
for pin in [RGB1_R_PIN, RGB1_G_PIN, RGB1_B_PIN, RGB2_R_PIN, RGB2_G_PIN, RGB2_B_PIN]:
    pin.freq(1000)

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

# Initialize variables
countdown = 9
led_state = 0
rgb_color = 0

def shift_out(data_pin, clock_pin, bit_order, value):
    """Emulate Arduino's shiftOut function"""
    for i in range(8):
        if bit_order == 1:  # MSBFIRST
            bit = (value >> (7 - i)) & 1
        else:  # LSBFIRST
            bit = (value >> i) & 1
        
        data_pin.value(bit)
        clock_pin.value(1)
        clock_pin.value(0)

def display_number(number):
    """Display a two-digit number on the 7-segment displays"""
    tens = number // 10
    ones = number % 10
    tens_code = segment_codes[tens]
    ones_code = segment_codes[ones]
    
    # Arduino's shiftOut is MSB first (MSBFIRST = 1)
    shift_out(DS_PIN, SHCP_PIN, 1, ones_code)
    shift_out(DS_PIN, SHCP_PIN, 1, tens_code)
    
    STCP_PIN.value(1)
    STCP_PIN.value(0)

def update_leds():
    """Update the state of the three LEDs"""
    global led_state
    
    LED1_PIN.value(1 if led_state == 0 else 0)
    LED2_PIN.value(1 if led_state == 1 else 0)
    LED3_PIN.value(1 if led_state == 2 else 0)
    
    led_state = (led_state + 1) % 3

def pwm_duty(value):
    """Convert 0-255 (Arduino) to 0-65535 (MicroPython PWM) for duty cycle"""
    return int(value * 65535 / 255)

def update_rgbs():
    """Update the color of the RGB LEDs"""
    global rgb_color
    
    if rgb_color == 0:  # Red
        RGB1_R_PIN.duty_u16(pwm_duty(255))
        RGB1_G_PIN.duty_u16(pwm_duty(0))
        RGB1_B_PIN.duty_u16(pwm_duty(0))
        RGB2_R_PIN.duty_u16(pwm_duty(255))
        RGB2_G_PIN.duty_u16(pwm_duty(0))
        RGB2_B_PIN.duty_u16(pwm_duty(0))
    elif rgb_color == 1:  # Green
        RGB1_R_PIN.duty_u16(pwm_duty(0))
        RGB1_G_PIN.duty_u16(pwm_duty(255))
        RGB1_B_PIN.duty_u16(pwm_duty(0))
        RGB2_R_PIN.duty_u16(pwm_duty(0))
        RGB2_G_PIN.duty_u16(pwm_duty(255))
        RGB2_B_PIN.duty_u16(pwm_duty(0))
    elif rgb_color == 2:  # Blue
        RGB1_R_PIN.duty_u16(pwm_duty(0))
        RGB1_G_PIN.duty_u16(pwm_duty(0))
        RGB1_B_PIN.duty_u16(pwm_duty(255))
        RGB2_R_PIN.duty_u16(pwm_duty(0))
        RGB2_G_PIN.duty_u16(pwm_duty(0))
        RGB2_B_PIN.duty_u16(pwm_duty(255))
    
    rgb_color = (rgb_color + 1) % 3

# Display initial value
display_number(99)

# Main loop
while True:
    display_number(countdown * 11)  # Shows numbers like 99, 88, 77...
    update_leds()
    update_rgbs()
    time.sleep(2)  # 2 seconds delay
    
    countdown -= 1
    
    if countdown == 0:
        countdown = 9
        LED1_PIN.value(0)
        LED2_PIN.value(0)
        LED3_PIN.value(0)
        
        # White color (all RGB components on)
        RGB1_R_PIN.duty_u16(pwm_duty(255))
        RGB1_G_PIN.duty_u16(pwm_duty(255))
        RGB1_B_PIN.duty_u16(pwm_duty(255))
        RGB2_R_PIN.duty_u16(pwm_duty(255))
        RGB2_G_PIN.duty_u16(pwm_duty(255))
        RGB2_B_PIN.duty_u16(pwm_duty(255))
        
        display_number(0)
        time.sleep(2)