import joblib
import gradio as gr
import pandas as pd

# Load the preprocessor and the best model
preprocessor = joblib.load('preprocessor.pkl')
best_model = joblib.load('best_model_gradient_boosting.pkl')

# Define the prediction function
def predict_income(age, workclass, fnlwgt, education, education_num, marital_status, occupation, relationship, race, sex, capital_gain, capital_loss, hours_per_week, native_country):
    # Create a DataFrame for the input data
    input_data = pd.DataFrame({
        'age': [age],
        'workclass': [workclass],
        'fnlwgt': [fnlwgt],
        'education': [education],
        'education-num': [education_num],
        'marital-status': [marital_status],
        'occupation': [occupation],
        'relationship': [relationship],
        'race': [race],
        'sex': [sex],
        'capital-gain': [capital_gain],
        'capital-loss': [capital_loss],
        'hours-per-week': [hours_per_week],
        'native-country': [native_country]
    })

    # Ensure input data types match the training data
    input_data = input_data.astype({
        'age': 'int64',
        'workclass': 'object',
        'fnlwgt': 'int64',
        'education': 'object',
        'education-num': 'int64',
        'marital-status': 'object',
        'occupation': 'object',
        'relationship': 'object',
        'race': 'object',
        'sex': 'object',
        'capital-gain': 'int64',
        'capital-loss': 'int64',
        'hours-per-week': 'int64',
        'native-country': 'object'
    })

    # Preprocess the input data
    input_data_preprocessed = preprocessor.transform(input_data)

    # Make prediction
    prediction = best_model.predict(input_data_preprocessed)
    prediction_proba = best_model.predict_proba(input_data_preprocessed)[:, 1]

    # Map prediction to class
    income_class = '>50K' if prediction[0] == 1 else '<=50K'
    probability = prediction_proba[0]

    return income_class, probability

# Create the Gradio interface
input_fields = [
    gr.Number(label="Age"),
    gr.Dropdown(label="Workclass", choices=['Private', 'Self-emp-not-inc', 'Self-emp-inc', 'Federal-gov', 'Local-gov', 'State-gov', 'Without-pay', 'Never-worked']),
    gr.Number(label="Fnlwgt"),
    gr.Dropdown(label="Education", choices=['Bachelors', 'Some-college', '11th', 'HS-grad', 'Prof-school', 'Assoc-acdm', 'Assoc-voc', '9th', '7th-8th', '12th', 'Masters', '1st-4th', '10th', 'Doctorate', '5th-6th', 'Preschool']),
    gr.Number(label="Education-num"),
    gr.Dropdown(label="Marital-status", choices=['Married-civ-spouse', 'Divorced', 'Never-married', 'Separated', 'Widowed', 'Married-spouse-absent', 'Married-AF-spouse']),
    gr.Dropdown(label="Occupation", choices=['Tech-support', 'Craft-repair', 'Other-service', 'Sales', 'Exec-managerial', 'Prof-specialty', 'Handlers-cleaners', 'Machine-op-inspct', 'Adm-clerical', 'Farming-fishing', 'Transport-moving', 'Priv-house-serv', 'Protective-serv', 'Armed-Forces']),
    gr.Dropdown(label="Relationship", choices=['Wife', 'Own-child', 'Husband', 'Not-in-family', 'Other-relative', 'Unmarried']),
    gr.Dropdown(label="Race", choices=['White', 'Asian-Pac-Islander', 'Amer-Indian-Eskimo', 'Other', 'Black']),
    gr.Dropdown(label="Sex", choices=['Female', 'Male']),
    gr.Number(label="Capital-gain"),
    gr.Number(label="Capital-loss"),
    gr.Number(label="Hours-per-week"),
    gr.Dropdown(label="Native-country", choices=['United-States', 'Cambodia', 'England', 'Puerto-Rico', 'Canada', 'Germany', 'Outlying-US(Guam-USVI-etc)', 'India', 'Japan', 'Greece', 'South', 'China', 'Cuba', 'Iran', 'Honduras', 'Philippines', 'Italy', 'Poland', 'Jamaica', 'Vietnam', 'Mexico', 'Portugal', 'Ireland', 'France', 'Dominican-Republic', 'Laos', 'Ecuador', 'Taiwan', 'Haiti', 'Columbia', 'Hungary', 'Guatemala', 'Nicaragua', 'Scotland', 'Thailand', 'Yugoslavia', 'El-Salvador', 'Trinadad&Tobago', 'Peru', 'Hong', 'Holand-Netherlands'])
]

output_fields = [
    gr.Textbox(label="Predicted Income Class"),
    gr.Textbox(label="Probability of >50K Income")
]

gr.Interface(fn=predict_income, inputs=input_fields, outputs=output_fields, title="Income Prediction App", description="Predict whether an individual makes over $50K a year based on various attributes.").launch(share=True,debug=True)