import gradio as gr
import joblib
import pandas as pd
import json

# Load the model, scaler, and mapping
scaler = joblib.load("scaler.joblib")
gmm = joblib.load("gmm_model.joblib")

with open("cluster_fatigue_map.json") as f:
    cluster_fatigue_map = json.load(f)

# Define expected input features
feature_cols = [
    'AVRR', 'SDNN', 'RMSSD', 'PNN50', 'Coefficient_of_Variation',
    'Age', 'Weight', 'Height'
]

# Inference function — accepts Python list of values
def predict_fatigue(
    AVRR, SDNN, RMSSD, PNN50, Coefficient_of_Variation,
    Age, Weight, Height
):
    try:
        input_dict = {
            'AVRR': AVRR,
            'SDNN': SDNN,
            'RMSSD': RMSSD,
            'PNN50': PNN50,
            'Coefficient_of_Variation': Coefficient_of_Variation,
            'Age': Age,
            'Weight': Weight,
            'Height': Height
        }

        df = pd.DataFrame([input_dict])[feature_cols]
        scaled = scaler.transform(df)
        cluster = gmm.predict(scaled)[0]
        fatigue_level = cluster_fatigue_map[str(cluster)]

        return {
            "cluster": int(cluster),
            "fatigue_level": fatigue_level
        }

    except Exception as e:
        return {"error": str(e)}

# Gradio app in REST-friendly format
with gr.Blocks() as demo:
    gr.Interface(
        fn=predict_fatigue,
        inputs=[
            gr.Number(label='AVRR'),
            gr.Number(label='SDNN'),
            gr.Number(label='RMSSD'),
            gr.Number(label='PNN50'),
            gr.Number(label='Coefficient of Variation'),
            gr.Number(label='Age'),
            gr.Number(label='Weight'),
            gr.Number(label='Height'),
        ],
        outputs="json",
        live=False
    )

demo.launch()