import gradio as gr
import pandas as pd
import skops.io as sio
from huggingface_hub import hf_hub_download
import spaces  # Import spaces module

# 1. Download model
model_path = hf_hub_download(
    repo_id="iqranaz/Drug-Classification", 
    filename="drug_pipeline.skops"
)

# 2. Load model
untrusted_types = sio.get_untrusted_types(file=model_path)
pipe = sio.load(model_path, trusted=untrusted_types)

# 3. Add the @spaces.GPU decorator to satisfy ZeroGPU requirement
@spaces.GPU
def predict_drug(age, sex, bp, cholesterol, na_to_k):
    df = pd.DataFrame([{
        "Age": age,
        "Sex": sex,
        "BP": bp,
        "Cholesterol": cholesterol,
        "Na_to_K": na_to_k
    }])
    pred = pipe.predict(df)[0]
    return f"Predicted Drug: {pred}"

# 4. Interface definition
demo = gr.Interface(
    fn=predict_drug,
    inputs=[
        gr.Number(label="Age", value=30),
        gr.Dropdown(choices=["M", "F"], label="Sex"),
        gr.Dropdown(choices=["HIGH", "NORMAL", "LOW"], label="BP"),
        gr.Dropdown(choices=["HIGH", "NORMAL"], label="Cholesterol"),
        gr.Number(label="Na_to_K Ratio", value=15.0),
    ],
    outputs="text",
    title="Drug Classification Model",
    description="Predict drug type based on patient health indicators."
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)