import gradio as gr
from collections import OrderedDict
import numpy as np
from huggingface_hub import from_pretrained_keras

def predict(
    age,workclass,fnlwgt,education,education_num,marital_status,occupation,relationship,
    race,gender,capital_gain,capital_loss,hours_per_week,native_country
):
    user_data = {}
    user_data['age'] = np.array([age],dtype=np.float32)
    user_data['workclass'] = np.array([f'{workclass}'],dtype="object")
    user_data['fnlwgt'] = np.array([fnlwgt],dtype=np.float32)
    user_data['education'] = np.array([f'{education}'],dtype="object")
    user_data['education_num'] = np.array([education_num],dtype=np.float32)
    user_data['marital_status'] = np.array([f'{marital_status}'],dtype="object")
    user_data['occupation'] = np.array([f'{occupation}'],dtype="object")
    user_data['relationship'] =  np.array([f'{relationship}'],dtype="object")
    user_data['race'] = np.array([f'{race}'],dtype="object")
    user_data['gender'] = np.array([f'{gender}'],dtype="object")
    user_data['capital_gain'] = np.array([capital_gain],dtype=np.float32)
    user_data['capital_loss'] = np.array([capital_loss],dtype=np.float32)
    user_data['hours_per_week'] = np.array([hours_per_week],dtype=np.float32)
    user_data['native_country'] = np.array([f'{native_country}'],dtype="object")
    test_user_data = OrderedDict(user_data)
    model = from_pretrained_keras("keras-io/neural-decision-forest")
    pred = model.predict(test_user_data)
    pred = np.argmax(pred,axis=1)
    return f"Outcome: {pred}"
    
work_class_list = [' Self-emp-not-inc', ' Private', ' State-gov', ' Federal-gov',' Local-gov', ' ?', ' Self-emp-inc', ' Without-pay',' Never-worked']
education_list = [' Bachelors', ' HS-grad', ' 11th', ' Masters', ' 9th',' Some-college', ' Assoc-acdm', ' Assoc-voc', ' 7th-8th',' Doctorate', ' Prof-school', ' 5th-6th', ' 10th', ' 1st-4th',' Preschool', ' 12th']
martial_list = [' Married-civ-spouse',' Divorced',' Married-spouse-absent',' Never-married',' Separated',' Married-AF-spouse',' Widowed']
race_list = [' White',' Black',' Asian-Pac-Islander',' Amer-Indian-Eskimo',' Other']
relation_list = [' Husband',' Not-in-family',' Wife',' Own-child',' Unmarried',' Other-relative']
occupation_list = [' Exec-managerial',' Handlers-cleaners',' Prof-specialty',' Other-service',' Adm-clerical',' Sales',' Craft-repair',' Transport-moving',' Farming-fishing',' Machine-op-inspct',' Tech-support',' ?',' Protective-serv',' Armed-Forces',' Priv-house-serv']
countries = [' United-States',' Cuba',' Jamaica',' India',' Mexico',' South',' Puerto-Rico',' Honduras',' England',' Canada',' Germany',' Iran',' Philippines',' Italy',' Poland',' Columbia',' Cambodia',' Thailand',' Ecuador',' Laos',' Taiwan',' Haiti',' Portugal',' Dominican-Republic',' El-Salvador',' France',' Guatemala',' China',' Japan',' Yugoslavia',' Peru',' Outlying-US(Guam-USVI-etc)',' Scotland',' Trinadad&Tobago',' Greece',' Nicaragua',' Vietnam',' Hong',' Ireland',' Hungary',' Holand-Netherlands']

title = "Deep Neural Decision Forest"
description = "This example uses the United States Census Income Dataset provided by the UC Irvine Machine Learning Repository. The task is binary classification to predict whether a person is likely to be making over USD 50,000 a year."
article = """<p style='text-align: center'>
        <a href='https://keras.io/examples/structured_data/deep_neural_decision_forests/' target='_blank'>Keras Example given by Khalid Salama</a>
        <br/>
        <a href="https://huggingface.co/lucifertrj">Space by @lucifertrj</a>
    </p>
    """

demo = gr.Interface(
    predict,
    [
        gr.Slider(12, 85, value=1),
        gr.Dropdown(work_class_list),
        gr.Slider(1260, 12225, value=200),
        gr.Dropdown(education_list),
        gr.Slider(1, 16, value=2),
        gr.Dropdown(martial_list),
        gr.Dropdown(occupation_list),
        gr.Dropdown(relation_list),
        gr.Dropdown(race_list),
        gr.Dropdown([' Male',' Female']),
        gr.Slider(0, 10000, value=100),
        gr.Slider(0, 4500, value=75),
        gr.Slider(1, 100, value=2),
        gr.Dropdown(countries),
    ],
    outputs = "text",
    title = title,
    description = description,
    article = article,
    examples=
    [
        [35,' Private',5000,' Masters',8,' Divorced',' Tech-support',' Husband',' White',' Male',6000,0,40,' Germany'],
        [27,' Self-emp-inc',2400,' Bachelors',6,' Separated',' Prof-specialty',' Wife',' Amer-Indian-Eskimo',' Female',4000,1050,32,' England'],
    ]
)

if __name__ == "__main__":
    demo.launch()