import pickle

import gradio as gr
import pandas as pd


def rf_predict(GDP, Unemployment, Medical_resources, Past_years_life_expectancy_growth_rate):
    # 读取模型
    with open('rf.pkl', 'rb') as f:
        rf = pickle.load(f)
    # 预测
    X_test = [GDP, Unemployment, Medical_resources, Past_years_life_expectancy_growth_rate]
    X_test = pd.DataFrame(X_test).T
    X_test.columns = ['GDP_per_capita', 'Unemployment_rate', 'Medical_resources_per_capita',
                      'Past_years_life_expectancy_growth_rate']
    y_pred = rf.predict(X_test)
    return y_pred


# def rf_predict(GDP, Unemployment, Medical_resources, Past_years_life_expectancy_growth_rate):
# print(rf_predict(85698, 5.5, 2120, 0.1))
gr.Interface(
    fn=rf_predict,
    inputs=[
        gr.Textbox(label="GDP"),
        gr.Textbox(label="Unemployment"),
        gr.Textbox(label="Medical_resources"),
        gr.Textbox(label="Past_years_life_expectancy_growth_rate"),
    ],
    outputs="number",
    title="Life Expectancy Prediction"
).launch()