from machine import Pin, I2C
import ssd1306
import utime

# ============================
# CONFIGURACIÓN DE HARDWARE
# ============================

# OLED I2C
i2c = I2C(0, scl=Pin(5), sda=Pin(4))
oled = ssd1306.SSD1306_I2C(128, 64, i2c)

# Sensores ultrasónicos
TRIG_A = Pin(2, Pin.OUT)
ECHO_A = Pin(3, Pin.IN)
TRIG_B = Pin(6, Pin.OUT)
ECHO_B = Pin(7, Pin.IN)

# Semáforo A
red_A = Pin(10, Pin.OUT)
yellow_A = Pin(11, Pin.OUT)
green_A = Pin(12, Pin.OUT)

# Semáforo B
red_B = Pin(13, Pin.OUT)
yellow_B = Pin(14, Pin.OUT)
green_B = Pin(15, Pin.OUT)

# ============================
# FUNCIONES
# ============================

def medir_distancia(trig, echo):
    """Mide la distancia con el sensor ultrasónico en cm."""
    trig.low()
    utime.sleep_us(2)
    trig.high()
    utime.sleep_us(10)
    trig.low()

    while echo.value() == 0:
        inicio = utime.ticks_us()
    while echo.value() == 1:
        fin = utime.ticks_us()

    duracion = fin - inicio
    distancia = (duracion * 0.0343) / 2
    return distancia

def muestrear_autos(trig, echo, umbral=2, tiempo_verde=5):
    """Cuenta autos mientras dura el verde."""
    contador = 0
    detectando = False
    inicio = utime.ticks_ms()

    while utime.ticks_diff(utime.ticks_ms(), inicio) < tiempo_verde * 1000:
        dist = medir_distancia(trig, echo)
        if dist <= umbral and not detectando:
            contador += 1
            detectando = True
        elif dist > umbral:
            detectando = False
        utime.sleep_ms(200)

    return contador

def clasificar_trafico(contador):
    """Clasifica tráfico según cantidad de autos detectados."""
    if contador <= 3:
        return "Ligero", 5
    elif contador <= 7:
        return "Pesado", 10
    else:
        return "Saturado", 15

def mostrar_trafico(estadoA, autosA, estadoB, autosB):
    """Muestra estado y autos en OLED."""
    oled.fill(0)
    oled.text("A: {} ({} autos)".format(estadoA, autosA), 0, 10)
    oled.text("B: {} ({} autos)".format(estadoB, autosB), 0, 30)
    oled.show()

def semaforo_A(t_verde):
    """Secuencia del semáforo A"""
    # Verde A
    red_A.off(); yellow_A.off(); green_A.on()
    red_B.on(); yellow_B.off(); green_B.off()

    autos = muestrear_autos(TRIG_A, ECHO_A, tiempo_verde=t_verde)

    # Amarillo A
    green_A.off(); yellow_A.on()
    utime.sleep(2)

    # Rojo A
    yellow_A.off(); red_A.on()
    return autos

def semaforo_B(t_verde):
    """Secuencia del semáforo B"""
    # Verde B
    red_B.off(); yellow_B.off(); green_B.on()
    red_A.on(); yellow_A.off(); green_A.off()

    autos = muestrear_autos(TRIG_B, ECHO_B, tiempo_verde=t_verde)

    # Amarillo B
    green_B.off(); yellow_B.on()
    utime.sleep(2)

    # Rojo B
    yellow_B.off(); red_B.on()
    return autos

# ============================
# PROGRAMA PRINCIPAL
# ============================

# Estados iniciales
estadoA, tiempoA, autosA = "Ligero", 5, 0
estadoB, tiempoB, autosB = "Ligero", 5, 0

try:
    while True:
        # Semáforo A
        autosA = semaforo_A(tiempoA)
        estadoA, tiempoA = clasificar_trafico(autosA)

        # Mostrar OLED
        mostrar_trafico(estadoA, autosA, estadoB, autosB)

        # Semáforo B
        autosB = semaforo_B(tiempoB)
        estadoB, tiempoB = clasificar_trafico(autosB)

        # Mostrar OLED
        mostrar_trafico(estadoA, autosA, estadoB, autosB)

except KeyboardInterrupt:
    oled.fill(0); oled.show()
    red_A.off(); yellow_A.off(); green_A.off()
    red_B.off(); yellow_B.off(); green_B.off()
    print("Programa detenido")