# train_models.py

import numpy as np
import pandas as pd
from sqlalchemy import create_engine
from Temp_Diag import MSET_Temp
import os

# ——— 配置 ———
windCode   = "WOF091200030"
# turbines = ["WOG01355"]
start, end = "2024-06-12 00:00:00", "2024-07-12 00:00:00"
model_root = "models"
channels   = [
    'main_bearing_temperature',
    'gearbox_oil_temperature',
    'generatordrive_end_bearing_temperature',
    'generatornon_drive_end_bearing_temperature'
]
engine = create_engine(
    "mysql+pymysql://root:admin123456@106.120.102.238:10336/energy_data_prod"
)
# ——————————

def fetch_channel_data(channel: str) -> np.ndarray:
    sql = f"""
      SELECT {channel}
      FROM {windCode}_minute
      WHERE time_stamp BETWEEN '{start}' AND '{end}'
    #   AND wind_turbine_number IN ({','.join(f"'{t}'" for t in turbines)})
      ORDER BY time_stamp ASC
    """
    df = pd.read_sql(sql, engine).dropna(subset=[channel])
    print(f"[TRAIN] {channel} 共 {len(df)} 条")
    return df[channel].values.reshape(-1,1)

if __name__ == "__main__":
    for ch in channels:
        data = fetch_channel_data(ch)
        if data.shape[0] < 65:
            print(f"[TRAIN] {ch} 样本不足，跳过")
            continue

        model = MSET_Temp(windCode, [], start, end)
        model.feature_weight = np.ones(data.shape[1],)
        model.alpha = 0.1; model.beta = 0.1

        if model.genDLMatrix(data, dataSize4D=60, dataSize4L=5) != 0:
            print(f"[TRAIN] {ch} D/L 构建失败")
            continue

        out_dir = os.path.join(model_root, windCode)
        os.makedirs(out_dir, exist_ok=True)
        path = os.path.join(out_dir, f"{ch}.pkl")
        model.save_model(path)
        print(f"[TRAIN] 已保存 {path}")

    print("\n[TRAIN] 模型文件列表:")
    for root, _, files in os.walk(model_root):
        for fn in files:
            print(" ", os.path.join(root, fn))
