Compressor/_tmp_cylinder_vibration/data.py

"""数据读取：首周期波形 + 粗采样 + 温度参考。

振动特征只使用 wave_file / wave_sample_one 的首周期；
温度 CSV 只在事后验证时读取，不参与特征和标签。
"""
from __future__ import annotations

import re
from datetime import datetime, timedelta
from pathlib import Path

import numpy as np
import pandas as pd

from models.common.db import connect

from . import config


def fetch_metadata(
    cylinders: list[str] | None = None,
    start: datetime | None = None,
    end: datetime | None = None,
) -> list[dict]:
    cylinders = cylinders or list(config.CYLINDERS)
    points = [config.CYLINDERS[name]["point"] for name in cylinders]
    placeholders = ", ".join(["%s"] * len(points))
    clauses = [f"point_name IN ({placeholders})", "rpm > 0"]
    params: list = list(points)
    if start is not None:
        clauses.append("sample_time >= %s")
        params.append(start)
    if end is not None:
        clauses.append("sample_time <= %s")
        params.append(end)
    sql = f"""
        SELECT id, device_part, point_name, sample_time, rpm, sample_count,
               sample_frequency_hz, cycle_start, cycle_end
        FROM wave_file
        WHERE {' AND '.join(clauses)}
        ORDER BY point_name, sample_time
    """
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(sql, params)
            return cursor.fetchall()
    finally:
        connection.close()


def coarse_metadata(rows: list[dict], minutes: int) -> list[dict]:
    """每个缸每个粗时间桶保留一个文件，先做轻量筛查。"""
    if not rows:
        return []
    frame = pd.DataFrame(rows)
    frame["sample_time"] = pd.to_datetime(frame["sample_time"])
    frame["bucket"] = frame["sample_time"].dt.floor(f"{minutes}min")
    frame = (
        frame.sort_values(["point_name", "sample_time"])
        .drop_duplicates(["point_name", "bucket"], keep="first")
        .drop(columns=["bucket"])
    )
    return frame.to_dict("records")


def fetch_samples(file_ids: list[int], chunk_size: int = 200) -> dict[int, np.ndarray]:
    """按块读取首周期波形。返回 id -> signal_value 数组。"""
    result: dict[int, np.ndarray] = {}
    if not file_ids:
        return result
    connection = connect()
    try:
        with connection.cursor() as cursor:
            for offset in range(0, len(file_ids), chunk_size):
                chunk = file_ids[offset:offset + chunk_size]
                placeholders = ", ".join(["%s"] * len(chunk))
                cursor.execute(
                    f"""
                    SELECT wave_file_id, signal_value
                    FROM wave_sample_one
                    WHERE wave_file_id IN ({placeholders})
                    ORDER BY wave_file_id, sample_index
                    """,
                    chunk,
                )
                current_id: int | None = None
                values: list[float] = []
                for row in cursor.fetchall():
                    file_id = int(row["wave_file_id"])
                    if current_id is None:
                        current_id = file_id
                    if file_id != current_id:
                        result[current_id] = np.asarray(values, dtype=float)
                        values = []
                        current_id = file_id
                    values.append(float(row["signal_value"]))
                if current_id is not None:
                    result[current_id] = np.asarray(values, dtype=float)
    finally:
        connection.close()
    return result


def load_temperature_reference() -> pd.DataFrame:
    """读取温度异常段，仅用于最后对照。"""
    path = config.TEMPERATURE_LABELS
    if not path.exists():
        return pd.DataFrame()
    frame = pd.read_csv(path, encoding="utf-8-sig")
    if frame.empty:
        return frame
    frame = frame[frame["机组"].astype(str).str.strip() == "5号机"].copy()
    frame["cylinder"] = frame["阶段部位"].map(_temperature_cylinder)
    frame["start_time"] = pd.to_datetime(frame["开始小时"], errors="coerce")
    frame["end_time"] = pd.to_datetime(frame["结束小时"], errors="coerce")
    return frame


def _temperature_cylinder(label: str) -> str | None:
    text = str(label or "").strip()
    match = re.match(r"^(一级|二级)(\d)#缸$", text)
    if not match:
        return None
    return config.STAGE_CYLINDER_TO_GLOBAL.get((match.group(1), f"{match.group(2)}#"))


def merge_intervals(intervals: list[tuple[datetime, datetime]], gap: timedelta = timedelta(0)) -> list[tuple[datetime, datetime]]:
    if not intervals:
        return []
    intervals = sorted(intervals)
    merged = [intervals[0]]
    for start, end in intervals[1:]:
        prev_start, prev_end = merged[-1]
        if start <= prev_end + gap:
            merged[-1] = (prev_start, max(prev_end, end))
        else:
            merged.append((start, end))
    return merged


def refine_windows(candidate_labels: pd.DataFrame, events: list[dict]) -> dict[str, list[tuple[datetime, datetime]]]:
    """二次评估窗口：事件窗口 + 粗筛候选窗口。"""
    windows: dict[str, list[tuple[datetime, datetime]]] = {name: [] for name in config.CYLINDERS}

    for event in events:
        event_time = pd.Timestamp(event["date"]).to_pydatetime()
        cylinder = event["cylinder"]
        windows[cylinder].append((
            event_time - timedelta(days=10),
            event_time + timedelta(days=3),
        ))

    if not candidate_labels.empty:
        for row in candidate_labels.itertuples(index=False):
            cylinder = getattr(row, "cylinder", None)
            if cylinder not in windows:
                continue
            start = pd.Timestamp(getattr(row, "start_time")).to_pydatetime() - timedelta(days=3)
            end = pd.Timestamp(getattr(row, "end_time")).to_pydatetime() + timedelta(days=1)
            windows[cylinder].append((start, end))

    # 二次评估需要同级对照缸来计算跨缸残差。
    expanded = {name: list(items) for name, items in windows.items()}
    for cylinder, items in windows.items():
        stage = config.CYLINDERS[cylinder]["stage"]
        for peer in config.STAGE_PEERS[stage]:
            if peer != cylinder:
                expanded[peer].extend(items)

    return {
        cylinder: merge_intervals(items, gap=timedelta(hours=6))
        for cylinder, items in expanded.items()
        if items
    }