import gradio as gr
import tensorflow as tf
import numpy as np

# Load the pickled model
model = tf.keras.models.load_model("census.h5")

# Mapping of categorical variables to encoded values
mapping = {
    'workclass': {' ?': 0, ' Federal-gov': 1, ' Local-gov': 2, ' Never-worked': 3, ' Private': 4, ' Self-emp-inc': 5, ' Self-emp-not-inc': 6, ' State-gov': 7, ' Without-pay': 8}, 
    'education': {' 10th': 0, ' 11th': 1, ' 12th': 2, ' 1st-4th': 3, ' 5th-6th': 4, ' 7th-8th': 5, ' 9th': 6, ' Assoc-acdm': 7, ' Assoc-voc': 8, ' Bachelors': 9, ' Doctorate': 10, ' HS-grad': 11, ' Masters': 12, ' Preschool': 13, ' Prof-school': 14, ' Some-college': 15}, 
    'marital_status': {' Divorced': 0, ' Married-AF-spouse': 1, ' Married-civ-spouse': 2, ' Married-spouse-absent': 3, ' Never-married': 4, ' Separated': 5, ' Widowed': 6}, 
    'occupation': {' ?': 0, ' Adm-clerical': 1, ' Armed-Forces': 2, ' Craft-repair': 3, ' Exec-managerial': 4, ' Farming-fishing': 5, ' Handlers-cleaners': 6, ' Machine-op-inspct': 7, ' Other-service': 8, ' Priv-house-serv': 9, ' Prof-specialty': 10, ' Protective-serv': 11, ' Sales': 12, ' Tech-support': 13, ' Transport-moving': 14}, 
    'relationship': {' Husband': 0, ' Not-in-family': 1, ' Other-relative': 2, ' Own-child': 3, ' Unmarried': 4, ' Wife': 5}, 
    'race': {' Amer-Indian-Eskimo': 0, ' Asian-Pac-Islander': 1, ' Black': 2, ' Other': 3, ' White': 4}, 
    'gender': {' Female': 0, ' Male': 1}, 
    'native_country': {' ?': 0, ' Cambodia': 1, ' Canada': 2, ' China': 3, ' Columbia': 4, ' Cuba': 5, ' Dominican-Republic': 6, ' Ecuador': 7, ' El-Salvador': 8, ' England': 9, ' France': 10, ' Germany': 11, ' Greece': 12, ' Guatemala': 13, ' Haiti': 14, ' Honduras': 15, ' Hong': 16, ' Hungary': 17, ' India': 18, ' Iran': 19, ' Ireland': 20, ' Italy': 21, ' Jamaica': 22, ' Japan': 23, ' Laos': 24, ' Mexico': 25, ' Nicaragua': 26, ' Outlying-US(Guam-USVI-etc)': 27, ' Peru': 28, ' Philippines': 29, ' Poland': 30, ' Portugal': 31, ' Puerto-Rico': 32, ' Scotland': 33, ' South': 34, ' Taiwan': 35, ' Thailand': 36, ' Trinadad&Tobago': 37, ' United-States': 38, ' Vietnam': 39, ' Yugoslavia': 40}
}

# Define the function for making predictions
def salarybracket(age, workclass, education, education_num, marital_status, occupation, relationship, race, gender, capital_gain, capital_loss, hours_per_week, native_country):
    # Get the value associated with the selected label
    workclass_value = next((v for k, v in mapping['workclass'].items() if k == workclass), None)
    education_value = next((v for k, v in mapping['education'].items() if k == education), None)
    marital_status_value = next((v for k, v in mapping['marital_status'].items() if k == marital_status), None)
    occupation_value = next((v for k, v in mapping['occupation'].items() if k == occupation), None)
    relationship_value = next((v for k, v in mapping['relationship'].items() if k == relationship), None)
    race_value = next((v for k, v in mapping['race'].items() if k == race), None)
    gender_value = next((v for k, v in mapping['gender'].items() if k == gender), None)
    native_country_value = next((v for k, v in mapping['native_country'].items() if k == native_country), None)
    
    inputs = np.array([[float(age), float(workclass_value), float(education_value), float(education_num), float(marital_status_value), float(occupation_value), float(relationship_value), float(race_value), float(gender_value), float(capital_gain), float(capital_loss), float(hours_per_week), float(native_country_value)]])
    prediction = model.predict(inputs)
    prediction_value = prediction[0][0]  # Assuming the prediction is a scalar
    result = "Income_bracket lesser than or equal to 50K 😊↓ " if prediction_value <= 0.5 else "Income_bracket greater than 50K 😄↑ "
    return f"{result}"

# Convert mapping to Gradio dropdown options
dropdown_options = {}
for column, values in mapping.items():
    options = []
    for label, value in values.items():
        options.append(label)
    dropdown_options[column] = options

# Create the Gradio interface
salarybracket_ga = gr.Interface(fn=salarybracket, 
                                inputs=[
                                    gr.Slider(17, 90, label="Age"),
                                    gr.Dropdown(dropdown_options['workclass'], label="Workclass"),
                                    gr.Dropdown(dropdown_options['education'], label="Education"),
                                    gr.Number(1, 16, label="Education Num [1 to 16]"),
                                    gr.Dropdown(dropdown_options['marital_status'], label="Marital Status"),
                                    gr.Dropdown(dropdown_options['occupation'], label="Occupation"),
                                    gr.Dropdown(dropdown_options['relationship'], label="Relationship"),
                                    gr.Dropdown(dropdown_options['race'], label="Race"),
                                    gr.Dropdown(dropdown_options['gender'], label="Gender"),
                                    gr.Number(0, 99999, label="Capital Gain [0 to 99999]"),
                                    gr.Number(0, 4356, label="Capital Loss [0 to 4356]"),
                                    gr.Number(1, 99, label="Hours per Week [1 to 99]"),
                                    gr.Dropdown(dropdown_options['native_country'], label="Native Country")                                        
                                ],
                                outputs="text",
                                title="Salary Bracket Prediction",
                                description="Predicting Income Bracket Using TensorFlow",
                                theme='dark'
                               )

salarybracket_ga.launch(share=True, debug=True)