import pickle
import pandas as pd

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

# 初始化模型
model = load_model()   # 模型会在model.py被导入时加载一次

# 预测函数
def predict(input_data: pd.DataFrame):
    # 假设输入数据已经准备好
    predictions = model.predict(input_data)
    return predictions.tolist()

if __name__ == '__main__':
    # 测试 predict 函数
    x = 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
    }])

    print(predict(x))  # 预测结果