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()


def train_and_save_model(dataset_id, model_type, model_name, model_description):
    dataset = get_dataset_by_id(dataset_id)
    if dataset.empty:
        raise ValueError(f"Dataset {dataset_id} is empty or not found.")

    # Step 1: 数据准备
    X = dataset.iloc[:, :-1]  # 特征数据
    y = dataset.iloc[:, -1]  # 目标变量

    # Step 2: 训练模型
    model = train_model_by_type(X, y, model_type)

    # Step 3: 保存模型到数据库
    # 使用提供的 model_name 和 model_description
    saved_model = save_model(model_name, model_type, model_description)

    # Step 4: 保存模型参数
    save_model_parameters(model, saved_model.ModelID)

    # Step 5: 计算评估指标（比如MSE）
    y_pred = model.predict(X)
    mse = mean_squared_error(y, y_pred)

    return saved_model, mse


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)  # 预测结果