Compressor/backend/LabelingPreTraining/generate_samples.py

"""Generate per-cycle pretraining samples from the wave database.

Reads 预训练之文件头.csv, groups the files by (point_name, measurement_type),
and for every wave file runs the project's cycle detector. Each detected cycle
is saved as one .npy matrix with columns
    [signal_value, second_value, 1, 1]
in time order. The file name is
    {wave_file_id}_{cycle_number}_{start_sample_index}_{end_sample_index_exclusive}.npy
"""

import csv
import sys
import traceback
from collections import defaultdict
from pathlib import Path

import numpy as np
from tqdm import tqdm

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

from app.algorithms.cycles import detect_cycles  # noqa: E402
from app.db import get_connection  # noqa: E402

HERE = Path(__file__).resolve().parent
CSV_PATH = ROOT / "预训练之文件头.csv"
OUTPUT_DIR = HERE / "inputdatas"
MANIFEST_PATH = HERE / "manifest.csv"


def load_file_rows() -> list[dict[str, str]]:
    with CSV_PATH.open(encoding="utf-8-sig") as handle:
        return list(csv.DictReader(handle))


def read_wave_samples(connection, file_id: int) -> np.ndarray:
    with connection.cursor() as cursor:
        cursor.execute(
            """
            SELECT sample_index, signal_value, second_value
            FROM wave_sample
            WHERE wave_file_id = %s
            ORDER BY wave_file_id, sample_index
            """,
            (file_id,),
        )
        rows = cursor.fetchall()
    return np.asarray(
        [
            (
                float(row["sample_index"]),
                float(row["signal_value"]),
                float(row["second_value"]) if row["second_value"] is not None else np.nan,
            )
            for row in rows
        ],
        dtype=float,
    )


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 cycle_sample_bounds(samples: np.ndarray, start_offset: int, end_offset: int) -> tuple[int, int]:
    start_sample_index = int(samples[start_offset, 0])
    if end_offset < len(samples):
        end_exclusive = int(samples[end_offset, 0])
    else:
        end_exclusive = int(samples[-1, 0]) + 1
    return start_sample_index, end_exclusive


def main() -> int:
    rows = load_file_rows()
    if not rows:
        print("CSV 为空，退出")
        return 1

    groups: dict[tuple[str, str], list[dict[str, str]]] = defaultdict(list)
    for row in rows:
        groups[(row["point_name"], row["measurement_type"])].append(row)

    manifest: list[dict[str, object]] = []
    total_cycles = 0
    connection = get_connection()
    try:
        group_progress = tqdm(sorted(groups.items()), desc="分组处理", unit="组")
        for group_index, (group_key, group_rows) in enumerate(group_progress, start=1):
            point_name, measurement_type = group_key
            folder = OUTPUT_DIR / f"{point_name}_{measurement_type}"
            folder.mkdir(parents=True, exist_ok=True)
            group_progress.set_postfix(group=f"{point_name}_{measurement_type}")

            for row in tqdm(
                sorted(group_rows, key=lambda item: int(item["id"])),
                desc=f"  {point_name}_{measurement_type}",
                unit="文件",
                leave=False,
            ):
                file_id = int(row["id"])
                manifest_row: dict[str, object] = {
                    "group": f"{point_name}_{measurement_type}",
                    "point_name": point_name,
                    "measurement_type": measurement_type,
                    "wave_file_id": file_id,
                    "rpm": row["rpm"],
                    "sample_time": row["sample_time"],
                    "cycle_count": 0,
                    "total_samples": 0,
                    "start_sample_index": "",
                    "end_sample_index": "",
                    "error": "",
                }
                try:
                    samples = read_wave_samples(connection, file_id)
                    if len(samples) == 0:
                        raise ValueError("没有采样数据")
                    cycles, diagnostics = detect_cycles(samples)
                    manifest_row["cycle_count"] = len(cycles)
                    saved = 0
                    total_cycle_rows = 0
                    first_start = ""
                    last_end = ""
                    for cycle in cycles:
                        start_sample_index, end_exclusive = cycle_sample_bounds(
                            samples, cycle.start_offset, cycle.end_offset
                        )
                        if not first_start:
                            first_start = str(start_sample_index)
                        last_end = str(end_exclusive)
                        matrix = build_cycle_matrix(
                            samples, cycle.start_offset, cycle.end_offset
                        )
                        if not np.all(np.isfinite(matrix[:, :2])):
                            raise ValueError("周期内存在非有限值")
                        path = folder / f"{file_id}_{cycle.number}_{start_sample_index}_{end_exclusive}.npy"
                        np.save(path, matrix)
                        saved += 1
                        total_cycle_rows += len(matrix)
                    total_cycles += saved
                    manifest_row["cycle_count"] = saved
                    manifest_row["total_samples"] = total_cycle_rows
                    manifest_row["start_sample_index"] = first_start
                    manifest_row["end_sample_index"] = last_end
                    if saved == 0:
                        manifest_row["error"] = (
                            f"未检出完整周期 (触发脉冲 {diagnostics.get('triggerRunCount', '?')} 个, "
                            f"零标记 {diagnostics.get('zeroMarkerCount', '?')} 个)"
                        )
                except Exception as error:
                    manifest_row["error"] = f"{error}: {traceback.format_exc(limit=1).strip()}"
                manifest.append(manifest_row)
                if manifest_row["error"]:
                    print(f"  id={file_id} 失败: {manifest_row['error'][:160]}")
    finally:
        connection.close()

    with MANIFEST_PATH.open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(
            handle,
            fieldnames=[
                "group",
                "point_name",
                "measurement_type",
                "wave_file_id",
                "rpm",
                "sample_time",
                "cycle_count",
                "total_samples",
                "start_sample_index",
                "end_sample_index",
                "error",
            ],
        )
        writer.writeheader()
        writer.writerows(manifest)

    ok = sum(1 for item in manifest if not item["error"])
    failed = sum(1 for item in manifest if item["error"])
    print(f"\n完成: {ok} 个文件成功, {failed} 个失败, 共生成 {total_cycles} 个周期样本")
    return 0


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