import machine
import time
from machine import Pin, PWM

# Pin definitions
LED1_PIN = 2
LED2_PIN = 3
LED3_PIN = 4
SHCP_PIN = 5
STCP_PIN = 6
DS_PIN = 7
SERVO1_PIN = 8
SERVO2_PIN = 9
SERVO3_PIN = 10

# Initialize LEDs
led1 = Pin(LED1_PIN, Pin.OUT)
led2 = Pin(LED2_PIN, Pin.OUT)
led3 = Pin(LED3_PIN, Pin.OUT)

# Initialize shift register pins
shcp = Pin(SHCP_PIN, Pin.OUT)
stcp = Pin(STCP_PIN, Pin.OUT)
ds = Pin(DS_PIN, Pin.OUT)

# Initialize servos
servo1 = PWM(Pin(SERVO1_PIN))
servo2 = PWM(Pin(SERVO2_PIN))
servo3 = PWM(Pin(SERVO3_PIN))

# Configure servos with 50Hz frequency (standard for most servos)
servo1.freq(50)
servo2.freq(50)
servo3.freq(50)

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

# Helper function to convert angle to duty cycle for servos
def angle_to_duty(angle):
    # Map angle 0-180 to duty cycle range for servos
    # Typically 0.5ms (min) to 2.5ms (max) in a 20ms period
    # This translates to about 2.5% to 12.5% duty cycle
    min_duty = 1638  # ~2.5% of 65535
    max_duty = 8192  # ~12.5% of 65535
    return int(min_duty + (max_duty - min_duty) * angle / 180)

# Function to shift out data to the 74HC595 shift registers (MSB first)
def shift_out(data_pin, clock_pin, byte_value):
    for i in range(7, -1, -1):  # Shift MSB first
        data_pin.value((byte_value >> i) & 1)
        clock_pin.value(0)
        clock_pin.value(1)  # Data shifts on rising edge

# Function to display a two-digit 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]
    
    # First shift out ones digit, then tens digit
    shift_out(ds, shcp, ones_code)
    shift_out(ds, shcp, tens_code)
    
    # Latch the data to output pins
    digitalWrite(stcp, 1)
    digitalWrite(stcp, 0)

# Helper function to mimic Arduino's digitalWrite behavior
def digitalWrite(pin, value):
    pin.value(value)

# Function to update LEDs based on number
def update_leds(number):
    units_digit = number % 10
    led1.value(1 if 0 <= units_digit <= 2 else 0)
    led2.value(1 if 3 <= units_digit <= 5 else 0)
    led3.value(1 if 6 <= units_digit <= 9 else 0)

# Function to update servos based on number
def update_servos(number):
    tens_digit = number // 10
    if 0 <= tens_digit <= 2:
        servo1.duty_u16(angle_to_duty(30))
        servo2.duty_u16(angle_to_duty(60))
        servo3.duty_u16(angle_to_duty(90))
    elif 3 <= tens_digit <= 5:
        servo1.duty_u16(angle_to_duty(60))
        servo2.duty_u16(angle_to_duty(90))
        servo3.duty_u16(angle_to_duty(120))
    elif 6 <= tens_digit <= 9:
        servo1.duty_u16(angle_to_duty(90))
        servo2.duty_u16(angle_to_duty(120))
        servo3.duty_u16(angle_to_duty(150))

# Main program
current_number = 0
last_update_time = 0

# Initial setup
display_number(current_number)
update_leds(current_number)
update_servos(current_number)

# Main loop
while True:
    current_time = time.ticks_ms()
    if time.ticks_diff(current_time, last_update_time) >= 2000:  # 2 seconds
        current_number += 1
        if current_number > 99:
            current_number = 0
        
        display_number(current_number)
        update_leds(current_number)
        update_servos(current_number)
        
        last_update_time = current_time