Compressor/_tmp_cylinder_vibration/run_full.py

"""全时间段：15 分钟首周期特征 -> 自适应标签。"""
from __future__ import annotations

import json
import sys
from pathlib import Path

import pandas as pd

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

from _tmp_cylinder_vibration import adaptive, config, data, features, labels, local_labels


def extract_rows(metadata: list[dict], chunk_size: int, tag: str) -> list[dict]:
    rows: list[dict] = []
    total = len(metadata)
    for offset in range(0, total, chunk_size):
        chunk = metadata[offset:offset + chunk_size]
        samples = data.fetch_samples([int(item["id"]) for item in chunk], chunk_size=chunk_size)
        for item in chunk:
            signal = samples.get(int(item["id"]))
            if signal is None:
                continue
            feature = features.extract_features(
                signal,
                float(item.get("sample_frequency_hz") or 0),
                float(item.get("rpm") or 0),
            )
            if feature is None:
                continue
            row = dict(item)
            row.update(feature)
            rows.append(row)
        print(f"[{tag}] {min(offset + chunk_size, total):,}/{total:,} files, features={len(rows):,}", flush=True)
    return rows


def load_or_extract(metadata: list[dict], path: Path, chunk_size: int, refresh: bool) -> list[dict]:
    if path.exists() and not refresh:
        print(f"[full] reuse {path}")
        return pd.read_csv(path).to_dict("records")
    rows = extract_rows(metadata, chunk_size, "full")
    if rows:
        pd.DataFrame(rows).to_csv(path, index=False)
        print(f"[full] wrote {path} ({len(rows):,})")
    return rows


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--refresh", action="store_true")
    parser.add_argument("--chunk-size", type=int, default=200)
    parser.add_argument("--limit-per-cylinder", type=int, default=0)
    args = parser.parse_args()

    metadata = data.fetch_metadata()
    if args.limit_per_cylinder > 0:
        limited = []
        for _, group in pd.DataFrame(metadata).groupby("point_name", sort=False):
            limited.extend(group.head(args.limit_per_cylinder).to_dict("records"))
        metadata = limited
    print(f"[full] running files={len(metadata):,}")

    features_path = config.OUTPUT_DIR / "vibration_features_full.csv"
    rows = load_or_extract(metadata, features_path, args.chunk_size, args.refresh)
    frame = labels.prepare_frame(rows, 15)
    baseline = labels.compute_baseline(frame)
    baseline_path = config.OUTPUT_DIR / "baseline_full.json"
    baseline_path.write_text(json.dumps(baseline, ensure_ascii=False, indent=2), encoding="utf-8")

    frame, point_labels, thresholds = local_labels.build_local_labels(frame)
    frame.to_csv(config.OUTPUT_DIR / "vibration_features_full_local.csv", index=False, date_format="%Y-%m-%d %H:%M:%S")
    point_labels_path = config.OUTPUT_DIR / "vibration_labels_full.csv"
    point_labels.to_csv(point_labels_path, index=False, date_format="%Y-%m-%d %H:%M:%S")

    intervals = adaptive.build_intervals(frame, point_labels, {})
    intervals_path = config.OUTPUT_DIR / "vibration_intervals_full.csv"
    intervals.to_csv(intervals_path, index=False, date_format="%Y-%m-%d %H:%M:%S")

    validation = adaptive.validate_events(intervals)
    validation_path = config.OUTPUT_DIR / "event_validation_full.csv"
    validation.to_csv(validation_path, index=False, date_format="%Y-%m-%d %H:%M:%S")


    summary = {
        "files": len(metadata),
        "feature_rows": len(frame),
        "point_labels": len(point_labels),
        "intervals": len(intervals),
        "confirmed_intervals": int(intervals["multi_feature_confirmed"].sum()) if not intervals.empty else 0,
        "event_validation": validation.to_dict("records"),
        "thresholds": thresholds,
    }
    (config.OUTPUT_DIR / "full_summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2, default=str), encoding="utf-8"
    )
    print(json.dumps({k: summary[k] for k in ("files", "feature_rows", "point_labels", "intervals", "confirmed_intervals")}, ensure_ascii=False, indent=2))
    print(validation.to_string())


if __name__ == "__main__":
    main()