from machine import Pin, ADC, I2C, PWM
from time import sleep, ticks_ms, ticks_diff
import dht
import ssd1306

# ============================================================
# SESMAF DIGITAL TWIN - FULL SYSTEM
# Momentary-button version
#
# SW1 = ACB ON / RESET (press once)
# SW2 = MANUAL ACB TRIP (press once)
# SW3 = ACB OFF (press once)
# SW4 = Transformer mode toggle AUTO/MANUAL (press once)
# SW5 = Transformer RESET (press once)
#
# Protection:
#   Rated current = 5 A
#   OC pickup = 80% = 4.0 A, TMS=0.1
#   EF pickup = 10% = 0.5 A, TMS=0.1
#   IEC Standard Inverse:
#       t = 0.14*TMS / ((I/Ipickup)^0.02 - 1)
# ============================================================

# ---------------- GPIO ----------------
OC_PIN = 32
EF_PIN = 33
TX_TEMP_PIN = 34
TX_HUM_PIN = 35
DHT_PIN = 4

SDA_PIN = 21
SCL_PIN = 22

RELAY1_PIN = 23
RELAY2_PIN = 25
RELAY3_PIN = 26
RELAY4_PIN = 27
RELAY5_PIN = 14

SERVO_PIN = 18
BUZZER_PIN = 19

SW1_PIN = 13
SW2_PIN = 16
SW3_PIN = 17
SW4_PIN = 5
SW5_PIN = 36

OLED1_ADDR = 0x3C
OLED2_ADDR = 0x3D
LCD_ADDR = 0x27

# ---------------- Settings ----------------
RATED_CURRENT = 5.0

OC_PICKUP = RATED_CURRENT * 0.80
OC_TMS = 0.1

EF_PICKUP = RATED_CURRENT * 0.10
EF_TMS = 0.1

TX_TEMP_TRIP = 80.0
TX_HUM_TRIP = 90.0

# Momentary buttons: released = 0, pressed = 1

# ---------------- ADC ----------------
oc_adc = ADC(Pin(OC_PIN))
ef_adc = ADC(Pin(EF_PIN))
tx_temp_adc = ADC(Pin(TX_TEMP_PIN))
tx_hum_adc = ADC(Pin(TX_HUM_PIN))

for adc in (oc_adc, ef_adc, tx_temp_adc, tx_hum_adc):
    adc.atten(ADC.ATTN_11DB)

# ---------------- DHT22 ----------------
dht_sensor = dht.DHT22(Pin(DHT_PIN))

# ---------------- I2C ----------------
i2c = I2C(0, scl=Pin(SCL_PIN), sda=Pin(SDA_PIN), freq=100000)
print("I2C DEVICES:", [hex(x) for x in i2c.scan()])

# ---------------- OLED ----------------
oled1 = ssd1306.SSD1306_I2C(128, 64, i2c, addr=OLED1_ADDR)
oled2 = ssd1306.SSD1306_I2C(128, 64, i2c, addr=OLED2_ADDR)

# ============================================================
# LCD 1602 I2C DRIVER
# ============================================================

LCD_RS = 0x01
LCD_RW = 0x02
LCD_EN = 0x04
LCD_BL = 0x08

def lcd_write(data):
    i2c.writeto(LCD_ADDR, bytes([data | LCD_BL]))

def lcd_pulse(data):
    lcd_write(data | LCD_EN)
    sleep(0.001)
    lcd_write(data & ~LCD_EN)
    sleep(0.001)

def lcd_send_nibble(nibble, rs):
    data = nibble & 0xF0
    if rs:
        data |= LCD_RS
    lcd_pulse(data)

def lcd_send(value, rs):
    lcd_send_nibble(value & 0xF0, rs)
    lcd_send_nibble((value << 4) & 0xF0, rs)

def lcd_command(command):
    lcd_send(command, False)
    sleep(0.002)

def lcd_data(data):
    lcd_send(data, True)

def lcd_init():
    sleep(0.05)
    lcd_send_nibble(0x30, False)
    sleep(0.005)
    lcd_send_nibble(0x30, False)
    sleep(0.001)
    lcd_send_nibble(0x30, False)
    sleep(0.001)
    lcd_send_nibble(0x20, False)
    sleep(0.001)
    lcd_command(0x28)
    lcd_command(0x0C)
    lcd_command(0x06)
    lcd_command(0x01)
    sleep(0.005)

def lcd_clear():
    lcd_command(0x01)
    sleep(0.005)

def lcd_set_cursor(row, col):
    lcd_command((0x80 if row == 0 else 0xC0) + col)

def lcd_print(text):
    for ch in text:
        lcd_data(ord(ch))

def lcd_print_line(row, text):
    text = text[:16]
    if len(text) < 16:
        text += " " * (16 - len(text))
    lcd_set_cursor(row, 0)
    lcd_print(text)

lcd_init()
lcd_clear()

# ---------------- Outputs ----------------
relay1 = Pin(RELAY1_PIN, Pin.OUT)
relay2 = Pin(RELAY2_PIN, Pin.OUT)
relay3 = Pin(RELAY3_PIN, Pin.OUT)
relay4 = Pin(RELAY4_PIN, Pin.OUT)
relay5 = Pin(RELAY5_PIN, Pin.OUT)

# Wokwi relay modules are treated as active-HIGH here.
def relay_set(relay, on):
    relay.value(1 if on else 0)

servo = PWM(Pin(SERVO_PIN), freq=50)

def servo_angle(angle):
    duty = int(26 + (angle / 180.0) * (128 - 26))
    servo.duty(duty)

buzzer = PWM(Pin(BUZZER_PIN), freq=1000)
buzzer.duty(0)

def buzzer_on():
    buzzer.duty(512)

def buzzer_off():
    buzzer.duty(0)

# ---------------- Momentary buttons ----------------
# Released = GPIO LOW, pressed = GPIO HIGH.
# Button connects GPIO to 3.3 V; internal pulldown holds LOW when released.
sw1 = Pin(SW1_PIN, Pin.IN, Pin.PULL_DOWN)
sw2 = Pin(SW2_PIN, Pin.IN, Pin.PULL_DOWN)
sw3 = Pin(SW3_PIN, Pin.IN, Pin.PULL_DOWN)
sw4 = Pin(SW4_PIN, Pin.IN, Pin.PULL_DOWN)
sw5 = Pin(SW5_PIN, Pin.IN)  # GPIO15 has a default pull-up

prev_sw1 = 0
prev_sw2 = 0
prev_sw3 = 0
prev_sw4 = 0
prev_sw5 = 0

# ---------------- States ----------------
acb_on = False
protection_tripped = False
transformer_tripped = False

trip_source = "NONE"
trip_time = 0.0

transformer_mode_auto = True

protection_timer_start = None
protection_timer_source = None

oc_current = 0.0
ef_current = 0.0
dht_temp = 0.0
dht_hum = 0.0
tx_temp = 0.0
tx_hum = 0.0

lcd_page = 0
last_lcd_page = ticks_ms()

# ============================================================
# MEASUREMENT
# ============================================================

def read_current(adc):
    return (adc.read() / 4095.0) * 10.0

def read_manual_temp():
    return (tx_temp_adc.read() / 4095.0) * 120.0

def read_manual_hum():
    return (tx_hum_adc.read() / 4095.0) * 100.0

def update_dht():
    global dht_temp, dht_hum
    try:
        dht_sensor.measure()
        dht_temp = dht_sensor.temperature()
        dht_hum = dht_sensor.humidity()
    except Exception as e:
        print("DHT ERROR:", e)

def update_transformer_values():
    global tx_temp, tx_hum
    if transformer_mode_auto:
        tx_temp = dht_temp
        tx_hum = dht_hum
    else:
        tx_temp = read_manual_temp()
        tx_hum = read_manual_hum()

# ============================================================
# IDMT
# ============================================================

def idmt_time(current, pickup, tms):
    if current <= pickup:
        return None
    ratio = current / pickup
    denominator = (ratio ** 0.02) - 1.0
    if denominator <= 0:
        return None
    return (0.14 * tms) / denominator

# ============================================================
# ACB
# ============================================================

def acb_turn_on():
    global acb_on

    if protection_tripped or transformer_tripped:
        return False

    if sw3.value() == 1:
        print("ACB ON BLOCKED - SW3 OFF SWITCH IS ON")
        return False

    acb_on = True

    # Relay2 = ACB ON indication
    relay_set(relay2, True)
    # Relay3 = ACB OFF indication
    relay_set(relay3, False)

    servo_angle(180)
    print("ACB ON")
    return True

def acb_turn_off():
    global acb_on

    acb_on = False

    relay_set(relay2, False)
    relay_set(relay3, True)

    servo_angle(0)
    print("ACB OFF")

def acb_trip(source, elapsed=0.0):
    global acb_on, protection_tripped, trip_source, trip_time

    acb_on = False
    protection_tripped = True
    trip_source = source
    trip_time = elapsed

    relay_set(relay2, False)
    relay_set(relay3, True)

    # Relay5 = common trip / 3-phase trip indication
    relay_set(relay5, True)

    servo_angle(0)
    buzzer_on()

    print("================================")
    print("ACB TRIPPED")
    print("SOURCE:", trip_source)
    print("TIME:", round(trip_time, 3), "s")
    print("================================")

def reset_acb():
    global protection_tripped, trip_source, trip_time
    global protection_timer_start, protection_timer_source

    if tx_fault_active():
        print("ACB RESET BLOCKED - TRANSFORMER FAULT ACTIVE")
        return

    protection_tripped = False
    trip_source = "NONE"
    trip_time = 0.0

    protection_timer_start = None
    protection_timer_source = None

    relay_set(relay5, False)
    buzzer_off()

    acb_turn_on()
    print("ACB RESET / ON")

# ============================================================
# TRANSFORMER
# ============================================================

def tx_fault_active():
    return tx_temp >= TX_TEMP_TRIP or tx_hum >= TX_HUM_TRIP

def tx_fault_source():
    if tx_temp >= TX_TEMP_TRIP and tx_hum >= TX_HUM_TRIP:
        return "TX T+H"
    if tx_temp >= TX_TEMP_TRIP:
        return "TX TEMP"
    if tx_hum >= TX_HUM_TRIP:
        return "TX HUM"
    return "NONE"

def transformer_trip():
    global transformer_tripped

    transformer_tripped = True

    # Transformer trip forces ACB OFF
    acb_turn_off()

    # Relay4 = transformer trip indication
    relay_set(relay4, True)

    buzzer_on()

    print("TRANSFORMER TRIPPED:", tx_fault_source())

def transformer_reset():
    global transformer_tripped

    if tx_fault_active():
        print("TX RESET BLOCKED - FAULT STILL ACTIVE")
        return

    transformer_tripped = False
    relay_set(relay4, False)

    if not protection_tripped:
        buzzer_off()

    print("TRANSFORMER RESET")

# ============================================================
# OC / EF PROTECTION
# ============================================================

def protection_update():
    global protection_timer_start, protection_timer_source
    global trip_time

    if protection_tripped or transformer_tripped or not acb_on:
        protection_timer_start = None
        protection_timer_source = None
        return

    oc_t = idmt_time(oc_current, OC_PICKUP, OC_TMS)
    ef_t = idmt_time(ef_current, EF_PICKUP, EF_TMS)

    fastest_source = None
    fastest_time = None

    if oc_t is not None:
        fastest_source = "OC"
        fastest_time = oc_t

    if ef_t is not None and (fastest_time is None or ef_t < fastest_time):
        fastest_source = "EF"
        fastest_time = ef_t

    # No active fault -> cancel pending timer
    if fastest_source is None:
        protection_timer_start = None
        protection_timer_source = None
        return

    # Start or restart if the fastest source changes
    if protection_timer_start is None:
        protection_timer_start = ticks_ms()
        protection_timer_source = fastest_source
        print("IDMT START:", fastest_source, "calculated:", round(fastest_time, 3), "s")

    elif protection_timer_source != fastest_source:
        protection_timer_start = ticks_ms()
        protection_timer_source = fastest_source
        print("IDMT SOURCE CHANGED:", fastest_source)

    elapsed = ticks_diff(ticks_ms(), protection_timer_start) / 1000.0

    if elapsed >= fastest_time:
        trip_time = elapsed
        acb_trip(fastest_source, elapsed)

        protection_timer_start = None
        protection_timer_source = None

# ============================================================
# MOMENTARY BUTTON LOGIC
# ============================================================

def process_switches():
    global prev_sw1, prev_sw2, prev_sw3, prev_sw4, prev_sw5
    global transformer_mode_auto

    s1 = sw1.value()
    s2 = sw2.value()
    s3 = sw3.value()
    s4 = sw4.value()
    raw5 = sw5.value()
    s5 = 1 if raw5 == 0 else 0  # SW5 is wired to GND; logical pressed = 1

    # Only act on a new button press (0 -> 1).
    press1 = (s1 == 1 and prev_sw1 == 0)
    press2 = (s2 == 1 and prev_sw2 == 0)
    press3 = (s3 == 1 and prev_sw3 == 0)
    press4 = (s4 == 1 and prev_sw4 == 0)
    press5 = (s5 == 1 and prev_sw5 == 0)

    # SW1 = ACB ON / RESET
    if press1:
        if protection_tripped:
            reset_acb()
        elif not transformer_tripped:
            acb_turn_on()

    # SW2 = Manual ACB TRIP
    if press2:
        if acb_on and not protection_tripped:
            acb_trip("MANUAL", 0.0)
        else:
            print("MANUAL TRIP IGNORED - ACB IS OFF")

    # SW3 = ACB OFF
    if press3:
        acb_turn_off()

    # SW4 = Toggle Transformer AUTO/MANUAL
    if press4:
        transformer_mode_auto = not transformer_mode_auto
        print("TRANSFORMER MODE:",
              "AUTO" if transformer_mode_auto else "MANUAL")

    # SW5 = Transformer RESET
    if press5:
        if transformer_tripped:
            transformer_reset()
        else:
            print("TX RESET IGNORED - NO TRANSFORMER TRIP")

    prev_sw1 = s1
    prev_sw2 = s2
    prev_sw3 = s3
    prev_sw4 = s4
    prev_sw5 = s5


# ============================================================
# DISPLAYS
# ============================================================

def update_oled1():
    oled1.fill(0)
    oled1.text("SESMAF", 0, 0)
    oled1.text("PROTECTION", 0, 10)
    oled1.text("OC:%4.1f A" % oc_current, 0, 22)
    oled1.text("EF:%4.1f A" % ef_current, 0, 34)

    if protection_tripped:
        oled1.text("ACB:TRIPPED", 0, 46)
        oled1.text(trip_source, 80, 46)
    elif acb_on:
        oled1.text("ACB:ON", 0, 46)
    else:
        oled1.text("ACB:OFF", 0, 46)

    oled1.show()

def update_oled2():
    oled2.fill(0)
    oled2.text("TRANSFORMER", 0, 0)
    oled2.text("TEMP:%4.1f C" % tx_temp, 0, 14)
    oled2.text("HUM :%4.1f %%" % tx_hum, 0, 28)
    oled2.text("MODE:" + ("AUTO" if transformer_mode_auto else "MANUAL"), 0, 42)

    if transformer_tripped:
        oled2.text("TX:TRIPPED", 0, 54)
    else:
        oled2.text("TX:NORMAL", 0, 54)

    oled2.show()

def update_lcd():
    global lcd_page, last_lcd_page

    now = ticks_ms()

    if ticks_diff(now, last_lcd_page) >= 3000:
        lcd_page = (lcd_page + 1) % 3
        last_lcd_page = now

    if lcd_page == 0:
        if protection_tripped:
            line1 = "TRIP:" + trip_source
            line2 = "TIME:%5.2fs" % trip_time
        elif transformer_tripped:
            line1 = "TRIP:TRANSFORM"
            line2 = tx_fault_source()
        else:
            line1 = "TRIP:NO TRIP"
            line2 = "READY"

    elif lcd_page == 1:
        line1 = "SW1:ON/RST"
        line2 = "SW2:MAN TRIP"

    else:
        line1 = "SW3:OFF SW4:MODE"
        line2 = "SW5:TX RESET"

    lcd_print_line(0, line1)
    lcd_print_line(1, line2)

# ============================================================
# SERIAL STATUS
# ============================================================

def print_status():
    print("--------------------------------")
    print("OC:", round(oc_current, 2), "A / Pickup:", OC_PICKUP, "A")
    print("EF:", round(ef_current, 2), "A / Pickup:", EF_PICKUP, "A")
    print("TX:", round(tx_temp, 1), "C /", round(tx_hum, 1), "%")
    print("TX MODE:", "AUTO" if transformer_mode_auto else "MANUAL")
    sw5_logical = 1 if sw5.value() == 0 else 0
    print("SW:", sw1.value(), sw2.value(), sw3.value(), sw4.value(), sw5_logical)
    print("ACB:", "ON" if acb_on else "OFF")
    print("PROTECTION:", trip_source if protection_tripped else "NORMAL")
    print("TX STATUS:", "TRIPPED" if transformer_tripped else "NORMAL")
    print("--------------------------------")

# ============================================================
# STARTUP
# ============================================================

print("")
print("========================================")
print(" SESMAF DIGITAL TWIN")
print(" MOMENTARY BUTTON VERSION")
print("========================================")
print("OC PICKUP:", OC_PICKUP, "A")
print("OC TMS:", OC_TMS)
print("EF PICKUP:", EF_PICKUP, "A")
print("EF TMS:", EF_TMS)
print("TX TEMP TRIP:", TX_TEMP_TRIP, "C")
print("TX HUM TRIP:", TX_HUM_TRIP, "%")
print("========================================")

# Safe startup state
acb_turn_off()
relay_set(relay1, False)
relay_set(relay4, False)
relay_set(relay5, False)
buzzer_off()

lcd_print_line(0, "SESMAF DIGITAL")
lcd_print_line(1, "SYSTEM READY")

sleep(1)

last_dht = ticks_ms()
last_display = ticks_ms()
last_serial = ticks_ms()

# ============================================================
# MAIN LOOP
# ============================================================

while True:
    now = ticks_ms()

    # Read protection currents continuously
    oc_current = read_current(oc_adc)
    ef_current = read_current(ef_adc)

    # DHT22 every 2 seconds
    if ticks_diff(now, last_dht) >= 2000:
        update_dht()
        last_dht = now

    # Transformer measurement
    update_transformer_values()

    # Transformer protection
    if not transformer_tripped and tx_fault_active():
        transformer_trip()

    # Slide switches
    process_switches()

    # OC / EF IDMT
    protection_update()

    # Displays
    if ticks_diff(now, last_display) >= 250:
        update_oled1()
        update_oled2()
        update_lcd()
        last_display = now

    # Serial status
    if ticks_diff(now, last_serial) >= 3000:
        print_status()
        last_serial = now

    sleep(0.02)