import gradio as gr
import joblib
from  tensorflow.keras.models import load_model

relationship_mapping = {
    'Own-child': 3,
    'Husband': 0,
    'Not-in-family': 1,
    'Unmarried': 4,
    'Wife': 5,
    'Other-relative': 2
}

marital_status_mapping = {
    'Never-married': 4,
    'Married-civ-spouse': 2,
    'Widowed': 6,
    'Divorced': 0,
    'Separated': 5,
    'Married-spouse-absent': 3,
    'Married-AF-spouse': 1
}

model_dl = load_model('model_dl_importance.h5')

scaler = joblib.load('scaler.pkl')


def predict_income_dl(relationship, fnlwgt, age, marital_status, model):
    rel = relationship_mapping.get(relationship, -1)
    mar = marital_status_mapping.get(marital_status, -1)
    input_data = [[rel, fnlwgt, age, mar]]
    input_data = scaler.transform(input_data)
    prediction = model_dl.predict(input_data, verbose=0)
    prediction = (prediction > 0.5)
    if prediction:
      return ">50"
    else:
      return "<50"
  
  
demo = gr.Interface(
    fn=predict_income_dl,
    inputs=[
        gr.Dropdown(choices=list(relationship_mapping.keys()), label="Relationship"),
        gr.Text(label="Final Weight"),
        gr.Text(label="Age"),
        gr.Dropdown(choices=list(marital_status_mapping.keys()), label="Marital Status"),
        # input berupa radio button
        gr.Radio(
            choices=["Random Forest", "Neural Network"],
            value="Random Forest",
            label="Model",
            info="Select model to use"
        )
    ],
    outputs="text")
# launch aplikasi
demo.launch(share=True)