import gradio as gr
import numpy as np
import joblib
from pytorch_tabnet.tab_model import TabNetClassifier

# Load model and preprocessing tools
model = TabNetClassifier()
model.load_model("tabnet_model.zip")
scaler = joblib.load("scaler.save")
encoder = joblib.load("encoder.save")

# Features used in the model
features = [f"{trait}{i}" for trait in ["EXT", "EST", "AGR", "CSN", "OPN"] for i in range(1, 11)]

def predict_personality(*inputs):
    X = np.array(inputs).reshape(1, -1).astype(np.float32)
    X_scaled = scaler.transform(X)
    y_pred = model.predict(X_scaled)
    label = encoder.inverse_transform(y_pred)[0]
    return f"Predicted Personality Type: {label}"

# Create Gradio interface
inputs = [gr.Slider(1, 5, step=0.1, label=f) for f in features]

demo = gr.Interface(
    fn=predict_personality,
    inputs=inputs,
    outputs=gr.Text(label="Personality Prediction"),
    title="Personality Type Classifier (Introvert vs. Extrovert)",
    description="This model predicts if a person is Introvert or Extrovert based on their IPIP-FFM scores."
)

demo.launch()