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

# ——— 配置 ———
windCode   = "WOF046400029"  # 张崾先：WOF091200030 七台河：WOF046400029
start, end = "2023-10-02 00:00:00", "2024-10-02 00:00:00"  # 张崾先：2023-10-20 00:00:00~2024-10-20 00:00:00 七台河：2023-10-02 00:00:00~2024-10-02 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"
    "mysql+pymysql://root:admin123456@192.168.50.235:30306/energy_data_prod"    
)
# ——————————

# def list_turbines() -> list[str]:
#     # 获取风场所有风机编号（有数据的）
#     sql = f"""
#       SELECT DISTINCT wind_turbine_number
#       FROM {windCode}_minute
#       WHERE time_stamp BETWEEN '{start}' AND '{end}'
#     """
#     return pd.read_sql(sql, engine)['wind_turbine_number'].tolist()

def list_turbines() -> list[str]:
    # 直接返回待训练的风机编号列表
    return ['WOG01344']

def fetch_channel_data(turbine: str, channel: str) -> np.ndarray:
    sql = f"""
      SELECT {channel}
      FROM {windCode}_minute
      WHERE wind_turbine_number = '{turbine}'
        AND time_stamp BETWEEN '{start}' AND '{end}'
      ORDER BY time_stamp ASC
    """
    df = pd.read_sql(sql, engine).dropna(subset=[channel])
    print(f"[TRAIN] 风机 {turbine} 通道 {channel} 拉取 {len(df)} 条")
    return df[channel].values.reshape(-1,1)

if __name__ == "__main__":
    turbines = list_turbines()
    print("[TRAIN] 本次训练风机列表：", turbines)

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

            model = MSET_Temp(windCode, [turbine], 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] {turbine}-{ch} D/L 构建失败")
                continue

            out_dir = os.path.join(model_root, windCode, turbine)
            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))