Compressor/backend/LabelingPreTraining/split_device_point.py

"""Split wave_file.point_name into device_part and device_point columns.

Every rpm > 0 point_name has the shape ``{N号机组}{部位}{测点}``, e.g.
``7号机组二缸压力盖侧`` -> ``7号机组二缸`` + ``压力盖侧``. The split is a
longest-known-suffix strip: the 测点 comes from a fixed whitelist of suffixes,
the remainder must match ``\\d+号机组(?:[一二三四五六]缸|曲轴箱|电机)``, and
reassembling must reproduce the original name. Any name failing these checks
is reported and the write is aborted, so no data is ever lost or mangled.

The mapping is computed locally from the distinct point_names, then written in
a single ``UPDATE ... CASE point_name WHEN ... END`` (one pass over the table),
covering every row whose point_name can be split (no rpm filter, so rpm<=0/NULL
rows are backfilled too; already-correct rows are idempotently rewritten with
the same values). ``device_point`` is added with an ``ALTER TABLE`` if it does
not exist yet. Run without ``--commit`` to only preview the mappings.

Usage:
    conda activate tspulse
    python split_device_point.py            # preview only
    python split_device_point.py --commit   # add column + write back
"""

import argparse
import re
import sys
from pathlib import Path

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

from predict import get_connection  # noqa: E402

KNOWN_POINTS = [
    "压力盖侧",
    "压力轴侧",
    "活塞杆沉降",
    "十字头振动",
    "自由端振动",
    "驱动端振动",
]
PART_RE = re.compile(r"^\d+号机组(?:[一二三四五六]缸|曲轴箱|电机)$")


def split_name(name: str) -> tuple[str, str] | None:
    """Return (device_part, device_point) or None if the name does not fit."""
    for suffix in sorted(KNOWN_POINTS, key=len, reverse=True):
        if name.endswith(suffix):
            part = name[: -len(suffix)]
            if PART_RE.match(part):
                return part, suffix
            return None
    return None


def build_mapping(connection) -> tuple[dict[str, tuple[str, str]], list[str]]:
    with connection.cursor() as cursor:
        cursor.execute("SELECT DISTINCT point_name FROM wave_file")
        names = [row["point_name"] for row in cursor.fetchall()]
    mapping: dict[str, tuple[str, str]] = {}
    failed: list[str] = []
    for name in names:
        split = split_name(name)
        if split is None:
            failed.append(name)
        else:
            mapping[name] = split
    return mapping, failed


def main() -> int:
    parser = argparse.ArgumentParser(description="拆分 point_name 到 device_part/device_point")
    parser.add_argument("--commit", action="store_true",
                        help="真正执行：新增 device_point 列并回写（默认仅预览）")
    args = parser.parse_args()

    with get_connection() as connection:
        mapping, failed = build_mapping(connection)

        if failed:
            print(f"以下 {len(failed)} 个 point_name 无法拆分，已中止，未做任何改动：")
            for name in failed:
                print(f"  {name!r}")
            return 1

        print(f"共 {len(mapping)} 个 point_name，全部拆分成功：")
        for name, (part, point) in sorted(mapping.items()):
            print(f"  {name:26} -> {part:12} | {point}")

        if not args.commit:
            print("\n未执行 --commit，仅预览。确认无误后加 --commit 回写。")
            return 0

        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT COUNT(*) AS n FROM information_schema.columns
                WHERE table_schema = DATABASE()
                  AND table_name = 'wave_file' AND column_name = 'device_point'
                """
            )
            if not cursor.fetchone()["n"]:
                cursor.execute(
                    "ALTER TABLE wave_file ADD COLUMN device_point varchar(150) NOT NULL DEFAULT ''"
                )
                print("已新增列 device_point varchar(150) NOT NULL DEFAULT ''")

        names = sorted(mapping)
        placeholders = ",".join(["%s"] * len(names))
        case_part = " ".join(
            f"WHEN %s THEN %s" for _ in names
        )
        case_point = " ".join(f"WHEN %s THEN %s" for _ in names)
        sql = f"""
            UPDATE wave_file
            SET device_part = CASE point_name {case_part} END,
                device_point = CASE point_name {case_point} END
            WHERE point_name IN ({placeholders})
        """
        params: list[str] = []
        for name in names:
            params.extend([name, mapping[name][0]])
        for name in names:
            params.extend([name, mapping[name][1]])
        params.extend(names)

        with connection.cursor() as cursor:
            affected = cursor.execute(sql, tuple(params))
        print(f"已回写 {affected} 行（point_name 可拆的全部行）。")
        return 0


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