import pickle
import pandas as pd


# 加载模型
def load_model(model_name):
    file_path = f'model_optimize/pkl/{model_name}.pkl'
    with open(file_path, 'rb') as f:
        return pickle.load(f)


# 模型预测
def predict(input_data: pd.DataFrame, model_name):
    # 初始化模型
    model = load_model(model_name)  # 根据指定的模型名加载模型
    predictions = model.predict(input_data)
    return predictions.tolist()


if __name__ == '__main__':
    # 反酸模型预测
    # 测试 predict 函数
    input_data = pd.DataFrame([{
        "organic_matter": 5.2,
        "chloride": 3.1,
        "cec": 25.6,
        "h_concentration": 0.5,
        "hn": 12.4,
        "al_concentration": 0.8,
        "free_alumina": 1.2,
        "free_iron": 0.9,
        "delta_ph": -0.2
    }])

    model_name = 'RF_filt'

    Acid_reflux_result = predict(input_data, model_name)
    print("Acid_reflux_result:", Acid_reflux_result)  # 预测结果


    # 降酸模型预测
    # 测试 predict 函数
    input_data = pd.DataFrame([{
        "pH": 5.2,
        "OM": 3.1,
        "CL": 25.6,
        "H": 0.5,
        "Al": 12.4
    }])

    model_name = 'rf_model_1214_1008'
    Acid_reduce_result = predict(input_data, model_name)
    print("Acid_reduce_result:", Acid_reduce_result)  # 预测结果