# UIFlow 2 (MicroPython): stand-alone fruit judge on the M5Stack, no laptop needed.
#
# Files on the M5Stack (upload with the UIFlow 2 file manager, same folder as this file):
#     mb_infer.py         mushroom body inference (from code/firmware/)
#     enose_model.json    trained model (from experiments/export_for_m5stack.py)
#
# Press button A -> purge 120 s -> put the fruit in, press A -> sample 240 s
# -> the screen shows fresh / ripe / overripe / rotten and the confidence.
# Sensor reading is the same as enose_logger.py; the heater steps, BME688 period and phase
# lengths are read from enose_model.json so they always match the training data.
import M5
from M5 import *
from hardware import *
from unit import ENVPROUnit, TVOCUnit, CO2Unit
import time
import mb_infer

MODEL_PATH = "enose_model.json"   # adjust if the file manager puts it in another folder
BOARD = "core2"                   # "core2" or "cores3"
STAGE_NAMES = {"fresh": "FRESH", "ripe": "RIPE", "overripe": "OVERRIPE", "rotten": "ROTTEN"}

M5.begin()
Widgets.fillScreen(0x111111)
big = Widgets.Label("E-nose", 8, 20, 1.0, 0x5EEAD4, 0x111111, Widgets.FONTS.DejaVu40)
small = Widgets.Label("loading model...", 8, 90, 1.0, 0xFFFFFF, 0x111111, Widgets.FONTS.DejaVu18)
hint = Widgets.Label("", 8, 200, 1.0, 0x98A3C0, 0x111111, Widgets.FONTS.DejaVu18)

if BOARD == "cores3":
    i2c0 = I2C(0, scl=Pin(1), sda=Pin(2), freq=100000)
else:
    i2c0 = I2C(0, scl=Pin(33), sda=Pin(32), freq=100000)
env = ENVPROUnit(i2c0)
tvoc = TVOCUnit(i2c0)
co2 = CO2Unit(i2c0)
co2.set_start_periodic_measurement()
model = mb_infer.MushroomBody(MODEL_PATH)
small.setText("model: " + "/".join(model.classes))
# Heater steps and timing come from the model file, i.e. from how the training data was recorded.
HEATER_STEPS_C = model.schedule["heater_steps_c"]
HEATER_MS = model.schedule["heater_ms"]
BME_PERIOD_S = model.schedule["bme_period_s"]
PURGE_S = model.schedule["purge_s"]
SAMPLE_S = model.schedule["sample_s"]
state = {"step": 0, "tick": 0, "co2": 0}
env.set_gas_heater(HEATER_STEPS_C[0], HEATER_MS)


def beep():
    try:
        Speaker.tone(2000, 150)
    except Exception:
        pass


def read_once(second):
    """One reading per second, same schedule as enose_logger.py."""
    _, _ = tvoc.measure_iaq()
    h2_raw, ethanol_raw = tvoc.measure_raw()
    if co2.is_data_ready():
        state["co2"] = co2.co2
    row = {"second": second, "ethanol": ethanol_raw, "h2": h2_raw, "co2": state["co2"]}
    for step in range(len(HEATER_STEPS_C)):
        row["gas{}".format(step)] = None
    if state["tick"] % BME_PERIOD_S == 0:
        row["gas{}".format(state["step"])] = env.get_gas_resistance() * 1000
        state["step"] = (state["step"] + 1) % len(HEATER_STEPS_C)
        env.set_gas_heater(HEATER_STEPS_C[state["step"]], HEATER_MS)
    state["tick"] += 1
    return row


def run_phase(name, seconds):
    rows = []
    next_ms = time.ticks_add(time.ticks_ms(), 1000)
    for second in range(seconds):
        M5.update()
        rows.append(read_once(second))
        big.setText(name)
        small.setText("{} s left".format(seconds - second))
        wait = time.ticks_diff(next_ms, time.ticks_ms())
        if wait > 0:
            time.sleep_ms(wait)
        next_ms = time.ticks_add(next_ms, 1000)
    return rows


def wait_button(message):
    hint.setText(message)
    beep()
    while True:
        M5.update()
        if BtnA.wasPressed():
            hint.setText("")
            return
        time.sleep_ms(20)


while True:
    wait_button("empty, aired box -> A: start")
    purge = run_phase("PURGE", PURGE_S)
    wait_button("put fruit in, press A")
    sample = run_phase("SAMPLE", SAMPLE_S)
    values = mb_infer.features(purge, sample, len(HEATER_STEPS_C))
    stage, confidence, _ = model.predict(values)
    big.setText(STAGE_NAMES.get(stage, stage))
    small.setText("confidence {:.0%}".format(confidence))
    print("RESULT,{},{:.3f},{}".format(stage, confidence, ",".join(str(v) for v in values)))
    beep()