Compressor/models/realtime_service.py

#!/usr/bin/env python3
"""Realtime PKS alarm scanner.

Runs the same pressure/vibration algorithm as the full scan over a rolling
window ending at the current (settled) hour, then merges the resulting stage
segments into ``compressor_alarm``.  A growing same-level segment is merged into
one row (never split), while different levels remain separate rows.

Modes
-----
* ``--once``                      run a single tick at the current local time.
* ``--loop --interval 300``       run a tick every ``interval`` seconds (CLI timer).
* ``--simulate "t1,t2,t3"``       smoke test: run ticks at the given times and
                                  (by default) roll the database back afterwards.

Examples
--------
    python models/realtime_service.py --once
    python models/realtime_service.py --loop --interval 300
    python models/realtime_service.py --simulate "2026-02-03 20:00:00,2026-02-04 00:00:00"
"""
from __future__ import annotations

import argparse
import sys
import time
from collections import Counter
from datetime import datetime, timedelta
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, queries
from models.generalized import pressure_algorithm
from models.vibration import vibration_algorithm

REALTIME_CHUNK_DAYS = 7


def run_tick(now: datetime, device_list: list[str], lookback_days: int, write: bool, journal: list | None, with_model: bool = False) -> dict:
    device_batches = devices.load_device_batches(device_list)
    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)

    ranges = devices.device_data_range(device_batches)
    latest = max(mx for _, mx in ranges.values())
    end = min(devices.floor_hour(now), devices.floor_hour(latest) + timedelta(hours=1))
    start = end - timedelta(days=lookback_days)
    if start >= end:
        return {"tick": now, "end": end, "action": "no_data", "inserted": 0, "updated": 0}

    rows = queries.fetch_all_hourly(start, end, REALTIME_CHUNK_DAYS, unit_by_batch, configs, sorted(unit_by_batch))
    hourly_levels, hour_segments, cycles = pressure_algorithm.run(rows, configs)
    events = vibration_algorithm.run(rows, configs)
    if with_model:
        model_predict.annotate("pressure", rows, hour_segments)
        model_predict.annotate("vibration", rows, events)
    alarms = alarm_store.build_alarms(io.clean_rows(hour_segments), io.clean_rows(events), batch_device)
    result = alarm_store.merge_realtime(alarms, start, end, journal=journal if write else None, dry_run=not write)
    return {
        "tick": now, "window": (start, end), "hourly_rows": len(rows),
        "pressure_segments": len(hour_segments), "vibration_events": len(events),
        "segments_to_store": len(alarms), **result,
    }


def format_tick(summary: dict) -> str:
    window = summary.get("window")
    window_text = f"{window[0]}..{window[1]}" if window else "-"
    return (
        f"[{summary['tick']}] window={window_text} 压力段={summary.get('pressure_segments')} "
        f"振动段={summary.get('vibration_events')} 写入段={summary.get('segments_to_store')} "
        f"insert={summary.get('inserted')} update={summary.get('updated')} merged_extra={summary.get('merged_extra')}"
    )


def snapshot(device_list: list[str], limit: int = 8) -> str:
    rows = alarm_store.fetch_alarms(device_list)
    counter = Counter((r["device_code"], r["alarm_type"], r["device_point"], r["alarm_level"]) for r in rows)
    return "; ".join(f"{k}={v}" for k, v in sorted(counter.items(), key=lambda item: str(item[0]))[:limit])


def main() -> None:
    parser = argparse.ArgumentParser(description="PKS 实时告警扫描（结果写入 compressor_alarm）")
    parser.add_argument("--device", nargs="+", default=list(devices.DEVICE_CODES))
    parser.add_argument("--lookback-days", type=int, default=45)
    parser.add_argument("--once", action="store_true", help="运行一次（默认）")
    parser.add_argument("--loop", action="store_true", help="按 --interval 循环运行")
    parser.add_argument("--interval", type=int, default=300, help="循环间隔秒，默认 300")
    parser.add_argument("--simulate", default=None, help="逗号分隔的时间列表，冒烟测试")
    parser.add_argument("--with-model", action="store_true", help="对故障段做模型滑动预测（默认关闭）")
    parser.add_argument("--no-rollback", action="store_true", help="冒烟测试后不回滚数据库")
    args = parser.parse_args()
    device_list = devices.normalize_devices(args.device)
    if args.lookback_days < 1:
        parser.error("--lookback-days 必须大于 0")

    if args.simulate:
        times = [queries.parse_dt(item.strip()) for item in args.simulate.split(",") if item.strip()]
        journal: list = []
        for moment in times:
            summary = run_tick(moment, device_list, args.lookback_days, write=True, journal=journal, with_model=args.with_model)
            print(format_tick(summary), flush=True)
            print("  当前告警分布：" + snapshot(device_list), flush=True)
        if not args.no_rollback:
            undone = alarm_store.rollback(journal)
            print(f"冒烟测试回滚：{undone}", flush=True)
        return

    if args.loop:
        print(f"实时服务启动，每 {args.interval} 秒运行一次；Ctrl+C 退出。", flush=True)
        while True:
            summary = run_tick(datetime.now(), device_list, args.lookback_days, write=True, journal=None, with_model=args.with_model)
            print(format_tick(summary), flush=True)
            time.sleep(args.interval)

    summary = run_tick(datetime.now(), device_list, args.lookback_days, write=True, journal=None, with_model=args.with_model)
    print(format_tick(summary))
    print("当前告警分布：" + snapshot(device_list))


if __name__ == "__main__":
    main()