Compressor/models/scan_all.py

#!/usr/bin/env python3
"""Full PKS scan (pressure + vibration) with device-code input.

Modes
-----
* ``--store local`` (default): delete previous intermediate/result files, then
  rebuild them under ``models/datas`` / ``models/results``.  Output is byte
  compatible with the reference CSVs.
* ``--store db``: rebuild intermediate files locally, then replace the
  ``compressor_alarm`` rows (pressure + vibration) for the selected devices.

Device input
------------
``--device 7# 8# 9#`` (default).  Each device resolves to its PKS import
batches; hourly aggregation keeps the original batch grouping so the result
matches the existing ``润滑油压力小时阶段段.csv`` / ``振动异常段.csv``.

Examples
--------
    python models/scan_all.py --device 7# 8# 9#
    python models/scan_all.py --auto-range --store db
    python models/scan_all.py --start "2025-01-01 00:00:00" --end "2026-09-01 00:00:00"
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

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

from models.common import alarm_store, db, devices, io, model_predict, paths, queries
from models.generalized import pressure_algorithm
from models.vibration import vibration_algorithm


SCAN_OUTPUT_FILES = {
    paths.COMMON_RESULTS: ["报警阈值配置.csv"],
    paths.PRESSURE_RESULTS: ["润滑油压力小时等级.csv", "润滑油压力小时阶段段.csv", "润滑油压力演化周期.csv"],
    paths.VIBRATION_RESULTS: ["振动异常段.csv"],
}


def clear_dir_files(directory: Path, names: list[str] | None = None) -> int:
    """Delete files in a directory; if ``names`` given, only those file names."""
    directory = Path(directory)
    if not directory.exists():
        return 0
    removed = 0
    targets = [directory / name for name in names] if names is not None else [entry for entry in directory.iterdir() if entry.is_file()]
    for entry in targets:
        if entry.is_file():
            entry.unlink()
            removed += 1
    return removed


def clean_intermediates() -> dict:
    """Remove previous scan intermediate/result files only (keep model outputs)."""
    removed = {"hourly": clear_dir_files(paths.HOURLY_DIR)}
    for directory, names in SCAN_OUTPUT_FILES.items():
        removed[directory.name] = clear_dir_files(directory, names)
    return removed


def main() -> None:
    parser = argparse.ArgumentParser(description="PKS 压力+振动全量扫描（按设备号）")
    parser.add_argument("--device", nargs="+", default=list(devices.DEVICE_CODES), help="压缩机编号，如 7# 8# 9#")
    parser.add_argument("--start", default=None)
    parser.add_argument("--end", default=None)
    parser.add_argument("--auto-range", action="store_true", help="按数据自动确定起止时间（默认在未指定 --start/--end 时启用）")
    parser.add_argument("--chunk-days", type=int, default=14)
    parser.add_argument("--hourly-csv", default=None, help="复用已落盘的小时中间表，不查询数据库")
    parser.add_argument("--store", choices=("local", "db"), default="local")
    parser.add_argument("--with-model", action="store_true", help="对故障段做模型滑动预测（默认关闭）")
    parser.add_argument("--dry-run", action="store_true", help="只统计不写库（store=db 时）")
    args = parser.parse_args()

    devices_selected = devices.normalize_devices(args.device)
    device_batches = devices.load_device_batches(devices_selected)
    unit_by_batch = devices.batch_to_unit(device_batches)
    batch_device = devices.batch_to_device(device_batches)
    configs = db.load_alarm_config(unit_by_batch)
    print(f"设备：{devices_selected}")
    print(f"批次映射：{ {d: bs for d, bs in device_batches.items()} }")

    rows = None
    if args.hourly_csv:
        print(f"复用小时中间表：{args.hourly_csv}")
        rows = queries.load_hourly_csv(Path(args.hourly_csv))

    removed = clean_intermediates()
    print(f"已清理中间/结果文件：{removed}")

    if args.start and args.end:
        start, end = queries.parse_dt(args.start), queries.parse_dt(args.end)
    else:
        start, end = devices.auto_range(device_batches)
        print(f"自动时间范围：{start} .. {end}")

    if rows is None:
        rows = queries.fetch_all_hourly(start, end, args.chunk_days, unit_by_batch, configs, sorted(unit_by_batch))
    queries.write_hourly_csv(paths.HOURLY_FULL, rows)
    print(f"小时中间表已缓存：{paths.HOURLY_FULL}")

    io.write_thresholds(paths.COMMON_RESULTS / "报警阈值配置.csv", configs)

    hourly_levels, hour_segments, cycles = pressure_algorithm.run(rows, configs)
    events = vibration_algorithm.run(rows, configs)
    print(f"小时特征：{len(rows):,} 行")
    print(f"润滑油压力小时等级：{len(hourly_levels):,} 行")
    print(f"润滑油压力小时阶段段：{len(hour_segments):,} 段")
    print(f"润滑油压力周期：{len(cycles):,} 个")
    print(f"振动异常段：{len(events):,} 段")

    extra_fields = ()
    if args.with_model:
        pressure_stats = model_predict.annotate("pressure", rows, hour_segments)
        vibration_stats = model_predict.annotate("vibration", rows, events)
        print(f"模型预测：压力 {pressure_stats}，振动 {vibration_stats}")
        extra_fields = tuple(model_predict.MODEL_FIELDS) + (model_predict.NOTE_FIELD,)

    pressure_algorithm.write_results(paths.PRESSURE_RESULTS, hourly_levels, hour_segments, cycles, extra_fields=extra_fields)
    vibration_algorithm.write_results(paths.VIBRATION_RESULTS, events, extra_fields=extra_fields)
    print(f"结果已写入本地：{paths.RESULTS}")

    if args.store == "db":
        clean_segments = io.clean_rows(hour_segments)
        clean_events = io.clean_rows(events)
        alarms = alarm_store.build_alarms(clean_segments, clean_events, batch_device)
        result = alarm_store.write_full(alarms, devices_selected, start, end, dry_run=args.dry_run)
        print(f"数据库写入：{result}")


if __name__ == "__main__":
    main()