import xgboost as xgb
import pandas as pd
import gradio as gr

from sklearn.model_selection import train_test_split
train = pd.read_csv('california_housing_train.csv') ## https://developers.google.com/machine-learning/crash-course/california-housing-data-description
train['total_bedrooms'].fillna(train['total_bedrooms'].mean(), inplace=True)
train.drop('households', axis=1, inplace=True)
train.drop('total_rooms', axis=1, inplace=True)
train.drop('total_bedrooms', axis=1, inplace=True)
train=train.loc[train['median_house_value']<=500000,:]

X=train[['longitude',
 'latitude',
 'housing_median_age',
 'population',
 'median_income']]
Y=train['median_house_value']

X_train,X_test,Y_train,Y_test = train_test_split(X,Y,test_size = 0.2, random_state=0)

xg_reg = xgb.XGBRegressor(objective ='reg:squarederror', eta=0.03, max_depth = 5, n_estimators = 200)
xg_reg.fit(X_train,Y_train)

def pred_house(longitude,latitude,housing_median_age,population,median_income):
    to_predict = pd.DataFrame({'longitude': [longitude],
                               'latitude': [latitude], 
                               'housing_median_age': [housing_median_age],
                               'population': [population],
                               'median_income': [median_income]})
    return round(xg_reg.predict(to_predict)[0],2)

demo = gr.Interface(
    fn=pred_house,
    inputs=[gr.Slider(-120, -114), gr.Slider(32, 39), gr.Slider(0, 50), "number", "number", ],
    outputs="number",
)
demo.launch()