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

# Define the repository ID and the model filename
# Make sure this matches your deployed model on Hugging Face
repo_id = "farooqhasanDA/logistic-regression-sklearn-model"
model_filename = "models/logistic_regression_model.joblib"

# Global variable to store the loaded model
loaded_model = None

def predict_proba_wrapper(feature1_num, feature2_text, feature3_dropdown, feature4_checkbox):
    global loaded_model

    # Download and load the model if not already loaded
    if loaded_model is None:
        print("Downloading and loading model for the first time...")
        try:
            model_path_local = hf_hub_download(repo_id=repo_id, filename=model_filename)
            loaded_model = joblib.load(model_path_local)
            print("Model loaded successfully.")
        except Exception as e:
            return f"Error loading model: {e}"

    try:
        # Convert text to a numerical value (e.g., length, or some simple encoding)
        f2_val = len(feature2_text) if feature2_text else 0
        # Convert dropdown choice to a numerical value
        f3_val = {'Option A': 0, 'Option B': 1, 'Option C': 2}.get(feature3_dropdown, 0)
        # Convert checkbox (boolean) to 0 or 1
        f4_val = 1 if feature4_checkbox else 0

        input_array = np.array([[float(feature1_num), float(f2_val), float(f3_val), float(f4_val)]])

        # Make a prediction and get prediction probabilities
        prediction = loaded_model.predict(input_array)
        probabilities = loaded_model.predict_proba(input_array)

        # Return formatted string
        return (f"Prediction: Class {prediction[0]}, "
                f"Probabilities: Class 0 = {probabilities[0][0]:.4f}, Class 1 = {probabilities[0][1]:.4f}")
    except ValueError:
        return "Error: Please ensure Feature 1 is a valid number."
    except Exception as e:
        return f"Prediction error: {e}"


# Create gradio input components with different types
inputs = [
    gr.Number(label="Numerical Feature 1 (e.g., 0.5)", value=0.5),
    gr.Textbox(label="Text Feature 2 (e.g., 'Hello world')", value="Sample text"),
    gr.Dropdown(choices=["Option A", "Option B", "Option C"], label="Categorical Feature 3"),
    gr.Checkbox(label="Boolean Feature 4 (checked/unchecked)", value=True)
]

# Create a gradio output component
output = gr.Textbox(label="Prediction Result")

# Instantiate gr.Interface
iface = gr.Interface(fn=predict_proba_wrapper, inputs=inputs, outputs=output,
                     title="Logistic Regression Model Predictor (Hosted on HF Space)",
                     description="Enter different types of features to get a binary classification prediction and probabilities. 
(Note: Text/Dropdown/Checkbox inputs are converted to numbers for this demo model.)")

# Launch the Gradio interface
iface.launch(debug=True)