import gradio as gr
import joblib
import numpy as np
from huggingface_hub import hf_hub_download

# Download models from Hugging Face Hub
svr_path = hf_hub_download(repo_id="iamomtiwari/Nutrition-regression-models", filename="svr_model.pkl")
lr_path = hf_hub_download(repo_id="iamomtiwari/Nutrition-regression-models", filename="linear_regression.pkl")

# Load models
svr_model = joblib.load(svr_path)
linear_reg = joblib.load(lr_path)

# Selected 10 important features
features = ['Caloric Value', 'Fat', 'Saturated Fats', 'Carbohydrates', 'Sugars',
            'Protein', 'Cholesterol', 'Sodium', 'Calcium', 'Iron']

# Define prediction function
def predict(model_name, *inputs):
    input_data = np.array([inputs]).reshape(1, -1)
    
    if model_name == "SVR":
        prediction = svr_model.predict(input_data)[0]
    else:
        prediction = linear_reg.predict(input_data)[0]
    
    return round(prediction, 4)

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Nutrition Density Prediction")
    
    model_choice = gr.Radio(["SVR", "Linear Regression"], label="Select Model")
    input_widgets = [gr.Slider(minimum=0, maximum=100, step=0.1, label=feature) for feature in features]
    predict_button = gr.Button("Predict")
    clear_button = gr.Button("Clear")
    output_label = gr.Textbox(label="Prediction")
    
    predict_button.click(predict, inputs=[model_choice] + input_widgets, outputs=output_label)
    
    # Reset sliders to their default value (0) on "Clear"
    def reset_sliders():
        return [0] * len(features)
    
    clear_button.click(reset_sliders, inputs=[], outputs=input_widgets)

# Run the app
demo.launch()