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

# Load model, scaler, and encoder
model = TabNetClassifier()
model.load_model('tabnet_model.zip')  # Must be .zip file
scaler = joblib.load('scaler.save')
encoder = joblib.load('encoder.save')

# Full form trait mapping
trait_prefixes = {
    'Extraversion': 'EXT',
    'Emotional Stability': 'EST',
    'Agreeableness': 'AGR',
    'Conscientiousness': 'CSN',
    'Openness': 'OPN'
}

# Create full feature names with full form labels
feature_labels = []
for trait in trait_prefixes:
    for i in range(10):
        feature_labels.append(f"{trait} Q{i+1}")

# Inference function
def predict_personality(*inputs):
    input_array = np.array(inputs).reshape(1, -1)
    scaled_input = scaler.transform(input_array)
    pred = model.predict(scaled_input)
    personality = encoder.inverse_transform(pred)[0]
    return f"Predicted Personality: **{personality}**"

# Gradio UI: 50 sliders with full trait names
inputs = [gr.Slider(1.0, 5.0, value=3.0, label=label) for label in feature_labels]
output = gr.Textbox(label="Prediction")

demo = gr.Interface(
    fn=predict_personality,
    inputs=inputs,
    outputs=output,
    title="Personality Type Classifier (Introvert vs. Extrovert)",
    description="Provide scores (1–5) for 50 questions from the IPIP-FFM questionnaire. The model will predict whether the person is an Introvert or an Extrovert."
)

demo.launch()