Compressor/backend/LabelingPreTraining/detect_cycle_index.py

"""Backfill the first detected cycle's sample_index range into wave_file.

The original full re-run encoded the first complete cycle of every file but
never wrote back where that cycle starts and ends. This script re-detects the
first complete cycle for every pressure file (rpm > 0) and records the real
``sample_index`` bounds into two columns:

* ``cycle_start`` / ``cycle_end`` >= 0 : first complete cycle, half-open
  ``[cycle_start, cycle_end)`` in actual sample_index values;
* ``-1``                             : no complete cycle found (matches the
  ``tspluse_status = -1`` convention);
* ``NULL``                           : not yet processed.

It reuses predict.py's adaptive multi-stage batch reading (STAGE_LIMITS =
3500 -> 7000 -> 14000, parallel reader connections, tuple watermark) but needs
no model: only the trigger channel is used, so it runs on CPU only. The fetch
also keeps the real ``sample_index`` per row because ``detect_cycles`` returns
array offsets, which differ from database indices for sparse columns.

Only rows still holding ``cycle_start IS NULL`` are selected, so the
``(sample_time, id)`` watermark is naturally idempotent: a file already
written (0/-1) is never re-selected on resume. Writes are flushed in one
``CASE id`` UPDATE per ``--write-batch`` rows.

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

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

import numpy as np

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

from app.algorithms.cycles import detect_cycles  # noqa: E402
import predict  # noqa: E402
from predict import MEASUREMENT_TYPE, STAGE_LIMITS  # noqa: E402


def _fetch_chunk(connection, file_ids: list[int], limit: int) -> dict[int, np.ndarray]:
    """Fetch the first ``limit`` rows of a chunk of files in one query.

    Returns {file_id: (n, 4) array} with columns
    [array_offset(0..n-1), signal_value, second_value, sample_index].
    """
    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,
                   sample_index
            FROM wave_sample
            WHERE wave_file_id IN ({placeholders}) AND sample_index < %s
            ORDER BY wave_file_id, sample_index
            """,
            (*file_ids, limit),
        )
        rows = cursor.fetchall()
    grouped: dict[int, tuple[list[float], list[float], list[int]]] = {}
    for row in rows:
        signal, second, index = 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)
        index.append(int(row["sample_index"]))
    out: dict[int, np.ndarray] = {}
    for file_id, (signal, second, index) in grouped.items():
        count = len(signal)
        out[file_id] = np.column_stack(
            [
                np.arange(count, dtype=float),
                np.asarray(signal),
                np.asarray(second),
                np.asarray(index, dtype=float),
            ]
        )
    return out


def read_batch_samples(
    reader_connections: list,
    file_ids: list[int],
    workers: int,
    limit: int,
) -> dict[int, np.ndarray]:
    """Fetch the first ``limit`` rows of every file in one round-trip."""
    if not file_ids:
        return {}
    if workers <= 1 or len(file_ids) <= workers:
        return _fetch_chunk(reader_connections[0], file_ids, limit)
    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(), limit)
            for i, chunk in enumerate(chunks)
        ]
        merged: dict[int, np.ndarray] = {}
        for future in futures:
            merged.update(future.result())
    return merged


def first_cycle_bounds(samples: np.ndarray) -> tuple[int, int] | None:
    """Return ``(cycle_start, cycle_end)`` in real sample_index, half-open."""
    cycles, _ = detect_cycles(samples)
    if not cycles:
        return None
    cycle = cycles[0]
    return (int(samples[cycle.start_offset, 3]), int(samples[cycle.end_offset, 3]))


def resolve_batch_bounds(
    rows: list[dict],
    reader_connections: list,
    workers: int,
    first_stage: int = 0,
) -> tuple[dict[int, tuple[int, int] | None], int]:
    """Adaptively detect the first cycle per file, escalating window sizes.

    Returns (``{file_id: (start, end) | None}``, error_count). None means no
    complete cycle was found at any window (-> -1).
    """
    errors = 0
    resolved: dict[int, tuple[int, int] | None] = {}
    stage_files: dict[int, dict] = {int(row["id"]): row for row in rows}
    for stage in range(first_stage, len(STAGE_LIMITS)):
        if not stage_files:
            break
        limit = STAGE_LIMITS[stage]
        samples_map = read_batch_samples(reader_connections, list(stage_files), workers, limit)
        next_pending: dict[int, dict] = {}
        for file_id, row in stage_files.items():
            samples = samples_map.get(file_id)
            if samples is None or len(samples) == 0:
                resolved[file_id] = None
                continue
            try:
                bounds = first_cycle_bounds(samples)
            except Exception:
                errors += 1
                resolved[file_id] = None
                continue
            if bounds is None and int(row.get("sample_count") or 0) > limit:
                next_pending[file_id] = row
            else:
                resolved[file_id] = bounds
        stage_files = next_pending
    for file_id in stage_files:
        resolved[file_id] = None
    return resolved, errors


def flush_bounds(connection, buffer: list[tuple[int, int, int]]) -> None:
    """Write pending (file_id, cycle_start, cycle_end) rows in one UPDATE."""
    if not buffer:
        return
    case_start = " ".join("WHEN %s THEN %s" for _ in buffer)
    case_end = " ".join("WHEN %s THEN %s" for _ in buffer)
    placeholders = ",".join(["%s"] * len(buffer))
    params_start: list[int] = []
    params_end: list[int] = []
    for file_id, start, end in buffer:
        params_start.extend([file_id, start])
        params_end.extend([file_id, end])
    with connection.cursor() as cursor:
        cursor.execute(
            f"UPDATE wave_file SET cycle_start = CASE id {case_start} END, "
            f"cycle_end = CASE id {case_end} END "
            f"WHERE id IN ({placeholders})",
            tuple(params_start + params_end + [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, sample_count
                FROM wave_file
                WHERE rpm > 0 AND measurement_type = %s AND cycle_start IS NULL
                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 cycle_start IS NULL
                  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="回写首个完整周期的 sample_index 边界")
    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()

    start_time = args.start_time.strip()
    watermark_time: object | None = start_time or None
    watermark_id = args.start_id
    total = 0
    stats = {"有周期": 0, "无周期": 0}
    errors = 0
    last_reported = 0
    write_buffer: list[tuple[int, 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_bounds(connection, write_buffer)

    def summary_text() -> str:
        return (
            f"已处理 {total}, 有周期 {stats['有周期']}, 无周期 {stats['无周期']}, "
            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} 连接, 窗口: {STAGE_LIMITS}"
          + (", 干跑(不回写)" 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"无待处理文件，完成 [{summary_text()}]")
                    return 0

                resolved, batch_errors = resolve_batch_bounds(
                    rows, reader_connections, read_workers,
                )
                errors += batch_errors

                for row in rows:
                    file_id = int(row["id"])
                    total += 1
                    bounds = resolved[file_id]
                    if bounds is None:
                        stats["无周期"] += 1
                        write_buffer.append((file_id, -1, -1))
                    else:
                        stats["有周期"] += 1
                        write_buffer.append((file_id, bounds[0], bounds[1]))
                    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_bounds(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 detect_cycle_index.py --start-time \"{watermark_time}\" "
                  f"--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())