Compressor/backend/LabelingPreTraining/predict.py

"""Step 3: long-running TSPulse anomaly detection poller.

Polls the wave_file table for rows matching ``rpm > 0 AND measurement_type =
'压力'``, reads only the first 3000 wave_sample rows per file, locates the
first complete crankshaft cycle, and encodes it with the frozen TSPulse model
into a 128-dim fingerprint. The fingerprint is compared against the part's
centroid and radius from the baseline table and ``tspluse_status`` is written
back as a proportional integer of the centroid distance:

* 无周期 (no complete cycle found) -> ``-1``
* 正常 (distance <= 2 x radius)   -> ``0``
* 异常 (distance > 2 x radius)    -> ``min(127, round(10 x distance / radius))``

Every processed file is written back, so a full re-run from an empty watermark
cleanly replaces any previously stored levels.

The samples of a whole batch of files are read with one ``IN (...) AND
sample_index < 3000`` query and all candidate cycles are encoded in a single
forward pass, so the network round trips and MPS kernel launches are amortised
across the batch. Files are processed in ``(sample_time, id)`` ascending order
and a tuple watermark on those two columns prevents re-processing files that
have already been seen, so the poller first drains the backlog and then keeps
watching for newly inserted rows.

Usage:
    conda activate tspulse
    python predict.py [--start-time ""] [--start-id 0] [--interval 60]
                      [--batch 100] [--limit 0] [--report-every 100]
                      [--write-batch 10] [--dry-run]
"""

import argparse
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np
import torch

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

from app.algorithms.cycles import detect_cycles  # noqa: E402
from app.db import get_connection  # noqa: E402
from tspulse import TSPulse, describe, get_device  # noqa: E402
from tspulse.dataset import N_CHANNELS, SEQ_LEN  # noqa: E402

HERE = Path(__file__).resolve().parent
CHECKPOINT_PATH = HERE / "checkpoints" / "tspulse_frozen.pt"
BASELINE_PATH = HERE / "baseline" / "baseline.npz"

SAMPLE_LIMIT = 3000
MEASUREMENT_TYPE = "压力"

ANOMALY_FACTOR = 2.0
SCALE = 10.0
STATUS_MAX = 127


def load_model(device: torch.device) -> tuple[TSPulse, np.ndarray, np.ndarray]:
    checkpoint = torch.load(CHECKPOINT_PATH, map_location="cpu", weights_only=False)
    config = checkpoint["config"]
    model = TSPulse(dim=config["dim"], depth=config["depth"], heads=config["heads"]).to(device)
    model.load_state_dict(checkpoint["state_dict"])
    model.eval()
    return model, checkpoint["mean"], checkpoint["std"]


def load_baseline() -> tuple[dict[str, int], np.ndarray, np.ndarray]:
    data = np.load(BASELINE_PATH)
    part_names = data["part_names"].astype(str).tolist()
    part_index = {part: i for i, part in enumerate(part_names)}
    return part_index, data["centroids"], data["radii"]


def _fetch_chunk(connection, file_ids: list[int]) -> dict[int, np.ndarray]:
    """Fetch the first SAMPLE_LIMIT rows of a chunk of files in one query."""
    if not file_ids:
        return {}
    placeholders = ",".join(["%s"] * len(file_ids))
    with connection.cursor() as cursor:
        cursor.execute(
            f"""
            SELECT wave_file_id,
                   CAST(signal_value AS FLOAT) AS sig,
                   CAST(second_value AS FLOAT) AS sec
            FROM wave_sample
            WHERE wave_file_id IN ({placeholders}) AND sample_index < %s
            ORDER BY wave_file_id, sample_index
            """,
            (*file_ids, SAMPLE_LIMIT),
        )
        rows = cursor.fetchall()
    grouped: dict[int, tuple[list[float], list[float]]] = {}
    for row in rows:
        signal, second = grouped.setdefault(row["wave_file_id"], ([], []))
        signal.append(float(row["sig"]))
        second.append(float(row["sec"]) if row["sec"] is not None else np.nan)
    out: dict[int, np.ndarray] = {}
    for file_id, (signal, second) in grouped.items():
        count = len(signal)
        out[file_id] = np.column_stack(
            [np.arange(count, dtype=float), np.asarray(signal), np.asarray(second)]
        )
    return out


def read_batch_samples(
    reader_connections: list,
    file_ids: list[int],
    workers: int,
) -> dict[int, np.ndarray]:
    """Fetch the first SAMPLE_LIMIT rows of every file in one round-trip.

    The batch is split into ``workers`` chunks fetched on parallel persistent
    connections. Returns {file_id: (n, 3) array} with columns
    [sample_index(0..n-1), signal_value, second_value].
    """
    if not file_ids:
        return {}
    if workers <= 1 or len(file_ids) <= workers:
        return _fetch_chunk(reader_connections[0], file_ids)
    chunks = np.array_split(file_ids, min(workers, len(file_ids)))
    with ThreadPoolExecutor(max_workers=len(chunks)) as executor:
        futures = [
            executor.submit(_fetch_chunk, reader_connections[i], chunk.tolist())
            for i, chunk in enumerate(chunks)
        ]
        merged: dict[int, np.ndarray] = {}
        for future in futures:
            merged.update(future.result())
    return merged


def build_cycle_matrix(samples: np.ndarray, start: int, end: int) -> np.ndarray:
    signal = samples[start:end, 1]
    second = samples[start:end, 2]
    count = end - start
    return np.column_stack([signal, second, np.ones(count), np.ones(count)])


def flush_status(connection, buffer: list[tuple[int, int]]) -> None:
    """Write pending (file_id, status) pairs in a single UPDATE statement."""
    if not buffer:
        return
    case_sql = " ".join("WHEN %s THEN %s" for _ in buffer)
    placeholders = ",".join(["%s"] * len(buffer))
    params: list[int] = []
    for file_id, level in buffer:
        params.extend([file_id, level])
    with connection.cursor() as cursor:
        cursor.execute(
            f"UPDATE wave_file SET tspluse_status = CASE id {case_sql} END "
            f"WHERE id IN ({placeholders})",
            tuple(params + [file_id for file_id, _ in buffer]),
        )
    buffer.clear()


def fetch_batch(
    connection,
    watermark_time: object | None,
    watermark_id: int,
    batch: int,
) -> list[dict]:
    with connection.cursor() as cursor:
        if watermark_time is None:
            cursor.execute(
                """
                SELECT id, point_name, measurement_type, sample_time
                FROM wave_file
                WHERE rpm > 0 AND measurement_type = %s
                ORDER BY sample_time ASC, id ASC
                LIMIT %s
                """,
                (MEASUREMENT_TYPE, batch),
            )
        else:
            cursor.execute(
                """
                SELECT id, point_name, measurement_type, sample_time
                FROM wave_file
                WHERE rpm > 0 AND measurement_type = %s
                  AND (sample_time > %s OR (sample_time = %s AND id > %s))
                ORDER BY sample_time ASC, id ASC
                LIMIT %s
                """,
                (MEASUREMENT_TYPE, watermark_time, watermark_time, watermark_id, batch),
            )
        return cursor.fetchall()


def compute_status(distance: float, radius: float) -> int:
    """Map a fingerprint distance to a proportional integer status.

    Normal (distance <= ANOMALY_FACTOR x radius) returns 0; anomalies are
    recorded as ``min(STATUS_MAX, round(SCALE x distance / radius))`` so the
    stored tinyint is proportional to the centroid distance and can be queried
    with ``tspluse_status >= N``. 无周期 files are handled separately as -1.
    """
    ratio = distance / radius
    if ratio <= ANOMALY_FACTOR:
        return 0
    return min(STATUS_MAX, int(round(SCALE * ratio)))


def prepare_cycle(
    samples: np.ndarray,
    part: str,
    part_index: dict[str, int],
) -> tuple[str, np.ndarray | None]:
    """Locate the first complete cycle and build its [signal, second, 1, 1] matrix.

    Returns (status, matrix); status is "ok" or a skip reason.
    """
    cycles, _ = detect_cycles(samples)
    if not cycles:
        return "无周期", None

    cycle = cycles[0]
    matrix = build_cycle_matrix(samples, cycle.start_offset, cycle.end_offset)
    if len(matrix) > SEQ_LEN:
        return "周期过长", None
    if not np.all(np.isfinite(matrix[:, :2])):
        return "坏数据", None
    if part not in part_index:
        return "无基准", None
    return "ok", matrix


def encode_batch(
    model: TSPulse,
    mean: np.ndarray,
    std: np.ndarray,
    device: torch.device,
    matrices: list[np.ndarray],
) -> np.ndarray:
    """Encode many cycle matrices in a single forward pass -> (B, dim)."""
    if not matrices:
        return np.empty((0, model.dim), dtype=np.float32)
    batch = np.zeros((len(matrices), SEQ_LEN, N_CHANNELS), dtype=np.float32)
    for index, matrix in enumerate(matrices):
        normalised = (matrix - mean) / std
        batch[index, : len(normalised)] = normalised
    tensor = torch.from_numpy(batch).to(device)
    with torch.no_grad():
        fingerprints, _ = model(tensor)
    return fingerprints.cpu().numpy()


def main() -> int:
    parser = argparse.ArgumentParser(description="TSPulse 长期异常预测轮询脚本")
    parser.add_argument("--start-time", type=str, default="",
                        help="起始水位时间(含)，如 2026-01-01 00:00:00，留空从头开始")
    parser.add_argument("--start-id", type=int, default=0,
                        help="起始水位 id(含)，同一 sample_time 时作为次序")
    parser.add_argument("--interval", type=int, default=60, help="无新数据时的轮询间隔秒数")
    parser.add_argument("--batch", type=int, default=100, help="每轮读取的文件数")
    parser.add_argument("--read-workers", type=int, default=4, help="并行读样本的连接数(0=单连接)")
    parser.add_argument("--limit", type=int, default=0, help="最多处理文件数，0 表示不限")
    parser.add_argument("--report-every", type=int, default=100, help="每 N 个文件打印一次汇总")
    parser.add_argument("--write-batch", type=int, default=100,
                        help="回写攒满 N 条才批量写一次(每个文件都会回写)")
    parser.add_argument("--dry-run", action="store_true", help="只预测不回写数据库")
    args = parser.parse_args()

    if not CHECKPOINT_PATH.exists():
        print(f"未找到模型 {CHECKPOINT_PATH}, 请先运行 tspulse/train.py")
        return 1
    if not BASELINE_PATH.exists():
        print(f"未找到基准表 {BASELINE_PATH}, 请先运行 build_baseline.py")
        return 1

    device = get_device()
    print(f"设备: {describe(device)}")
    model, mean, std = load_model(device)
    part_index, centroids, radii = load_baseline()
    print(f"基准表: {len(part_index)} 个部位, 异常判定: 距离 > {ANOMALY_FACTOR:g}×半径, "
          f"回写 tspluse_status = min({STATUS_MAX}, round({SCALE:g} × 距离/半径)), 无周期=-1")

    start_time = args.start_time.strip()
    watermark_time: object | None = start_time or None
    watermark_id = args.start_id
    total = 0
    stats = {"正常": 0, "异常": 0, "无周期": 0}
    skipped = {"无基准": 0, "坏数据": 0, "周期过长": 0, "无样本": 0}
    errors = 0
    last_reported = 0
    write_buffer: list[tuple[int, int]] = []

    read_workers = max(1, min(args.read_workers, args.batch))
    reader_connections = [get_connection() for _ in range(read_workers)]

    def flush(connection) -> None:
        if args.dry_run:
            write_buffer.clear()
        else:
            flush_status(connection, write_buffer)

    def summary_text() -> str:
        return (
            f"已处理 {total}, 正常 {stats['正常']}, 异常 {stats['异常']}, "
            f"无周期 {stats['无周期']}, 跳过 {sum(skipped.values())}({dict(skipped)}), "
            f"错误 {errors}, 水位 sample_time={watermark_time} id={watermark_id}"
        )

    print(f"水位起点: sample_time={start_time or '从头'}, id={args.start_id}, "
          f"批大小: {args.batch}, 读并行: {read_workers} 连接, 间隔: {args.interval}s"
          + (", 干跑(不回写)" if args.dry_run else ""))

    try:
        while True:
            connection = get_connection()
            try:
                rows = fetch_batch(connection, watermark_time, watermark_id, args.batch)
                if not rows:
                    flush(connection)
                    connection.close()
                    print(
                        f"无新数据(sample_time>{watermark_time})，等待 {args.interval}s 后继续... "
                        f"[{summary_text()}]",
                        flush=True,
                    )
                    time.sleep(args.interval)
                    continue

                file_ids = [int(row["id"]) for row in rows]
                part_of = {int(row["id"]): f"{row['point_name']}_{row['measurement_type']}" for row in rows}
                sample_map = read_batch_samples(reader_connections, file_ids, read_workers)

                candidates: list[tuple[int, str, np.ndarray]] = []
                for row in rows:
                    file_id = int(row["id"])
                    part = part_of[file_id]
                    samples = sample_map.get(file_id)
                    if samples is None or len(samples) == 0:
                        status, matrix = "无样本", None
                    else:
                        try:
                            status, matrix = prepare_cycle(samples, part, part_index)
                        except Exception:
                            errors += 1
                            status, matrix = "错误", None

                    total += 1

                    if status == "ok":
                        candidates.append((file_id, part, matrix))
                    elif status == "无周期":
                        stats["无周期"] += 1
                        write_buffer.append((file_id, -1))
                    else:
                        if status in skipped:
                            skipped[status] += 1
                        write_buffer.append((file_id, 0))

                if candidates:
                    fingerprints = encode_batch(
                        model, mean, std, device, [item[2] for item in candidates]
                    )
                    for (file_id, part, _matrix), fingerprint in zip(candidates, fingerprints):
                        index = part_index[part]
                        distance = float(np.linalg.norm(fingerprint - centroids[index]))
                        radius = float(radii[index])
                        status = compute_status(distance, radius)
                        stats["异常" if status > 0 else "正常"] += 1
                        write_buffer.append((file_id, status))
                        if len(write_buffer) >= args.write_batch:
                            flush(connection)

                if total // args.report_every > last_reported:
                    last_reported = total // args.report_every
                    print(f"  汇总: {summary_text()}", flush=True)

                if args.limit and total >= args.limit:
                    flush(connection)
                    watermark_time = rows[-1]["sample_time"]
                    watermark_id = int(rows[-1]["id"])
                    connection.close()
                    print(f"已达 --limit={args.limit}，退出")
                    print(f"  汇总: {summary_text()}")
                    return 0

                flush(connection)
                watermark_time = rows[-1]["sample_time"]
                watermark_id = int(rows[-1]["id"])
            finally:
                try:
                    connection.close()
                except Exception:
                    pass

    except KeyboardInterrupt:
        if write_buffer and not args.dry_run:
            try:
                with get_connection() as connection:
                    flush_status(connection, write_buffer)
            except Exception:
                pass
        for connection in reader_connections:
            try:
                connection.close()
            except Exception:
                pass
        print(f"\n已中断。{summary_text()}")
        if watermark_time is not None:
            print(f"续跑命令: python predict.py --start-time \"{watermark_time}\" --start-id {watermark_id}")
        return 0

    for connection in reader_connections:
        try:
            connection.close()
        except Exception:
            pass
    return 0


if __name__ == "__main__":
    sys.exit(main())