# 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

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


# ---------- 分组方案工厂 ----------
def make_scheme(name: str):
    if name == "A物理":
        return {"heat": ["汽轮发电机功率", "汽机抽汽供热量"],
                "weather": ["环境温度", "大气湿度"],
                "operation": ["循泵运行方式数字编码", "风机运行台数", "冷却塔进水蝶阀开启数量"]}
    if name == "B打乱":
        return {"g1": ["汽轮发电机功率", "大气湿度", "冷却塔进水蝶阀开启数量"],
                "g2": ["汽机抽汽供热量", "循泵运行方式数字编码"],
                "g3": ["环境温度", "风机运行台数"]}
    if name == "C量纲":
        return {"energy": ["汽轮发电机功率", "汽机抽汽供热量", "环境温度"],
                "ratio": ["大气湿度", "背压"],
                "count": ["循泵运行方式数字编码", "风机运行台数", "冷却塔进水蝶阀开启数量"]}
    if name == "D整体PCA":
        return {"all": ["汽轮发电机功率", "汽机抽汽供热量", "环境温度",
                        "大气湿度", "循泵运行方式数字编码",
                        "风机运行台数", "冷却塔进水蝶阀开启数量"]}
    raise ValueError(name)

# ---------- 交叉验证 ----------
def cv_rmse(scheme_name: str, k: int = 5):
    from sklearn.model_selection import KFold
    scheme = make_scheme(scheme_name)
    df = pl.read_csv(CSV_PATH, encoding="utf8").filter(pl.col("背压") > 0)
    X = df[[c for cols in scheme.values() for c in cols]].to_numpy()
    y = df["背压"].to_numpy()
    kf = KFold(n_splits=k, shuffle=True, random_state=42)
    rmse_list = []
    for train_idx, test_idx in kf.split(X):
        # 训练集降维
        X_train, X_test = X[train_idx], X[test_idx]
        y_train, y_test = y[train_idx], y[test_idx]
        # 整体 PCA 方案
        if "all" in scheme:
            scaler = StandardScaler().fit(X_train)
            pca = PCA(n_components=3).fit(scaler.transform(X_train))
            fac_train = pca.transform(scaler.transform(X_train))
            fac_test = pca.transform(scaler.transform(X_test))
        else:
            # 分组 PCA
            scalars, pcas = {}, {}
            fac_train = np.empty((len(X_train), 3))
            fac_test = np.empty((len(X_test), 3))
            for idx, (g, cols) in enumerate(scheme.items()):
                idx_list = [scheme["all"].index(c) if "all" in scheme else
                            [c for cols in scheme.values() for c in cols].index(c)
                            for c in cols]
                Xg_train, Xg_test = X_train[:, idx_list], X_test[:, idx_list]
                scalars[g] = StandardScaler().fit(Xg_train)
                Xstd_train = scalars[g].transform(Xg_train)
                Xstd_test = scalars[g].transform(Xg_test)
                pcas[g] = PCA(n_components=1).fit(Xstd_train)
                fac_train[:, idx] = pcas[g].transform(Xstd_train).ravel()
                fac_test[:, idx] = pcas[g].transform(Xstd_test).ravel()
        # 训练 KD-Tree
        kdt = KDTree(fac_train, leaf_size=40)
        dist, idx = kdt.query(fac_test, k=64)
        w = 1/(dist+1e-8)**2
        w /= w.sum(axis=1, keepdims=True)
        y_pred = (w * y_train[idx]).sum(axis=1)
        rmse_list.append(np.sqrt(np.mean((y_test - y_pred)**2)))
    return np.mean(rmse_list), np.std(rmse_list)

# ---------- 一键扫描 ----------
if __name__ == "__main__":
    for s in ["A物理", "B打乱", "C量纲", "D整体PCA"]:
        mean, std = cv_rmse(s)
        print(f"{s}: RMSE = {mean:.4f} ± {std:.4f} kPa")