Compressor/models/common/verify.py

#!/usr/bin/env python3
"""Automated verification of the reconstructed algorithms and 3h models.

Checks
------
1. Algorithm scan CSVs are byte-identical to the golden reference.
2. DB full rerun of the hourly intermediate matches the cached table.
3. Pressure / vibration 3h model metrics match the documented results.
4. Dataset candidate and split counts match the documented results.

Writes ``models/results/verification/verification_report.json`` and exits
non-zero if any check fails.
"""
from __future__ import annotations

import argparse
import csv
import json
import sys
from collections import Counter
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from models.common import alarm_store, devices, paths

ALGORITHM_FILES = {
    "pressure": ["润滑油压力小时等级.csv", "润滑油压力小时阶段段.csv", "润滑油压力演化周期.csv"],
    "vibration": ["振动异常段.csv"],
    "common": ["报警阈值配置.csv"],
}

EXPECTED = {
    "pressure": {
        "candidates_normal": 12789, "candidates_anomaly": 1016,
        "splits": {"train": (9095, 8439, 656), "validation": (2457, 2256, 201), "test": (2253, 2094, 159)},
        "threshold": 0.8625534176826477,
        "anomaly": {"precision": 0.9512, "recall": 0.9811, "f1": 0.9659},
        "normal_f1": 0.9974,
        "pr_auc": 0.9961,
        "event_recall": (10, 10),
        "by_unit_normal_fp_rate": {"7号机": 0.0032, "8号机": 0.0113, "9号机": 0.0},
    },
    "vibration": {
        "candidates_normal": 13908, "candidates_anomaly": 231,
        "splits": {"train": (9302, 9112, 190), "validation": (2657, 2637, 20), "test": (2180, 2159, 21)},
        "threshold": 0.9166295528411865,
        "anomaly": {"precision": 0.7143, "recall": 0.7143, "f1": 0.7143},
        "normal_f1": 0.9972,
        "pr_auc": 0.8270,
        "event_recall": (4, 4),
        "by_unit_normal_fp_rate": {"7号机": 0.0, "8号机": 0.0071, "9号机": 0.0013},
    },
}

METRIC_TOLERANCE = 1e-3
THRESHOLD_TOLERANCE = 1e-6


def _csv_header_rows(path: Path):
    with path.open(encoding="utf-8-sig", newline="") as handle:
        reader = csv.DictReader(handle)
        return list(reader.fieldnames or []), list(reader)


def compare_files(reference: Path, candidate: Path) -> dict:
    """Byte-compare when identical; otherwise compare the reference columns as a subset.

    The model-enabled run appends model-reference columns, so the base algorithm
    columns are compared while extra columns are allowed.
    """
    base = {"reference": str(reference), "candidate": str(candidate)}
    if reference.read_bytes() == candidate.read_bytes():
        return {**base, "identical": True, "mode": "byte"}
    ref_header, ref_rows = _csv_header_rows(reference)
    cand_header, cand_rows = _csv_header_rows(candidate)
    missing = [column for column in ref_header if column not in cand_header]
    extra = [column for column in cand_header if column not in ref_header]
    if missing or len(ref_rows) != len(cand_rows):
        return {**base, "identical": False, "mode": "subset", "missing_columns": missing,
                "row_counts": [len(ref_rows), len(cand_rows)]}
    identical = all(
        ref_rows[index].get(column) == cand_rows[index].get(column)
        for index in range(len(ref_rows))
        for column in ref_header
    )
    return {**base, "identical": identical, "mode": "subset", "extra_columns": extra}


def check_algorithm():
    checks = []
    for group, files in ALGORITHM_FILES.items():
        for name in files:
            reference = paths.REFERENCE / name
            candidate = (paths.COMMON_RESULTS if group == "common" else paths.RESULTS / group) / name
            checks.append({"group": group, "file": name, **compare_files(reference, candidate)})
    return {"passed": all(item["identical"] for item in checks), "checks": checks}


def manifest_candidate_counts(dataset_dir: Path) -> dict:
    counts = {"normal": 0, "anomaly": 0}
    with (dataset_dir / "manifest.csv").open(encoding="utf-8-sig", newline="") as handle:
        for row in csv.DictReader(handle):
            if row["label"] == "0":
                counts["normal"] += 1
            elif row["label"] == "1":
                counts["anomaly"] += 1
    return counts


def check_model(model_type: str):
    expected = EXPECTED[model_type]
    dataset_dir = paths.DATAS / model_type
    models_dir = paths.MODEL_ARTIFACTS / model_type
    results_dir = paths.RESULTS / model_type

    candidates = manifest_candidate_counts(dataset_dir)
    split_report = json.loads((dataset_dir / "split_report.json").read_text(encoding="utf-8"))[model_type]
    metadata = json.loads((models_dir / f"{model_type}_metadata.json").read_text(encoding="utf-8"))
    evaluation = json.loads((results_dir / f"{model_type}_test_evaluation.json").read_text(encoding="utf-8"))
    summary = evaluation["window_metrics"]

    details = []

    def check(name, actual, want, tolerance):
        passed = abs(actual - want) <= tolerance if isinstance(want, float) else actual == want
        details.append({"name": name, "actual": actual, "expected": want, "tolerance": tolerance, "passed": bool(passed)})

    check("candidates_normal", candidates["normal"], expected["candidates_normal"], 0)
    check("candidates_anomaly", candidates["anomaly"], expected["candidates_anomaly"], 0)
    for split, (total, normal, anomaly) in expected["splits"].items():
        got = split_report["sample_counts"][split]
        check(f"{split}_total", got["samples_total"], total, 0)
        check(f"{split}_normal", got["samples_normal"], normal, 0)
        check(f"{split}_anomaly", got["samples_anomaly"], anomaly, 0)
    check("threshold", metadata["threshold"], expected["threshold"], THRESHOLD_TOLERANCE)
    check("anomaly_precision", summary["class_metrics"]["anomaly"]["precision"], expected["anomaly"]["precision"], METRIC_TOLERANCE)
    check("anomaly_recall", summary["class_metrics"]["anomaly"]["recall"], expected["anomaly"]["recall"], METRIC_TOLERANCE)
    check("anomaly_f1", summary["class_metrics"]["anomaly"]["f1"], expected["anomaly"]["f1"], METRIC_TOLERANCE)
    check("normal_f1", summary["class_metrics"]["normal"]["f1"], expected["normal_f1"], METRIC_TOLERANCE)
    check("pr_auc", summary["pr_auc"], expected["pr_auc"], METRIC_TOLERANCE)
    detected, total = expected["event_recall"]
    check("event_recall", evaluation["event_metrics"]["detected_anomaly_events"], detected, 0)
    check("event_total", evaluation["event_metrics"]["anomaly_events"], total, 0)
    for unit, want in expected["by_unit_normal_fp_rate"].items():
        got = 1 - evaluation["by_unit"][unit]["specificity"]
        check(f"{unit}_normal_fp_rate", got, want, METRIC_TOLERANCE)

    return {"passed": all(item["passed"] for item in details), "checks": details}


def check_db_parity():
    report_path = paths.RESULTS / "verification" / "db_parity.json"
    if not report_path.is_file():
        return {"passed": False, "reason": "db_parity.json 缺失"}
    report = json.loads(report_path.read_text(encoding="utf-8"))
    return {"passed": bool(report.get("identical")), "report": report}


def _reference_counts(path: Path, batch_to_device: dict[int, str]) -> Counter:
    counts: Counter = Counter()
    if not path.is_file():
        return counts
    with path.open(encoding="utf-8-sig", newline="") as handle:
        for row in csv.DictReader(handle):
            counts[batch_to_device[int(row["批次"])]] += 1
    return counts


def check_alarm_store():
    """Compare compressor_alarm rows to the reference stage CSVs (1 row per segment)."""
    device_list = list(devices.DEVICE_CODES)
    device_batches = devices.load_device_batches(device_list)
    batch_device = devices.batch_to_device(device_batches)
    pressure_ref = _reference_counts(paths.REFERENCE / "润滑油压力小时阶段段.csv", batch_device)
    vibration_ref = _reference_counts(paths.REFERENCE / "振动异常段.csv", batch_device)
    expected: Counter = Counter()
    for device in device_list:
        expected[(device, "pressure")] = pressure_ref.get(device, 0)
        expected[(device, "vibration")] = vibration_ref.get(device, 0)

    rows = alarm_store.fetch_alarms(device_list)
    actual: Counter = Counter()
    for row in rows:
        kind = "pressure" if row["alarm_type"] == alarm_store.OIL_ALARM_TYPE else "vibration"
        actual[(row["device_code"], kind)] += 1

    checks = []
    for key in sorted(set(expected) | set(actual), key=str):
        checks.append({"name": f"{key[0]}_{key[1]}", "actual": actual.get(key, 0), "expected": expected.get(key, 0),
                       "passed": actual.get(key, 0) == expected.get(key, 0)})
    status_ok = all(int(row["status"]) == 1 for row in rows)
    json_ok = True
    des_ok = True
    for row in rows:
        try:
            json.loads(row["alarm_info_json"])
        except (TypeError, ValueError):
            json_ok = False
        if len(row["alarm_des"] or "") > 4000:
            des_ok = False
    checks.append({"name": "status_all_1", "actual": status_ok, "expected": True, "passed": status_ok})
    checks.append({"name": "json_valid", "actual": json_ok, "expected": True, "passed": json_ok})
    checks.append({"name": "des_within_4000", "actual": des_ok, "expected": True, "passed": des_ok})
    return {"passed": all(item["passed"] for item in checks), "checks": checks, "rows": len(rows)}


def main() -> None:
    parser = argparse.ArgumentParser(description="PKS 重构算法/模型自动校验")
    parser.add_argument("--with-db", action="store_true", help="额外校验 compressor_alarm 与参考 CSV 的对应关系")
    args = parser.parse_args()

    report = {
        "algorithm": check_algorithm(),
        "db_parity": check_db_parity(),
        "models": {mt: check_model(mt) for mt in ("pressure", "vibration")},
    }
    if args.with_db:
        report["alarm_store"] = check_alarm_store()
    report["passed"] = (
        report["algorithm"]["passed"]
        and report["db_parity"]["passed"]
        and all(item["passed"] for item in report["models"].values())
        and report.get("alarm_store", {"passed": True})["passed"]
    )
    output = paths.RESULTS / "verification"
    output.mkdir(parents=True, exist_ok=True)
    (output / "verification_report.json").write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")

    print(f"算法扫描结果一致性：{'通过' if report['algorithm']['passed'] else '失败'}")
    print(f"数据库全量重跑一致性：{'通过' if report['db_parity']['passed'] else '失败'}")
    for model_type, result in report["models"].items():
        failed = [item["name"] for item in result["checks"] if not item["passed"]]
        print(f"{model_type} 模型复现：{'通过' if result['passed'] else '失败'}" + (f" 差异={failed}" if failed else ""))
    if "alarm_store" in report:
        failed = [item["name"] for item in report["alarm_store"]["checks"] if not item["passed"]]
        print(f"compressor_alarm 对应：{'通过' if report['alarm_store']['passed'] else '失败'}"
              + (f" 差异={failed}" if failed else "") + f"（{report['alarm_store']['rows']} 行）")
    print(f"总体：{'全部通过' if report['passed'] else '存在失败项'}")
    print(f"报告：{output / 'verification_report.json'}")
    if not report["passed"]:
        sys.exit(1)


if __name__ == "__main__":
    main()