import tflite_runtime.interpreter as tflite
import numpy as np
import as7265x
import smbus
import joblib
import RPi.GPIO as GPIO
from luma.core.interface.serial import i2c
from luma.core.render import canvas
from luma.oled.device import ssd1306
from time import sleep

# ======================== Load TF Model ========================
interpreter = tflite.Interpreter(model_path="model.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

print("Input details:", input_details)
print("Output details:", output_details)

# ======================== Load Scalers ========================
scaler_X = joblib.load('scaler_X.pkl')
scaler_y = joblib.load('scaler_y.pkl')

# ======================== GPIO Button Setup ========================
BUTTON_PIN = 22  # Gunakan GPIO 22
GPIO.setmode(GPIO.BCM)
GPIO.setup(BUTTON_PIN, GPIO.IN, pull_up_down=GPIO.PUD_DOWN)

# ======================== Create Sensor Instance ========================
i2c_bus = smbus.SMBus(1)
sensor = as7265x.AS7265X(i2c_bus)

# ======================== OLED Display Setup ========================
serial = i2c(port=1, address=0x3C)
device = ssd1306(serial, rotate=0)

# ======================== Scan Function ========================
def scan():
    sensor.begin()
    sensor.enableBulb(as7265x.LED_WHITE)
    sensor.enableBulb(as7265x.LED_IR)
    sensor.enableBulb(as7265x.LED_UV)
    sensor.setIntegrationCycles(1)
    sensor.takeMeasurements()

    data = [
        sensor.getCalibratedA(), sensor.getCalibratedB(), sensor.getCalibratedC(),
        sensor.getCalibratedD(), sensor.getCalibratedE(), sensor.getCalibratedF(),
        sensor.getCalibratedG(), sensor.getCalibratedH(), sensor.getCalibratedR(),
        sensor.getCalibratedI(), sensor.getCalibratedS(), sensor.getCalibratedJ(),
        sensor.getCalibratedT(), sensor.getCalibratedU(), sensor.getCalibratedV(),
        sensor.getCalibratedW(), sensor.getCalibratedK(), sensor.getCalibratedL()
    ]

    sensor.disableBulb(as7265x.LED_WHITE)
    sensor.disableBulb(as7265x.LED_IR)
    sensor.disableBulb(as7265x.LED_UV)

    data = np.array(data).reshape(1, -1)
    return data

# ======================== Predict Function ========================
def predict(data):
    scaled_data = scaler_X.transform(data)
    scaled_data = scaled_data.reshape((1, 18, 1))

    interpreter.set_tensor(input_details[0]['index'], scaled_data)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]['index'])
    output_original = scaler_y.inverse_transform(output_data)

    print("Model output (scaled):", output_data)
    print("Model output (original):", output_original)

    with canvas(device) as draw:
        draw.rectangle(device.bounding_box, outline="white", fill="black")
        draw.text((5, 20), "Prediction Result:", fill="white")
        draw.text((5, 40), str(np.round(output_original[0], 2)), fill="white")

# ======================== Main Loop ========================
print("Model berhasil diload")
print("Tekan tombol untuk memulai prediksi...")

try:
    while True:
        if GPIO.input(BUTTON_PIN) == GPIO.HIGH:
            print("Tombol ditekan. Memulai scan dan prediksi...")
            data = scan()
            predict(data)
            sleep(1.5)  # Debounce delay
except KeyboardInterrupt:
    print("Program dihentikan oleh pengguna.")
finally:
    GPIO.cleanup()
    print("GPIO dibersihkan.")