import gradio as gr
import numpy as np
import tensorflow as tf

# Load the trained model
model = tf.keras.models.load_model("sleep_cognition_model.h5")  

# Define prediction function
def predict(TST, SE, WASO, REM_Sleep, Deep_Sleep, Number_of_Awakenings):
    input_data = np.array([[TST, SE, WASO, REM_Sleep, Deep_Sleep, Number_of_Awakenings]])
    prediction = model.predict(input_data)
    return {
        "Reaction Time": float(prediction[0][0]),
        "Memory Recall Accuracy": float(prediction[0][1]),
        "Attention Score": float(prediction[0][2]),
        "Executive Function Score": float(prediction[0][3]),
        "Mental Fatigue Index": float(prediction[0][4]),
    }

# Define Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Number(label="Total Sleep Time (TST)"),
        gr.Number(label="Sleep Efficiency (SE)"),
        gr.Number(label="Wake After Sleep Onset (WASO)"),
        gr.Number(label="REM Sleep (%)"),
        gr.Number(label="Deep Sleep (%)"),
        gr.Number(label="Number of Awakenings"),
    ],
    outputs="json",
    title="Sleep & Cognitive Function Predictor",
    description="Enter sleep parameters to predict cognitive function scores.",
)

# Run the app
iface.launch()