Compressor/backend/import_compressor_labels.py

"""Import algorithm stage labels into compressor_label.

Phase 2: algorithm labels are written as segment records. Manual rows are
never overwritten by a re-import.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import re
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any

from app.db import ensure_label_table, get_connection

ROOT = Path(__file__).resolve().parents[1]
VIBRATION_CSV = ROOT / "models" / "datas" / "reference" / "振动异常段.csv"
OIL_CSV = ROOT / "models" / "datas" / "reference" / "润滑油压力小时阶段段.csv"
TEMPERATURE_CSV = ROOT / "_tmp_temperature" / "out" / "温度异常段.csv"

SEVERITY = {"轻微": 1, "异常": 2, "严重异常": 3}
VIBRATION_SIDE_TO_METRIC = {"联轴器端": "YSJ_10", "链轮端": "YSJ_11"}

INSERT_SQL = """
INSERT INTO compressor_label (
    source_type, origin, target_type, target_id, target_key, channel_key, segment_key,
    unit, device_code, device_part, device_point, metric_key, metric_label,
    start_time, end_time, algorithm_start, algorithm_end,
    label_type, label_subtype, severity, confirm_status, usable_for_train,
    source_file, source_row, algorithm_version, features_json
) VALUES (
    %s, %s, %s, %s, %s, %s, %s,
    %s, %s, %s, %s, %s, %s,
    %s, %s, %s, %s,
    %s, %s, %s, %s, %s,
    %s, %s, %s, %s
)
ON DUPLICATE KEY UPDATE
    label_type = IF(compressor_label.origin = 'manual', compressor_label.label_type, VALUES(label_type)),
    label_subtype = IF(compressor_label.origin = 'manual', compressor_label.label_subtype, VALUES(label_subtype)),
    severity = IF(compressor_label.origin = 'manual', compressor_label.severity, VALUES(severity)),
    confirm_status = IF(compressor_label.origin = 'manual', compressor_label.confirm_status, VALUES(confirm_status)),
    usable_for_train = IF(compressor_label.origin = 'manual', compressor_label.usable_for_train, VALUES(usable_for_train)),
    features_json = IF(compressor_label.origin = 'manual', compressor_label.features_json, VALUES(features_json)),
    updated_at = CURRENT_TIMESTAMP
"""


def read_csv(path: Path) -> list[dict[str, str]]:
    if not path.exists():
        print(f"跳过：文件不存在 {path}")
        return []
    with path.open("r", encoding="utf-8-sig", newline="") as handle:
        return list(csv.DictReader(handle))


def parse_unit(raw: str) -> str | None:
    match = re.match(r"\s*(\d+)\s*号机", raw or "")
    return match.group(1) if match else None


def parse_time(raw: str) -> datetime | None:
    text = (raw or "").strip()
    if not text:
        return None
    for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M", "%Y-%m-%d"):
        try:
            return datetime.strptime(text, fmt)
        except ValueError:
            continue
    return None


def time_text(value: datetime) -> str:
    return value.strftime("%Y-%m-%d %H:%M:%S")


def metric_label(site_points: dict[str, str], unit: str, metric: str) -> str:
    return site_points.get(f"YSJ{unit}_{metric.split('_')[-1]}", metric)


def confirm_status(row: dict[str, str]) -> str:
    return "已确认" if (row.get("是否形成确认等级") or "").strip() == "是" else "未确认"


def make_segment(
    *,
    source_type: str,
    unit: str,
    metric: str,
    start: datetime,
    end: datetime,
    label_type: str,
    label_subtype: str | None,
    row: dict[str, str],
    source_file: str,
    site_points: dict[str, str],
) -> tuple[Any, ...]:
    device_code = f"{unit}#"
    target_key = f"pks:{unit}:{device_code}:{metric}"
    source_row = (row.get("周期编号") or "").strip()
    segment_key = hashlib.md5(
        f"{source_type}|{unit}|{metric}|{source_row}|{time_text(start)}|{time_text(end)}".encode("utf-8")
    ).hexdigest()
    label = metric_label(site_points, unit, metric)
    return (
        source_type,
        "algorithm",
        "pks",
        None,
        target_key,
        target_key,
        segment_key,
        unit,
        device_code,
        (row.get("机组") or "").strip(),
        label,
        metric,
        label,
        time_text(start),
        time_text(end),
        time_text(start),
        time_text(end),
        label_type,
        label_subtype,
        SEVERITY.get(label_type),
        confirm_status(row),
        1,
        source_file,
        source_row,
        "csv-v1",
        json.dumps(row, ensure_ascii=False),
    )


def build_rows(site_points: dict[str, str]) -> list[tuple[Any, ...]]:
    rows: list[tuple[Any, ...]] = []

    for row in read_csv(VIBRATION_CSV):
        unit = parse_unit(row.get("机组", ""))
        start = parse_time(row.get("开始小时", ""))
        end = parse_time(row.get("结束小时", ""))
        side = (row.get("最强异常侧") or "").strip()
        metric = VIBRATION_SIDE_TO_METRIC.get(side)
        if not unit or not start or not end or not metric:
            continue
        rows.append(
            make_segment(
                source_type="vibration",
                unit=unit,
                metric=metric,
                start=start,
                end=end + timedelta(hours=1),
                label_type=(row.get("振动阶段") or "").strip(),
                label_subtype=side or None,
                row=row,
                source_file=str(VIBRATION_CSV.relative_to(ROOT)),
                site_points=site_points,
            )
        )

    for row in read_csv(OIL_CSV):
        unit = parse_unit(row.get("机组", ""))
        start = parse_time(row.get("开始小时", ""))
        end = parse_time(row.get("结束小时", ""))
        if not unit or not start or not end:
            continue
        metric = "YSJ_5" if unit in {"7", "8", "9"} else "YSJ_1"
        rows.append(
            make_segment(
                source_type="oil_pressure",
                unit=unit,
                metric=metric,
                start=start,
                end=end + timedelta(hours=1),
                label_type=(row.get("压力阶段") or "").strip(),
                label_subtype="润滑油压力",
                row=row,
                source_file=str(OIL_CSV.relative_to(ROOT)),
                site_points=site_points,
            )
        )

    for row in read_csv(TEMPERATURE_CSV):
        unit = parse_unit(row.get("机组", ""))
        start = parse_time(row.get("开始小时", ""))
        end = parse_time(row.get("结束小时", ""))
        metric = (row.get("阶段点位") or "").strip()
        if not unit or not start or not end or not re.fullmatch(r"YSJ_\d+", metric):
            continue
        rows.append(
            make_segment(
                source_type="temperature",
                unit=unit,
                metric=metric,
                start=start,
                end=end + timedelta(hours=1),
                label_type=(row.get("温度阶段") or "").strip(),
                label_subtype=(row.get("阶段部位") or row.get("特征类型") or "").strip() or None,
                row=row,
                source_file=str(TEMPERATURE_CSV.relative_to(ROOT)),
                site_points=site_points,
            )
        )

    return [row for row in rows if row[17] in SEVERITY]


def load_site_points() -> dict[str, str]:
    with get_connection() as connection:
        with connection.cursor() as cursor:
            cursor.execute("SELECT ItemName, ItemDescription FROM site_point")
            return {str(row["ItemName"]): str(row["ItemDescription"] or "") for row in cursor.fetchall()}


def main() -> None:
    parser = argparse.ArgumentParser(description="导入算法阶段标签到 compressor_label")
    parser.add_argument("--dry-run", action="store_true", help="只统计，不写库")
    args = parser.parse_args()

    ensure_label_table()
    site_points = load_site_points()
    rows = build_rows(site_points)
    counts: dict[str, int] = {}
    for row in rows:
        counts[row[0]] = counts.get(row[0], 0) + 1
    print(f"待导入 {len(rows)} 条：{counts}")
    if args.dry_run:
        return

    with get_connection() as connection:
        with connection.cursor() as cursor:
            cursor.executemany(INSERT_SQL, rows)
            cursor.execute(
                "SELECT source_type, COUNT(*) AS n FROM compressor_label "
                "WHERE origin = 'algorithm' GROUP BY source_type ORDER BY source_type"
            )
            print("库内算法标签：", {row["source_type"]: int(row["n"]) for row in cursor.fetchall()})


if __name__ == "__main__":
    main()