import machine, time
from machine import Pin, ADC, I2C
from lcd_api import LcdApi
from i2c_lcd import I2cLcd

# Pin setup
SAFE_LED = Pin(18, Pin.OUT)
DANGER_LED = Pin(5, Pin.OUT)
RELAY = Pin(23, Pin.OUT)
BUZZER = Pin(19, Pin.OUT)
pot = ADC(Pin(34))
pot.atten(ADC.ATTN_11DB)

# LCD Setup (20x4)
i2c = I2C(0, scl=Pin(22), sda=Pin(21), freq=400000)
lcd = I2cLcd(i2c, 0x27, 4, 20)

# ---- AI Predictor Variables ----
baseline_samples = []
baseline_mean = 0
baseline_var = 0
trained = False

def train_ai():
    global baseline_mean, baseline_var, trained
    lcd.clear()
    lcd.putstr("Training AI...")
    time.sleep(1)

    for i in range(20):  # Collect 20 safe samples
        val = pot.read()
        baseline_samples.append(val)
        time.sleep(0.1)
    
    baseline_mean = sum(baseline_samples) / len(baseline_samples)
    baseline_var = sum((x - baseline_mean)**2 for x in baseline_samples) / len(baseline_samples)
    trained = True
    
    lcd.clear()
    lcd.putstr("AI Training Done")

def ai_predict(value):
    if not trained:
        return "Not trained"
    deviation = abs(value - baseline_mean)
    if deviation > 2 * (baseline_var**0.5):
        return "RISK"
    return "SAFE"

# ---- Main Loop ----
train_ai()

while True:
    sensor_value = pot.read()
    prediction = ai_predict(sensor_value)
    
    lcd.clear()
    lcd.move_to(0,0)
    lcd.putstr("AI Prediction Mode")
    lcd.move_to(0,1)
    lcd.putstr("Current: {}".format(sensor_value))
    lcd.move_to(0,2)
    lcd.putstr("Prediction: {}".format(prediction))
    
    if prediction == "RISK":
        SAFE_LED.value(0)
        DANGER_LED.value(1)
        RELAY.value(0)  # cut off power
        BUZZER.value(1)
        lcd.move_to(0,3)
        lcd.putstr("ALERT: CHECK PHONE")
        print("MOBILE ALERT: Insulation failure risk detected!")
    else:
        SAFE_LED.value(1)
        DANGER_LED.value(0)
        RELAY.value(1)
        BUZZER.value(0)
        lcd.move_to(0,3)
        lcd.putstr("System Stable")
    
    time.sleep(1)