# build_grid_bp.py
import polars as pl
from tqdm import tqdm
import numpy as np
from pathlib import Path
from pca import predict_backpressure  # 你之前实现的函数
import csv

# ---------- 参数 ----------
OUT_FILE = Path("grid_bp_抽凝.csv")
BATCH_SIZE = 50_000  # 每多少行写一次盘/进度条更新
# 网格定义
GRID = {
    "汽轮发电机功率": range(60, 295, 10),  # 24 档
    "汽机抽汽供热量": range(-30, 1911, 100),  # 20 档
    "环境温度": range(-10, 43, 5),  # 11 档
    "大气湿度": range(5, 100, 20),  # 5 档
    "循泵运行方式数字编码": [1, 2, 3, 4, 5, 6],  # 6 档
    "风机运行台数": range(0, 9),  # 9 档
    "冷却塔进水蝶阀开启数量": range(0, 9),  # 9 档
}
N_TOTAL = np.prod([len(v) for v in GRID.values()])


# ---------- 工具 ----------

def append_batch(df_list: list[pl.DataFrame]):
    """批量追加到 CSV（Polars 兼容写法）"""
    if not df_list:
        return
    df_batch = pl.concat(df_list)
    # 首次写带表头，之后追加
    if not OUT_FILE.exists():
        df_batch.write_csv(OUT_FILE)
    else:
        # 手动追加，不带表头
        with OUT_FILE.open("a", newline="", encoding="utf-8") as f:
            csv.writer(f).writerows(df_batch.rows())


# ---------- 主循环 ----------
def build_grid():
    keys = list(GRID.keys())
    iterables = [GRID[k] for k in keys]
    batch_rows = []
    pbar = tqdm(total=N_TOTAL, desc="构造网格")
    for values in np.array(np.meshgrid(*iterables)).T.reshape(-1, len(keys)):
        sample = dict(zip(keys, values))
        # 用之前写好的 predict_backpressure
        new_df = pl.DataFrame([sample])
        bp = predict_backpressure(new_df)[0]
        sample["背压"] = bp
        batch_rows.append(pl.DataFrame([sample]))
        pbar.update(1)
        # 批量写盘
        if len(batch_rows) >= BATCH_SIZE:
            append_batch(batch_rows)
            batch_rows.clear()
    # 末尾不足一批
    append_batch(batch_rows)
    pbar.close()
    print(f"✅ 完成，共 {N_TOTAL} 行 → {OUT_FILE.resolve()}")


# ---------- 入口 ----------
if __name__ == "__main__":
    build_grid()
