Compressor/backend/LabelingPreTraining/predict_re.py

"""Re-run TSPulse detection for files currently marked 无周期 (tspluse_status = -1).

The original full re-run used a fixed 3000-sample prefix per file. Files that
start mid-revolution have too few trigger pulses inside that prefix for the
zero-marker segmentation to find a complete cycle, so they were mislabelled -1
even though the client and the training generator (which read the whole file)
detect them normally.

This script re-processes ONLY the rows still holding ``tspluse_status = -1``,
using predict.py's adaptive multi-stage window but starting at the 7000-sample
stage (these files already failed the smaller windows). Files that still yield
no cycle stay -1; the rest are written their corrected status (0 or the
proportional anomaly value).

The query filters on ``tspluse_status = -1``, so the ``(sample_time, id)``
watermark is naturally idempotent: a file already fixed is no longer selected
on resume. Run it once after the full re-run finishes; no column reset needed.

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

import argparse
import sys

import numpy as np
import torch

import predict
from predict import (
    MEASUREMENT_TYPE,
    STAGE_LIMITS,
    compute_status,
    encode_batch,
    flush_status,
    load_baseline,
    load_model,
    resolve_batch_cycles,
)


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, sample_count
                FROM wave_file
                WHERE rpm > 0 AND measurement_type = %s AND tspluse_status = -1
                ORDER BY sample_time ASC, id ASC
                LIMIT %s
                """,
                (MEASUREMENT_TYPE, batch),
            )
        else:
            cursor.execute(
                """
                SELECT id, point_name, measurement_type, sample_time, sample_count
                FROM wave_file
                WHERE rpm > 0 AND measurement_type = %s AND tspluse_status = -1
                  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 main() -> int:
    parser = argparse.ArgumentParser(description="重新处理 tspluse_status = -1 的文件")
    parser.add_argument("--start-time", type=str, default="",
                        help="起始水位时间(含)，留空从头开始")
    parser.add_argument("--start-id", type=int, default=0,
                        help="起始水位 id(含)")
    parser.add_argument("--batch", type=int, default=100, help="每轮读取的文件数")
    parser.add_argument("--read-workers", type=int, default=4, help="并行读样本的连接数")
    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()

    device = predict.get_device()
    print(f"设备: {predict.describe(device)}")
    model, mean, std = load_model(device)
    part_index, centroids, radii = load_baseline()
    print(f"基准表: {len(part_index)} 个部位, 重新处理 tspluse_status = -1 的文件, "
          f"自适应窗口: {STAGE_LIMITS[1:]} (跳过 {STAGE_LIMITS[0]})")

    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 = [predict.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} 连接"
          + (", 干跑(不回写)" if args.dry_run else ""))

    try:
        while True:
            connection = predict.get_connection()
            try:
                rows = fetch_batch(connection, watermark_time, watermark_id, args.batch)
                if not rows:
                    flush(connection)
                    connection.close()
                    print(f"无待处理的 -1 文件，处理完成 [{summary_text()}]")
                    return 0

                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}
                resolved, batch_errors = resolve_batch_cycles(
                    rows, part_of, part_index, reader_connections, read_workers,
                    first_stage=1,
                )
                errors += batch_errors

                candidates: list[tuple[int, str, np.ndarray]] = []
                for row in rows:
                    file_id = int(row["id"])
                    total += 1
                    status, matrix = resolved[file_id]
                    if status == "ok":
                        candidates.append((file_id, part_of[file_id], 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 predict.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_re.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())