import gradio as gr
import pickle
import numpy as np
import pandas as pd

# Load models
with open('linear_regression_model.pkl', 'rb') as f:
    lin_reg = pickle.load(f)
with open('logistic_regression_model.pkl', 'rb') as f:
    log_reg = pickle.load(f)

# Define feature names (must match training)
FEATURES = ['Glucose', 'BMI', 'Age']

def predict_model(model_type, glucose, bmi, age):
    X = np.array([[glucose, bmi, age]])
    if model_type == 'Linear Regression':
        pred = lin_reg.predict(X)[0]
        return f"Linear Regression Prediction: {pred:.2f}"
    else:
        pred = log_reg.predict(X)[0]
        prob = log_reg.predict_proba(X)[0][1]
        result = 'Diabetic' if pred == 1 else 'Not Diabetic'
        return f"Logistic Regression Prediction: {result} (Probability: {prob:.2f})"

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Diabetes Prediction App\nSelect model and enter values to predict.")
    model_type = gr.Radio(["Linear Regression", "Logistic Regression"], label="Choose Model")
    glucose = gr.Number(label="Glucose")
    bmi = gr.Number(label="BMI")
    age = gr.Number(label="Age")
    predict_btn = gr.Button("Predict")
    output = gr.Textbox(label="Prediction Result")

    def run_predict(model_type, glucose, bmi, age):
        return predict_model(model_type, glucose, bmi, age)

    predict_btn.click(run_predict, inputs=[model_type, glucose, bmi, age], outputs=output)

demo.launch()

# Hugging Face Spaces requirements:
# 1. Add a requirements.txt file with gradio and scikit-learn
# 2. Make sure app.py is the entry point
# 3. Place your pickle files in the same directory
# Example requirements.txt:
# gradio
# scikit-learn
# numpy
# pandas