Compressor/backend/export_pks_daily_10min.py

#!/usr/bin/env python3
"""只读导出 PKS 全场点位数据到 Excel（每机组一个 sheet，每天固定 10 分钟）。

第 1 行为英文字段名，第 2 行为中文名（取自 site_point.ItemDescription）。
用法示例：

    python backend/export_pks_daily_10min.py \
        --start 2026-07-12 --end 2026-07-21 \
        --window 16:14:00 --minutes 10 \
        --output 7-8-9号机_20260712-20260721_1614-1624.xlsx
"""
from __future__ import annotations

import argparse
import os
import re
import sys
from datetime import date, datetime, time, timedelta
from pathlib import Path

import pymysql
from openpyxl import Workbook
from openpyxl.styles import Font

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from app.config import settings  # noqa: E402

UNITS = ("7#", "8#", "9#")
SITE_POINT_COLUMNS = tuple(f"YSJ_{index}" for index in range(1, 42))
SAMPLE_SECONDS = 5


def connect():
    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=settings.db_connect_timeout,
        read_timeout=600,
        write_timeout=60,
        autocommit=True,
    )


def parse_date(value: str) -> date:
    return datetime.strptime(value, "%Y-%m-%d").date()


def parse_time(value: str) -> time:
    return datetime.strptime(value, "%H:%M:%S").time()


def unit_number(unit: str) -> str:
    match = re.search(r"([789])", unit)
    if not match:
        raise ValueError(f"无法从 {unit!r} 解析机组号")
    return match.group(1)


def load_chinese_names(units: list[str]) -> dict[str, dict[str, str]]:
    """读取每个机组 YSJ_1..41 的中文名，返回 {机组: {列名: 中文名}}。"""
    names = [f"YSJ{unit_number(unit)}_{index}" for unit in units for index in range(1, 42)]
    placeholders = ",".join(["%s"] * len(names))
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"SELECT ItemName, ItemDescription FROM site_point WHERE ItemName IN ({placeholders})",
                names,
            )
            description = {row["ItemName"]: row["ItemDescription"] for row in cursor.fetchall()}
    finally:
        connection.close()

    result: dict[str, dict[str, str]] = {}
    for unit in units:
        unit_names = {}
        for index in range(1, 42):
            column = f"YSJ_{index}"
            item_name = f"YSJ{unit_number(unit)}_{index}"
            unit_names[column] = description.get(item_name) or column
        result[unit] = unit_names
    return result


def fetch_unit_rows(unit: str, days: list[date], window: time, minutes: int) -> list[dict]:
    """按天查询指定机组固定 10 分钟窗口内的全部点位数据。"""
    columns = ", ".join(f"t.{column}" for column in SITE_POINT_COLUMNS)
    sql = f"""
        SELECT t.sample_time, {columns}
        FROM pks_long_sample t
        INNER JOIN compressor_data_import i ON i.id = t.import_batch_id
        WHERE i.directory_path = %s
          AND i.data_type = 'PKS'
          AND t.sample_time >= %s
          AND t.sample_time <= %s
        ORDER BY t.sample_time
    """
    rows: list[dict] = []
    connection = connect()
    try:
        with connection.cursor() as cursor:
            for day in days:
                start = datetime.combine(day, window)
                end = start + timedelta(minutes=minutes) - timedelta(seconds=SAMPLE_SECONDS)
                cursor.execute(sql, (unit, start, end))
                rows.extend(cursor.fetchall())
    finally:
        connection.close()
    return rows


def write_sheet(workbook: Workbook, unit: str, chinese_names: dict[str, str], rows: list[dict]) -> None:
    sheet = workbook.create_sheet(title=f"{unit_number(unit)}号机")
    header_en = ["sample_time", *SITE_POINT_COLUMNS]
    header_cn = ["采样时间", *(chinese_names[column] for column in SITE_POINT_COLUMNS)]

    sheet.append(header_en)
    sheet.append(header_cn)
    for cell in sheet[1]:
        cell.font = Font(bold=True)
    for cell in sheet[2]:
        cell.font = Font(bold=True)

    for row in rows:
        sample_time = row["sample_time"]
        values = [sample_time.strftime("%Y-%m-%d %H:%M:%S") if isinstance(sample_time, datetime) else sample_time]
        values.extend(row[column] for column in SITE_POINT_COLUMNS)
        sheet.append(values)

    sheet.freeze_panes = "A3"


def main() -> None:
    parser = argparse.ArgumentParser(description="只读导出 PKS 全场点位固定 10 分钟数据到 Excel")
    parser.add_argument("--start", default="2026-07-12", help="起始日期 YYYY-MM-DD（含）")
    parser.add_argument("--end", default="2026-07-21", help="结束日期 YYYY-MM-DD（含）")
    parser.add_argument("--window", default="16:14:00", help="每天窗口起始时刻 HH:MM:SS")
    parser.add_argument("--minutes", type=int, default=10, help="每天窗口分钟数")
    parser.add_argument("--units", nargs="+", default=list(UNITS), help="机组列表，默认 7# 8# 9#")
    parser.add_argument(
        "--output",
        default=str(Path(__file__).resolve().parents[1] / "7-8-9号机_20260712-20260721_1614-1624.xlsx"),
    )
    args = parser.parse_args()

    start, end = parse_date(args.start), parse_date(args.end)
    if end < start:
        parser.error("结束日期不能早于起始日期")
    if args.minutes < 1:
        parser.error("窗口分钟数必须大于 0")
    window = parse_time(args.window)

    days = [start + timedelta(days=offset) for offset in range((end - start).days + 1)]
    expected_per_day = args.minutes * 60 // SAMPLE_SECONDS
    print(f"只读模式：导出 {', '.join(args.units)}，{start} ~ {end}，每天 {args.window} 起 {args.minutes} 分钟")
    print(f"每天应有 {expected_per_day} 条（5 秒采样），共 {len(days)} 天")

    chinese_names = load_chinese_names(args.units)
    workbook = Workbook()
    workbook.remove(workbook.active)

    exported = []
    for unit in args.units:
        rows = fetch_unit_rows(unit, days, window, args.minutes)
        if not rows:
            print(f"{unit}：该时间段无数据，跳过")
            continue
        expected_total = expected_per_day * len(days)
        flag = "OK" if len(rows) == expected_total else "注意"
        print(f"{unit}：{len(rows)} 条（预期 {expected_total}）[{flag}]")
        write_sheet(workbook, unit, chinese_names[unit], rows)
        exported.append((unit, len(rows)))

    if not exported:
        parser.error("所选机组在指定时间段均无数据，未生成文件")

    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    workbook.save(output)
    print(f"输出文件：{output}")
    for unit, count in exported:
        print(f"  sheet {unit_number(unit)}号机：{count} 行")


if __name__ == "__main__":
    main()