# Import necessary libraries for hardware control and timing
from machine import Pin, PWM
import utime

#LCD1602 Class to interface with the LCD display
class Lcd1602:
    def __init__(self, rs, en, d4, d5, d6, d7):
        self.rs = Pin(rs, Pin.OUT)
        self.en = Pin(en, Pin.OUT)
        self.data_pins = [Pin(d4, Pin.OUT), Pin(d5, Pin.OUT),
                          Pin(d6, Pin.OUT), Pin(d7, Pin.OUT)]
        self.init_lcd()
        self.clear()

    # Send enable pulse to latch data into LCD
    def pulse_en(self):
        self.en.low()
        utime.sleep_us(2)        # Wait briefly to settle
        self.en.high()
        utime.sleep_us(2)        # Short wait while EN is high
        self.en.low()
        utime.sleep_us(100)      # Wait for LCD to process the data

    # Send data or command to LCD
    def send(self, data, mode):
        self.rs.value(mode)  # Set RS pin (0 for command, 1 for data)
        # Send high nibble (upper 4 bits)
        for i in range(4):
            self.data_pins[i].value((data >> (i + 4)) & 0x01)
        self.pulse_en()
        # Send low nibble (lower 4 bits)
        for i in range(4):
            self.data_pins[i].value((data >> i) & 0x01)
        self.pulse_en()

    # Send command byte to LCD
    def command(self, cmd):
        self.send(cmd, 0)
        if cmd <= 3: # Commands like clear display or return home need extra time
            utime.sleep_ms(5)

    def putchar(self, char):
        self.send(ord(char), 1)

    def putstr(self, string):
        for char in string:
            self.putchar(char)

    def clear(self):
        self.command(0x01)  # Clear display command
        self.command(0x02)  # Return home command

    def move_to(self, col, row):
        row_offsets = [0x00, 0x40] # DDRAM addresses for line 0 and line 1
        addr = 0x80 + row_offsets[row] + col # Base DDRAM address + offset
        self.command(addr)

    # Initialize LCD with proper sequence
    def init_lcd(self):
        utime.sleep_ms(50)  # Wait for LCD to power up
        self.rs.low()
        self.en.low()
        # Initialization sequence for 4-bit mode
        self.send(0x33, 0)  # 8-bit mode init (repeated for robustness)
        self.send(0x32, 0)  # Set to 4-bit mode
        self.command(0x28)  # Function set: 4-bit, 2 line, 5x8 dots
        self.command(0x0C)  # Display ON, cursor OFF, blink OFF
        self.command(0x06)  # Entry mode: increment cursor, no shift
        self.command(0x01)  # Clear display
        utime.sleep_ms(5)   # Extra delay after clear

# Hardware device setup
lcd = Lcd1602(rs=12, en=11, d4=10, d5=9, d6=8, d7=7)  # LCD pin mapping
trigger = Pin(14, Pin.OUT)  # Ultrasonic sensor trigger pin
echo = Pin(16, Pin.IN)      # Ultrasonic sensor echo pin
servo = PWM(Pin(27))        # Servo motor PWM control
servo.freq(50)              # Set PWM frequency for servo

# Servo motor positions (duty cycle values for SG90 servo)
SERVO_CLOSED_POS = 4000 # Calibrated for container 1 / closed position
SERVO_OPEN_POS = 6000   # Calibrated for container 2 / open position

# Power button setup
power_button = Pin(15, Pin.IN, Pin.PULL_UP)

# Keypad matrix setup
matrix_keys = [['1', '2', '3'],
               ['4', '5', '6'],
               ['7', '8', '9'],
               ['*', '0', '#']]
keypad_rows = [26, 22, 21, 20]
keypad_columns = [19, 18, 17]
row_pins = [Pin(pin, Pin.OUT) for pin in keypad_rows]
col_pins = [Pin(pin, Pin.IN, Pin.PULL_DOWN) for pin in keypad_columns]

# Global variables to track system state
system_on = False
object_count = 0
object_currently_detected = False
pill_target1 = 0
pill_target2 = 0
# container_step is not used in the provided code, could be removed if not planned for future use.

# Function to measure distance using ultrasonic sensor
def measure_distance():
    trigger.low()
    utime.sleep_us(2)
    trigger.high()
    utime.sleep_us(10) # Trigger pulse
    trigger.low()

    # Wait for echo pin to go high (pulse start)
    start_time = utime.ticks_us()
    while echo.value() == 0:
        if utime.ticks_diff(utime.ticks_us(), start_time) > 50000: # 50ms timeout
            return 999  # Timeout protection, indicates error or no echo
    pulse_start = utime.ticks_us()

    # Wait for echo pin to go low (pulse end)
    start_time = utime.ticks_us() # Reset start_time for this loop's timeout
    while echo.value() == 1:
        if utime.ticks_diff(utime.ticks_us(), start_time) > 50000: # 50ms timeout
            return 999 # Timeout protection
    pulse_end = utime.ticks_us()

    pulse_duration = utime.ticks_diff(pulse_end, pulse_start)
    # Calculate distance in cm (Speed of sound = 343m/s or 0.0343 cm/us)
    # Distance = (duration * speed_of_sound) / 2 (round trip)
    return (pulse_duration * 0.0343) / 2

# Function to get user input from keypad with prompt
def scankeypad(prompt):
    user_input = ""
    lcd.clear()
    lcd.putstr(prompt[:16])  # Display first line of prompt
    if len(prompt) > 16:
        lcd.move_to(0, 1)
        lcd.putstr(prompt[16:32]) # Display second line of prompt if it exists
    else:
        lcd.move_to(0, 1) # Move to second line for input

    input_display_col = len(prompt) % 16 if len(prompt) <=16 else 0 # Determine start column for input on line 2
    if len(prompt) > 16 : # if prompt occupied part of line 2
        input_display_col = len(prompt[16:32])

    lcd.move_to(input_display_col, 1)


    while True:
        key_pressed = None
        for row_num, row_pin in enumerate(row_pins):
            # Drive one row high at a time
            for r_pin_idx, r_pin_val in enumerate(row_pins): # Ensure all other rows are low
                if r_pin_idx == row_num:
                    r_pin_val.high()
                else:
                    r_pin_val.low()
            utime.sleep_us(50) # Short delay for signal to stabilize

            for col_num, col_pin in enumerate(col_pins):
                if col_pin.value(): # If a column pin is high, a key is pressed
                    key_pressed = matrix_keys[row_num][col_num]
                    while col_pin.value(): # Wait for key release (debounce)
                        utime.sleep_ms(20)
                    break # Break from column scan
            row_pin.low() # Set row low again
            if key_pressed:
                break # Break from row scan

        if key_pressed:
            if key_pressed == '#': # Enter/confirm key
                if user_input: # Ensure some input was made
                    return int(user_input)
            elif key_pressed == '*': # Backspace/delete key
                if user_input:
                    user_input = user_input[:-1]
                    lcd.move_to(input_display_col, 1)
                    lcd.putstr(" " * (16 - input_display_col)) # Clear previous input on LCD
                    lcd.move_to(input_display_col, 1)
                    lcd.putstr(user_input)
            elif key_pressed.isdigit() and len(user_input) < 4: # Max 4 digits
                user_input += key_pressed
                lcd.putstr(key_pressed)
        utime.sleep_ms(50) # Overall keypad scan delay

# Function to select 1 or 2 containers
def wait_for_container_selection():
    lcd.clear()
    lcd.putstr("1 or 2 containers")
    lcd.move_to(0,1)
    lcd.putstr("Keypad: 1 or 2")
    while True:
        # Simplified keypad scan for just '1' or '2'
        key_pressed = None
        for row_num, row_pin in enumerate(row_pins):
            for r_pin_idx, r_pin_val in enumerate(row_pins):
                if r_pin_idx == row_num: r_pin_val.high()
                else: r_pin_val.low()
            utime.sleep_us(50)
            for col_num, col_pin in enumerate(col_pins):
                if col_pin.value():
                    key_pressed = matrix_keys[row_num][col_num]
                    while col_pin.value(): utime.sleep_ms(20)
                    break
            row_pin.low()
            if key_pressed: break

        if key_pressed in ["1", "2"]:
            return int(key_pressed)
        utime.sleep_ms(50)

# Function to count pills using ultrasonic detection
def count_pills(target):
    global object_count, object_currently_detected
    object_count = 0
    object_currently_detected = False # True if an object is currently in the detection zone
    lcd.clear()
    lcd.putstr(f"Count: 0/{target}")
    start_time_timeout = utime.ticks_ms() # For timeout if no pill is detected

    while object_count < target:
        dist = measure_distance()
        if dist < 5 and not object_currently_detected: # Threshold for pill detection
            object_count += 1
            object_currently_detected = True
            lcd.clear()
            lcd.putstr(f"Count: {object_count}/{target}")
            start_time_timeout = utime.ticks_ms()  # Reset timeout as a pill was detected
        elif dist >= 5 and object_currently_detected:
            object_currently_detected = False # Pill has passed

        # Timeout if no pill detected for 5 seconds (e.g., dispenser empty)
        if utime.ticks_diff(utime.ticks_ms(), start_time_timeout) > 5000:
            lcd.clear()
            lcd.putstr("Error: No pill")
            lcd.move_to(0,1)
            lcd.putstr("Timeout")
            utime.sleep(3)
            return False # Error: No pill detected in time
        utime.sleep_ms(100) # Polling interval for sensor
    return True # Successfully counted target pills

# Main Loop
lcd.clear()
lcd.putstr("Press button to")
lcd.move_to(0, 1)
lcd.putstr("start")

while True:
    if not system_on:
        if not power_button.value():  # Button is active low (PULL_UP)
            utime.sleep_ms(20)  # Debounce delay
            if not power_button.value(): # Check again after debounce
                while not power_button.value(): # Wait for button release
                    utime.sleep_ms(20)
                system_on = True

                # Initial prompt for first pill amount
                pill_target1 = scankeypad("Select Pill     Amount 1: ")
                if pill_target1 is None: # User might cancel or input error handling (not implemented here)
                    system_on = False # Reset
                    # Display initial message again
                    lcd.clear()
                    lcd.putstr("Press button to")
                    lcd.move_to(0, 1)
                    lcd.putstr("start")
                    continue


                container_choice = wait_for_container_selection()

                pill_target2 = 0 # Initialize pill_target2
                if container_choice == 2:
                    # Prompt for second pill amount if 2 containers selected
                    pill_target2 = scankeypad("Select Pill     Amount 2: ")
                    if pill_target2 is None:
                        system_on = False
                        lcd.clear()
                        lcd.putstr("Press button to")
                        lcd.move_to(0, 1)
                        lcd.putstr("start")
                        continue


                # Dispense for first container
                lcd.clear()
                lcd.putstr("Dispensing C1...")
                servo.duty_u16(SERVO_CLOSED_POS) # Ensure servo is at position for container 1
                utime.sleep(0.5) # Allow servo to move

                if count_pills(pill_target1):
                    if container_choice == 2 and pill_target2 > 0:
                        lcd.clear()
                        lcd.putstr("Switching to C2")
                        servo.duty_u16(SERVO_OPEN_POS)  # Move servo to position for container 2
                        utime.sleep(1) # Allow servo to move
                        lcd.clear()
                        lcd.putstr("Dispensing C2...")
                        if not count_pills(pill_target2):
                            lcd.clear()
                            lcd.putstr("C2 Error") # Error during second dispense
                            utime.sleep(3)
                            # System will reset below
                    elif container_choice == 2 and pill_target2 == 0:
                        lcd.clear()
                        lcd.putstr("C2 count is 0.")
                        lcd.move_to(0,1)
                        lcd.putstr("Skipping C2.")
                        utime.sleep(2)


                    lcd.clear()
                    lcd.putstr("Dispense Done!")
                else:
                    lcd.clear()
                    lcd.putstr("Dispense Error") # Error during first dispense
                
                utime.sleep(3) # Display message for 3 seconds

                # Reset system state for next operation
                lcd.clear()
                lcd.putstr("Press button to")
                lcd.move_to(0, 1)
                lcd.putstr("start")
                servo.duty_u16(SERVO_CLOSED_POS) # Reset servo to default position
                system_on = False
    utime.sleep_ms(100) # Main loop delay