# =============================================================================
# SMART NUMERICAL RELAY — VERSION 4 FINAL (مصحح كامل)
# ANSI 50/51/50N/51N/79/81 | MT 30kV Algérie
# =============================================================================

from machine import Pin, ADC, SoftI2C
from umqtt.simple import MQTTClient
from ssd1306 import SSD1306_I2C
import network
import time
import math
import ujson

# =============================================================================
# CONFIGURATION
# =============================================================================

WIFI_SSID     = "Wokwi-GUEST"
WIFI_PASSWORD = ""
MQTT_BROKER   = "1133c5f5d79b4a5db07d54df4e5806e3.s1.eu.hivemq.cloud"
MQTT_PORT     = 8883
MQTT_USER     = "admin"
MQTT_PASSWORD = "Admin12345"
CLIENT_ID     = "Sonelgaz_SmartRelay_V4_PFE"

T_MEASURE = b"mt/feeder1/measure"
T_ALARM   = b"mt/feeder1/alarm"
T_STATE   = b"mt/feeder1/state"
T_EVENT   = b"mt/feeder1/event"
T_RECLOSE = b"mt/feeder1/reclose"
T_CMD     = b"mt/feeder1/cmd"

# =============================================================================
# PARAMÈTRES ÉLECTRIQUES
# =============================================================================

CT_RATIO     = 40.0
I_NOMINAL    = 80.0
I_PICKUP_51  = 100.0
I_PICKUP_50  = 180.0
I_PICKUP_51N = 20.0
I_PICKUP_50N = 40.0

IEC_K        = 0.14
IEC_ALPHA    = 0.02
TMS_51       = 0.10
TMS_51N      = 0.05

FREQ_NOMINAL       = 50.0
FREQ_UF_WARNING    = 49.5
FREQ_UF_TRIP       = 48.5
FREQ_OF_WARNING    = 50.5
FREQ_OF_TRIP       = 51.5
FREQ_WARNING_DELAY = 3

RECLOSE_ENABLED      = True
MAX_RECLOSE_ATTEMPTS = 3
RECLOSE_DELAYS       = [2, 5, 10]
RECLOSE_WOKWI_FACTOR = 0.3

# =============================================================================
# HARDWARE
# =============================================================================

ct_adc   = ADC(Pin(34))
freq_adc = ADC(Pin(35))
ct_adc.atten(ADC.ATTN_11DB)
freq_adc.atten(ADC.ATTN_11DB)

relay_pin  = Pin(4,  Pin.OUT)
led_green  = Pin(2,  Pin.OUT)
led_yellow = Pin(16, Pin.OUT)
led_red    = Pin(17, Pin.OUT)
buzzer_pin = Pin(27, Pin.OUT)

i2c  = SoftI2C(scl=Pin(22), sda=Pin(21))
oled = SSD1306_I2C(128, 64, i2c)

relay_pin.value(1)
led_green.value(1)
led_yellow.value(0)
led_red.value(0)
buzzer_pin.value(0)

# =============================================================================
# VARIABLES GLOBALES
# =============================================================================

relay_state       = "NORMAL"
reclose_count     = 0
lockout_active    = False
remote_mode       = True

pickup_51_active  = False
pickup_51_start   = 0
pickup_51N_active = False
pickup_51N_start  = 0

freq_uf_warn_active = False
freq_uf_warn_start  = 0
freq_of_warn_active = False
freq_of_warn_start  = 0

mqtt_client = None

# =============================================================================
# HELPERS
# =============================================================================

def log_event(code, description, i_phase=0.0, i_earth=0.0, freq=0.0):
    print("[{}] {} | I={:.1f}A Ie={:.1f}A F={:.2f}Hz | State={}".format(
        code, description, i_phase, i_earth, freq, relay_state))

def set_leds(state):
    led_green.value(0)
    led_yellow.value(0)
    led_red.value(0)
    if state == "NORMAL":
        led_green.value(1)
    elif "PICKUP" in state or "WARN" in state or state == "TIMING":
        led_yellow.value(1)
    elif "TRIP" in state or "LOCKOUT" in state or state == "RECLOSE":
        led_red.value(1)
        led_yellow.value(1)

def buzzer_pulse(pattern="single"):
    if pattern == "single":
        buzzer_pin.value(1); time.sleep_ms(120); buzzer_pin.value(0)
    elif pattern == "double":
        for _ in range(2):
            buzzer_pin.value(1); time.sleep_ms(100)
            buzzer_pin.value(0); time.sleep_ms(80)
    elif pattern == "long":
        buzzer_pin.value(1); time.sleep_ms(400); buzzer_pin.value(0)

def oled_show(i_phase, i_earth, freq, state, rc):
    load_pct = int((i_phase / I_NOMINAL) * 100)
    oled.fill(0)
    oled.text("SMART RELAY V4", 0, 0)
    oled.text("I:{:.0f}A {}%".format(i_phase, load_pct), 0, 14)
    oled.text("Ie:{:.1f}A".format(i_earth), 0, 26)
    oled.text("F:{:.2f}Hz".format(freq), 0, 38)
    oled.text("{} RC:{}".format(state[:9], rc), 0, 52)
    oled.show()

# =============================================================================
# MESURES
# =============================================================================

def read_current():
    samples = sum(ct_adc.read() for _ in range(10)) / 10
    i_sec = (samples / 4095.0) * 5.0
    return round(max(0.0, i_sec * CT_RATIO), 2)

def calc_earth_current(i_phase):
    if i_phase >= I_PICKUP_50:
        return round(i_phase * 0.28, 2)
    elif i_phase >= I_PICKUP_51:
        return round(i_phase * 0.08, 2)
    return round(i_phase * 0.03, 2)

def read_frequency():
    val = freq_adc.read()
    return round(45.0 + (val / 4095.0) * 10.0, 2)

def calc_idmt_time(i_measured, i_pickup, tms):
    if i_measured <= i_pickup:
        return 999.0
    ratio = i_measured / i_pickup
    denom = math.pow(ratio, IEC_ALPHA) - 1.0
    if denom < 0.001:
        return 999.0
    return round(max(0.05, (IEC_K * tms) / denom), 3)

# =============================================================================
# ✅ CONNEXION WIFI — مصحح
# =============================================================================

def connect_wifi():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    if wlan.isconnected():
        print("WiFi déjà connecté | IP:", wlan.ifconfig()[0])
        return True

    wlan.connect(WIFI_SSID, WIFI_PASSWORD)

    oled.fill(0)
    oled.text("WiFi...", 0, 28)
    oled.show()

    print("WiFi connexion", end="")
    timeout = 0
    while not wlan.isconnected() and timeout < 20:
        print(".", end="")
        time.sleep(0.5)
        timeout += 1

    if wlan.isconnected():
        ip = wlan.ifconfig()[0]
        print(" OK | IP:", ip)
        oled.fill(0)
        oled.text("WiFi OK", 0, 20)
        oled.text(ip, 0, 36)
        oled.show()
        time.sleep(1)
        return True
    else:
        print(" ECHEC — vérifier SSID/mot de passe")
        oled.fill(0)
        oled.text("WiFi ECHEC", 0, 28)
        oled.show()
        return False

# =============================================================================
# ✅ CONNEXION MQTT — مصحح (ssl_params complets + retry)
# =============================================================================

def connect_mqtt():
    global mqtt_client

    # ─── المشكلة الأولى كانت هنا: ssl=True بدون ssl_params ───
    # HiveMQ Cloud يحتاج server_hostname + cert_reqs=0
    try:
        client = MQTTClient(
            CLIENT_ID,
            MQTT_BROKER,
            port     = MQTT_PORT,
            user     = MQTT_USER,
            password = MQTT_PASSWORD,
            ssl      = True,
            ssl_params = {
                "server_hostname": MQTT_BROKER,   # ← ضروري لـ TLS SNI
                "cert_reqs"      : 0              # ← يتجاوز فحص الشهادة (Wokwi)
            }
        )
        client.set_callback(on_mqtt_command)
        client.connect()
        client.subscribe(T_CMD)
        mqtt_client = client

        print("MQTT OK | Broker:", MQTT_BROKER)
        log_event("INFO", "MQTT HiveMQ Cloud connecté — protection active")

        oled.fill(0)
        oled.text("MQTT OK", 0, 20)
        oled.text("HiveMQ Cloud", 0, 36)
        oled.show()
        time.sleep(1)
        return True

    except Exception as e:
        print("MQTT ERREUR:", e)
        mqtt_client = None

        oled.fill(0)
        oled.text("MQTT ERREUR", 0, 20)
        oled.show()
        return False

# =============================================================================
# ✅ PUBLICATION MQTT — مصحح (safe publish موحد)
# =============================================================================

def mqtt_publish(topic, payload):
    global mqtt_client
    if not mqtt_client:
        return
    try:
        mqtt_client.publish(topic, ujson.dumps(payload))
    except Exception as e:
        print("MQTT publish error:", e)
        mqtt_client = None   # ← يجبر إعادة الاتصال في الحلقة

def on_mqtt_command(topic, msg):
    global remote_mode, lockout_active
    try:
        cmd    = ujson.loads(msg)
        action = cmd.get("action", "").upper()
        log_event("CMD", "Commande reçue: " + action)

        if action == "TRIP":
            if remote_mode:
                open_breaker("CMD distante TRIP")
            else:
                log_event("CMD", "TRIP refusé — mode LOCAL")

        elif action == "CLOSE":
            if not lockout_active:
                close_breaker("CMD distante CLOSE")
            else:
                log_event("CMD", "CLOSE refusé — LOCKOUT actif")

        elif action == "RESET":
            reset_lockout()

        elif action == "MODE_LOCAL":
            remote_mode = False
            log_event("CMD", "Mode LOCAL activé")

        elif action == "MODE_REMOTE":
            remote_mode = True
            log_event("CMD", "Mode DISTANT activé")

    except Exception as e:
        print("Commande invalide:", e)

# =============================================================================
# COMMANDES DISJONCTEUR
# =============================================================================

def open_breaker(reason, i_phase=0, i_earth=0, freq=0):
    global relay_state
    relay_pin.value(0)
    relay_state = "TRIP"
    set_leds("TRIP")
    buzzer_pulse("double")
    log_event("TRIP", reason, i_phase, i_earth, freq)
    mqtt_publish(T_ALARM, {
        "code" : "TRIP",
        "msg"  : reason,
        "I"    : i_phase,
        "Ie"   : i_earth,
        "F"    : freq,
        "state": relay_state
    })

def close_breaker(reason="Fermeture normale"):
    global relay_state
    relay_pin.value(1)
    relay_state = "NORMAL"
    set_leds("NORMAL")
    buzzer_pin.value(0)
    log_event("RESET", reason)

def reset_lockout():
    global lockout_active, reclose_count
    lockout_active = False
    reclose_count  = 0
    close_breaker("Lockout réinitialisé")
    log_event("CMD", "Reset lockout exécuté")

# =============================================================================
# AUTO-RECLOSER ANSI 79
# =============================================================================

def auto_reclose_sequence(i_phase, i_earth, freq=0.0):
    global reclose_count, lockout_active, relay_state

    if not RECLOSE_ENABLED or lockout_active:
        return

    while reclose_count < MAX_RECLOSE_ATTEMPTS:
        reclose_count += 1
        real_delay = RECLOSE_DELAYS[reclose_count - 1]
        sim_delay  = max(1, int(real_delay * RECLOSE_WOKWI_FACTOR))

        relay_state = "RECLOSE"
        set_leds("RECLOSE")
        buzzer_pulse("single")

        log_event("RECLOSE",
                  "Tentative {}/{} delai {}s".format(
                      reclose_count, MAX_RECLOSE_ATTEMPTS, real_delay),
                  i_phase, i_earth, freq)

        mqtt_publish(T_RECLOSE, {
            "attempt": reclose_count,
            "max"    : MAX_RECLOSE_ATTEMPTS,
            "status" : "WAITING",
            "delay_s": real_delay
        })

        time.sleep(sim_delay)

        relay_pin.value(1)
        log_event("RECLOSE", "Disjoncteur refermé", i_phase, i_earth)

        mqtt_publish(T_RECLOSE, {
            "attempt": reclose_count,
            "max"    : MAX_RECLOSE_ATTEMPTS,
            "status" : "RECLOSED",
            "delay_s": 0
        })

        # Surveillance post-reclose
        fault_persist = False
        for _ in range(3):
            time.sleep(1)
            i_new  = read_current()
            ie_new = calc_earth_current(i_new)
            if i_new >= I_PICKUP_50 or ie_new >= I_PICKUP_50N:
                open_breaker("Defaut persistant reclose {}".format(reclose_count),
                             i_new, ie_new)
                fault_persist = True
                break
            elif i_new >= I_PICKUP_51 or ie_new >= I_PICKUP_51N:
                open_breaker("Surcharge persistante reclose {}".format(reclose_count),
                             i_new, ie_new)
                fault_persist = True
                break

        if not fault_persist:
            relay_state   = "NORMAL"
            reclose_count = 0
            set_leds("NORMAL")
            log_event("RESET", "Reclosing reussi — defaut fugitif", i_phase, i_earth)
            mqtt_publish(T_RECLOSE, {"attempt":0,"max":3,"status":"SUCCESS"})
            return

    # LOCKOUT
    lockout_active = True
    relay_state    = "LOCKOUT"
    relay_pin.value(0)
    set_leds("LOCKOUT")
    buzzer_pulse("long")
    log_event("LOCKOUT",
              "LOCKOUT — {} tentatives echouees".format(MAX_RECLOSE_ATTEMPTS),
              i_phase, i_earth)
    mqtt_publish(T_RECLOSE, {
        "attempt": reclose_count,
        "max"    : MAX_RECLOSE_ATTEMPTS,
        "status" : "LOCKOUT"
    })

# =============================================================================
# PROTECTION ANSI 81 — Fréquence
# =============================================================================

def protection_81(freq, i_phase, i_earth):
    global relay_state, freq_uf_warn_active, freq_uf_warn_start
    global freq_of_warn_active, freq_of_warn_start

    # 81U TRIP
    if freq < FREQ_UF_TRIP:
        if relay_state not in ("TRIP", "LOCKOUT"):
            open_breaker("ANSI 81U TRIP {:.2f}Hz".format(freq), i_phase, i_earth, freq)
            relay_state = "FREQ_TRIP"
        freq_uf_warn_active = False
        return

    # 81U WARNING
    if freq < FREQ_UF_WARNING:
        if not freq_uf_warn_active:
            freq_uf_warn_active = True
            freq_uf_warn_start  = time.time()
            relay_state = "FREQ_UF_WARN"
            log_event("FREQ_WARN", "81U Warning {:.2f}Hz".format(freq),
                      i_phase, i_earth, freq)
            buzzer_pulse("single")
            mqtt_publish(T_ALARM, {"code":"W81U","msg":"Under Freq Warning","F":freq})
    else:
        if freq_uf_warn_active:
            freq_uf_warn_active = False
            if relay_state == "FREQ_UF_WARN":
                relay_state = "NORMAL"

    # 81O TRIP
    if freq > FREQ_OF_TRIP:
        if relay_state not in ("TRIP", "LOCKOUT"):
            open_breaker("ANSI 81O TRIP {:.2f}Hz".format(freq), i_phase, i_earth, freq)
            relay_state = "FREQ_TRIP"
        freq_of_warn_active = False
        return

    # 81O WARNING
    if freq > FREQ_OF_WARNING:
        if not freq_of_warn_active:
            freq_of_warn_active = True
            freq_of_warn_start  = time.time()
            relay_state = "FREQ_OF_WARN"
            log_event("FREQ_WARN", "81O Warning {:.2f}Hz".format(freq),
                      i_phase, i_earth, freq)
            buzzer_pulse("single")
            mqtt_publish(T_ALARM, {"code":"W81O","msg":"Over Freq Warning","F":freq})
    else:
        if freq_of_warn_active:
            freq_of_warn_active = False
            if relay_state == "FREQ_OF_WARN":
                relay_state = "NORMAL"

# =============================================================================
# PROTECTION LOGIC PRINCIPALE
# =============================================================================

def protection_logic(i_phase, i_earth, freq):
    global relay_state, pickup_51_active, pickup_51_start
    global pickup_51N_active, pickup_51N_start

    if lockout_active:
        relay_state = "LOCKOUT"
        return

    # Priorité 1 — ANSI 50N instantané
    if i_earth >= I_PICKUP_50N:
        open_breaker("ANSI 50N Earth Fault Inst.", i_phase, i_earth, freq)
        auto_reclose_sequence(i_phase, i_earth, freq)
        pickup_51N_active = False
        return

    # Priorité 2 — ANSI 50 instantané
    if i_phase >= I_PICKUP_50:
        open_breaker("ANSI 50 Short Circuit Inst.", i_phase, i_earth, freq)
        auto_reclose_sequence(i_phase, i_earth, freq)
        pickup_51_active = False
        return

    # Priorité 3 — ANSI 51N IDMT
    if i_earth >= I_PICKUP_51N:
        if not pickup_51N_active:
            pickup_51N_active = True
            pickup_51N_start  = time.time()
            relay_state = "PICKUP"
            t = calc_idmt_time(i_earth, I_PICKUP_51N, TMS_51N)
            log_event("PICKUP", "51N Pickup Ie={:.1f}A t={:.2f}s".format(i_earth, t),
                      i_phase, i_earth, freq)
            mqtt_publish(T_ALARM, {"code":"PICKUP_51N","msg":"Earth Fault Pickup",
                                   "I":i_phase,"Ie":i_earth})
            buzzer_pulse("single")
        else:
            t_req  = calc_idmt_time(i_earth, I_PICKUP_51N, TMS_51N)
            t_elap = time.time() - pickup_51N_start
            relay_state = "TIMING"
            if t_elap >= t_req:
                pickup_51N_active = False
                open_breaker("ANSI 51N IDMT t={:.2f}s".format(t_req),
                             i_phase, i_earth, freq)
                mqtt_publish(T_ALARM, {"code":"F51N","msg":"Earth Fault IDMT",
                                       "I":i_phase,"Ie":i_earth})
                auto_reclose_sequence(i_phase, i_earth, freq)
    else:
        if pickup_51N_active:
            pickup_51N_active = False
            log_event("RESET", "51N Pickup annule", i_phase, i_earth, freq)

    # Priorité 4 — ANSI 51 IDMT
    if i_phase >= I_PICKUP_51:
        if not pickup_51_active:
            pickup_51_active = True
            pickup_51_start  = time.time()
            relay_state = "PICKUP"
            t = calc_idmt_time(i_phase, I_PICKUP_51, TMS_51)
            log_event("PICKUP", "51 Pickup I={:.1f}A t={:.2f}s".format(i_phase, t),
                      i_phase, i_earth, freq)
            mqtt_publish(T_ALARM, {"code":"PICKUP_51","msg":"Overcurrent Pickup",
                                   "I":i_phase,"Ie":i_earth})
            buzzer_pulse("single")
        else:
            t_req  = calc_idmt_time(i_phase, I_PICKUP_51, TMS_51)
            t_elap = time.time() - pickup_51_start
            relay_state = "TIMING"
            if t_elap >= t_req:
                pickup_51_active = False
                open_breaker("ANSI 51 IDMT t={:.2f}s".format(t_req),
                             i_phase, i_earth, freq)
                mqtt_publish(T_ALARM, {"code":"F51","msg":"Overcurrent IDMT",
                                       "I":i_phase,"Ie":i_earth})
                auto_reclose_sequence(i_phase, i_earth, freq)
    else:
        if pickup_51_active:
            pickup_51_active = False
            log_event("RESET", "51 Pickup annule", i_phase, i_earth, freq)
        if relay_state not in ("TRIP","RECLOSE","LOCKOUT","FREQ_TRIP",
                               "FREQ_UF_WARN","FREQ_OF_WARN"):
            relay_state = "NORMAL"

    # Priorité 5 — ANSI 81 (uniquement si pas de défaut courant)
    if relay_state in ("NORMAL", "FREQ_UF_WARN", "FREQ_OF_WARN"):
        protection_81(freq, i_phase, i_earth)

# =============================================================================
# DÉMARRAGE
# =============================================================================

print("=" * 50)
print("SMART RELAY V4 | ANSI 50/51/50N/51N/79/81")
print("MT 30kV | CT 200/5 | IEC 60255-3")
print("=" * 50)

connect_wifi()
connect_mqtt()

log_event("INFO", "Systeme demarre — 6 protections ANSI actives")

# =============================================================================
# BOUCLE PRINCIPALE
# =============================================================================

loop_count = 0

while True:
    try:
        # ── Commandes MQTT entrantes ──────────────────────────
        if mqtt_client:
            try:
                mqtt_client.check_msg()
            except Exception:
                mqtt_client = None   # يجبر reconnexion

        # ── Mesures ──────────────────────────────────────────
        i_phase = read_current()
        i_earth = calc_earth_current(i_phase)
        freq    = read_frequency()

        # ── Protection ───────────────────────────────────────
        if not lockout_active:
            protection_logic(i_phase, i_earth, freq)

        # ── LEDs + OLED ──────────────────────────────────────
        set_leds(relay_state)
        oled_show(i_phase, i_earth, freq, relay_state, reclose_count)

        # ── MQTT Publish ─────────────────────────────────────
        mqtt_publish(T_MEASURE, {
            "I_phase"  : i_phase,
            "I_earth"  : i_earth,
            "frequency": freq,
            "load_pct" : round((i_phase / I_NOMINAL) * 100, 1),
            "state"    : relay_state,
            "breaker"  : bool(relay_pin.value()),
            "reclose"  : reclose_count,
            "lockout"  : lockout_active
        })

        if loop_count % 5 == 0:
            mqtt_publish(T_STATE, {
                "state"         : relay_state,
                "breaker_closed": bool(relay_pin.value()),
                "reclose_count" : reclose_count,
                "lockout"       : lockout_active,
                "remote_mode"   : remote_mode,
                "pickup_51_on"  : pickup_51_active,
                "pickup_51N_on" : pickup_51N_active
            })

        loop_count += 1

        print("I={:.0f}A Ie={:.0f}A F={:.2f}Hz | {} | RC={} LO={}".format(
            i_phase, i_earth, freq, relay_state, reclose_count, lockout_active))

        time.sleep(1.5)

    except Exception as e:
        print("Loop Error:", e)
        time.sleep(2)

        # ── Reconnexion automatique MQTT ──────────────────────
        if mqtt_client is None:
            print("Tentative reconnexion MQTT...")
            connect_mqtt()