Compressor/backend/indicator_features/run_segments.py

#!/usr/bin/env python3
"""Machine running segments from ``wave_file.rpm > 0`` (read-only).

The site's 开机停机表 is known to be inaccurate, so the running state is taken
from the data itself: every calendar day that has at least one pressure record
with ``rpm > 0`` is a running day.  Consecutive running days (allowing a small
merge gap for missing daily records) form one *running segment*.  ``event_bins``
uses these segments to pick the pre/post running periods instead of fixed
event-relative windows.
"""
from __future__ import annotations

import argparse
import csv
import sys
from datetime import date
from pathlib import Path

import pymysql

BACKEND = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(BACKEND))

from app.config import settings  # noqa: E402


def connect() -> pymysql.connections.Connection:
    return pymysql.connect(
        host=settings.db_host, port=settings.db_port, user=settings.db_user,
        password=settings.db_password, database=settings.db_name, charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor, connect_timeout=8, read_timeout=120,
    )


def running_days(unit: str, connection=None) -> list[date]:
    """Sorted distinct days with a running pressure record for one unit."""
    own = connection is None
    connection = connection or connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT DISTINCT DATE(sample_time) AS d
                FROM wave_file
                WHERE point_name LIKE %s AND measurement_type = '压力' AND rpm > 0
                ORDER BY d
                """,
                (f"{unit}号机组%",),
            )
            return [row["d"] for row in cursor.fetchall()]
    finally:
        if own:
            connection.close()


def build_segments(days: list[date], merge_gap: int = 2) -> list[tuple[date, date]]:
    """Merge running days into segments; a gap of at most ``merge_gap`` off-days is bridged."""
    segments: list[list[date]] = []
    for day in days:
        if segments and (day - segments[-1][1]).days <= merge_gap + 1:
            segments[-1][1] = day
        else:
            segments.append([day, day])
    return [(start, end) for start, end in segments]


def unit_segments(unit: str, merge_gap: int = 2, connection=None) -> list[tuple[date, date]]:
    return build_segments(running_days(unit, connection), merge_gap)


def main() -> None:
    parser = argparse.ArgumentParser(description="wave_file.rpm>0 统计开机日与运行段（只读）")
    parser.add_argument("--unit", required=True, choices=("7", "8", "9"))
    parser.add_argument("--merge-gap", type=int, default=2, help="允许跨过的停机天数（桥接）")
    parser.add_argument("--out", default="", help="运行段 CSV 输出路径（可选）")
    args = parser.parse_args()

    days = running_days(args.unit)
    segments = build_segments(days, args.merge_gap)
    print(f"{args.unit}号机 rpm>0 开机日 {len(days)} 天，运行段 {len(segments)} 个：")
    for start, end in segments:
        print(f"  {start} ~ {end}  ({(end - start).days + 1}天)")
    if args.out:
        out = Path(args.out)
        out.parent.mkdir(parents=True, exist_ok=True)
        with out.open("w", encoding="utf-8-sig", newline="") as handle:
            writer = csv.writer(handle)
            writer.writerow(["unit", "start", "end", "days"])
            for start, end in segments:
                writer.writerow([args.unit, start, end, (end - start).days + 1])
        print("->", out)


if __name__ == "__main__":
    main()