# backpressure_idw.py
import polars as pl
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.neighbors import KDTree
from pathlib import Path
import pickle

# ---------- 路径 ----------
CSV_PATH = Path(r"C:\wps同步\高安屯冷端历史\抽凝纯输入输出.csv")
MODEL_PATH = Path("bp2-chou.pkl")


class BackPressureModel:
    def __init__(self, model_path=MODEL_PATH):
        """初始化模型，从指定路径加载训练好的模型"""
        if model_path.exists():
            with open(model_path, 'rb') as f:
                model_data = pickle.load(f)
            print(f"模型已从 {model_path} 加载")
        else:
            # 如果模型文件不存在，自动训练并保存模型
            print(f"模型文件 {model_path} 不存在，正在重新训练模型...")
            model_data = self._build_model()
            self.save_model(model_data, model_path)
            print(f"模型已保存到 {model_path}")

        scalars, pca_models, train_factors, bp_train, kdt = model_data
        self.scalars = scalars
        self.pca_models = pca_models
        self.train_factors = train_factors
        self.bp_train = bp_train
        self.kdt = kdt

        # 定义列名映射
        self.groups = {
            "heat": ["汽轮发电机功率", "汽机抽汽供热量"],
            "weather": ["环境温度", "大气湿度"],
            "operation": ["循泵运行方式数字编码", "风机运行台数", "冷却塔进水蝶阀开启数量"],
        }

    def predict(self, input_data):
        """
        预测函数，接受列表或数组格式的输入数据
        输入: [汽轮发电机功率, 汽机抽汽供热量, 环境温度, 大气湿度, 循泵运行方式数字编码, 风机运行台数, 冷却塔进水蝶阀开启数量]
        """
        if isinstance(input_data, (list, tuple, np.ndarray)):
            # 如果输入是单个样本（一维），转换为二维
            if len(np.array(input_data).shape) == 1:
                input_data = [input_data]

        # 将输入数据转换为DataFrame格式
        df_input = pl.DataFrame({
            "汽轮发电机功率": [x[0] for x in input_data],
            "汽机抽汽供热量": [x[1] for x in input_data],
            "环境温度": [x[2] for x in input_data],
            "大气湿度": [x[3] for x in input_data],
            "循泵运行方式数字编码": [x[4] for x in input_data],
            "风机运行台数": [x[5] for x in input_data],
            "冷却塔进水蝶阀开启数量": [x[6] for x in input_data]
        })

        # 新数据 → 3 维因子
        factors = np.empty((len(df_input), 3))
        for idx, (g, cols) in enumerate(self.groups.items()):
            X = df_input.select(cols).to_numpy()
            Xstd = self.scalars[g].transform(X)
            factors[:, idx] = self.pca_models[g].transform(Xstd).ravel()

        # KD-Tree 取最近 64 点做 IDW
        dist, idx = self.kdt.query(factors, k=64)
        weights = 1.0 / (dist + 1e-8) ** 2
        weights /= weights.sum(axis=1, keepdims=True)
        bp_pred = (weights * self.bp_train[idx]).sum(axis=1)

        # 如果输入是单个样本，返回单个值
        if len(input_data) == 1:
            return bp_pred[0]
        else:
            return bp_pred.tolist()

    def _build_model(self):
        """内部方法：训练模型"""
        df = pl.read_csv(CSV_PATH)
        # 清洗
        df = df.filter(pl.col("背压").is_not_null() & (pl.col("背压") > 0))

        groups = {
            "heat": ["汽轮发电机功率", "汽机抽汽供热量"],
            "weather": ["环境温度", "大气湿度"],
            "operation": ["循泵运行方式数字编码", "风机运行台数", "冷却塔进水蝶阀开启数量"],
        }
        out_col = "背压"

        # 每组标准化 + PCA 降 1 维
        scalars, pca_models = {}, {}
        factors = np.empty((len(df), 3))
        for idx, (g, cols) in enumerate(groups.items()):
            X = df.select(cols).to_numpy()
            scalars[g] = StandardScaler().fit(X)
            Xstd = scalars[g].transform(X)
            pca_models[g] = PCA(n_components=1).fit(Xstd)
            factors[:, idx] = pca_models[g].transform(Xstd).ravel()

        # 全局变量
        bp_train = df[out_col].to_numpy().ravel()
        kdt = KDTree(factors, leaf_size=40)

        return scalars, pca_models, factors, bp_train, kdt

    def save_model(self, model_data, filepath=MODEL_PATH):
        """保存训练好的模型到pkl文件"""
        with open(filepath, 'wb') as f:
            pickle.dump(model_data, f)
        print(f"模型已保存到: {filepath}")

    @classmethod
    def load_model(cls, filepath=MODEL_PATH):
        """类方法：从pkl文件加载模型"""
        return cls(filepath)

    @classmethod
    def train_and_save_model(cls, csv_path=CSV_PATH, model_path=MODEL_PATH):
        """类方法：重新训练模型并保存"""
        print("正在重新训练模型...")
        instance = cls.__new__(cls)  # 创建实例但不调用__init__
        model_data = instance._build_model()
        instance.save_model(model_data, model_path)
        print("模型训练和保存完成")
        # 返回新的实例
        return cls(model_path)

    def test_prediction_error(self):
        """测试预测结果的误差"""
        df = pl.read_csv(CSV_PATH)
        # 将预测结果作为新列加到df中，并计算误差
        df_with_pred = df.with_columns(self._predict_dataframe(df))
        mean_error = df_with_pred.select(pl.col("背压") - pl.col("背压_预测")).mean()
        print(mean_error)
        df_with_pred.write_csv("backpressure_idw_pred.csv")
        return mean_error

    def _predict_dataframe(self, new: pl.DataFrame) -> pl.Series:
        """
        内部方法: 用于DataFrame的预测，与原函数功能相同
        """
        # 新数据 → 3 维因子
        factors = np.empty((len(new), 3))
        for idx, (g, cols) in enumerate(self.groups.items()):
            X = new.select(cols).to_numpy()
            Xstd = self.scalars[g].transform(X)
            factors[:, idx] = self.pca_models[g].transform(Xstd).ravel()

        # KD-Tree 取最近 64 点做 IDW
        dist, idx = self.kdt.query(factors, k=64)
        weights = 1.0 / (dist + 1e-8) ** 2
        weights /= weights.sum(axis=1, keepdims=True)
        bp_pred = (weights * self.bp_train[idx]).sum(axis=1)
        return pl.Series("背压_预测", bp_pred)


# ---------- 脚本自测 ----------
if __name__ == "__main__":
    # 测试新的API - 直接实例化
    model = BackPressureModel()
    model.test_prediction_error()
    result = model.predict([81, -17, 29, 12, 3, 4, 5])
    print(f"新API预测结果: {result}")

    # 测试批量预测
    batch_result = model.predict([[81, -17, 29, 12, 3, 4, 5], [82, -2, 13, 37, 3, 2, 3]])
    print(f"批量预测结果: {batch_result}")

    # 测试类方法
    model2 = BackPressureModel.load_model()
    result2 = model2.predict([81, -17, 29, 12, 3, 4, 5])
    print(f"使用load_model方法加载模型预测结果: {result2}")

    # 如果需要重新训练模型
    # trained_model = BackPressureModel.train_and_save_model()
