Compressor/_tmp_cylinder_vibration/run.py

"""5 号机缸体振动：粗筛 -> 二次细评 -> 事件验证。

用法：
  python3 _tmp_cylinder_vibration/run.py --stage all
  python3 _tmp_cylinder_vibration/run.py --stage coarse
  python3 _tmp_cylinder_vibration/run.py --stage all --refresh
"""
from __future__ import annotations

import argparse
import json
import sys
from datetime import datetime
from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.patches import Patch

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

from _tmp_cylinder_vibration import config, data, features, labels


def parse_time(value: str | None) -> datetime | None:
    if not value:
        return None
    return datetime.fromisoformat(value)


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)
        done = min(offset + chunk_size, total)
        print(f"[{tag}] {done:,}/{total:,} files, features={len(rows):,}", flush=True)
    return rows


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


def metadata_for_windows(
    windows: dict[str, list[tuple[datetime, datetime]]],
    chunk_size: int,
) -> list[dict]:
    rows: list[dict] = []
    for cylinder, intervals in windows.items():
        start = min(item[0] for item in intervals)
        end = max(item[1] for item in intervals)
        items = data.fetch_metadata([cylinder], start, end)
        for item in items:
            sample_time = pd.Timestamp(item["sample_time"]).to_pydatetime()
            if any(window_start <= sample_time <= window_end for window_start, window_end in intervals):
                rows.append(item)
        print(f"[refine] {cylinder}: windows={len(intervals)}, files={len(items):,}", flush=True)
    unique: dict[int, dict] = {}
    for item in rows:
        unique[int(item["id"])] = item
    return list(unique.values())


def write_frame(frame: pd.DataFrame, path: Path) -> None:
    if frame.empty:
        frame.to_csv(path, index=False)
        return
    frame.to_csv(path, index=False, date_format="%Y-%m-%d %H:%M:%S")


def _plot_gap_aware(ax, x, y, gap_minutes: int = 60, **kwargs) -> None:
    x = pd.Series(pd.to_datetime(x)).reset_index(drop=True)
    y = pd.Series(np.asarray(y, dtype=float)).reset_index(drop=True)
    gaps = x.diff() > pd.Timedelta(minutes=gap_minutes)
    y = y.mask(gaps)
    ax.plot(x, y, **kwargs)


def plot_events(frame: pd.DataFrame, vibration_labels: pd.DataFrame, temperature: pd.DataFrame) -> None:
    config.PLOT_DIR.mkdir(parents=True, exist_ok=True)
    for event in config.EVENTS:
        cylinder = event["cylinder"]
        center = pd.Timestamp(event["date"])
        window_start = center - pd.Timedelta(days=10)
        window_end = center + pd.Timedelta(days=3)
        part = frame[
            (frame["cylinder"] == cylinder)
            & (frame["sample_time"] >= window_start)
            & (frame["sample_time"] <= window_end)
        ].sort_values("sample_time")
        if part.empty:
            continue

        fig, axes = plt.subplots(3, 1, figsize=(14, 9), sharex=True)
        _plot_gap_aware(axes[0], part["sample_time"], part["score"], color="#c0392b", linewidth=1.1, label="score")
        if "threshold" in part:
            _plot_gap_aware(axes[0], part["sample_time"], part["threshold"], color="#7f8c8d", linestyle="--", label="threshold")
        axes[0].axvline(center, color="#2c3e50", linestyle="-", linewidth=1.2, label="fault date")
        axes[0].set_ylabel("anomaly score")
        axes[0].legend(loc="upper left", fontsize=8)
        axes[0].grid(alpha=0.2)

        if not vibration_labels.empty:
            segments = vibration_labels[vibration_labels["cylinder"] == cylinder]
            for row in segments.itertuples(index=False):
                if pd.Timestamp(row.end_time) < window_start or pd.Timestamp(row.start_time) > window_end:
                    continue
                color = {"高": "#e74c3c", "中": "#e67e22", "低": "#f1c40f"}.get(row.confidence, "#95a5a6")
                if pd.Timestamp(row.start_time) == pd.Timestamp(row.end_time):
                    axes[0].axvline(row.start_time, color=color, linestyle=":", linewidth=1.5, alpha=0.8)
                else:
                    axes[0].axvspan(row.start_time, row.end_time, color=color, alpha=0.15)

        if not temperature.empty:
            temp = temperature[
                (temperature["cylinder"] == cylinder)
                & (temperature["end_time"] >= window_start)
                & (temperature["start_time"] <= window_end)
            ]
            for row in temp.itertuples(index=False):
                axes[0].axvspan(row.start_time, row.end_time, color="#7f8c8d", alpha=0.12, hatch="//")

        for column, color in (("z_resid_rms_log", "#2980b9"), ("z_resid_kurt", "#8e44ad")):
            if column in part:
                _plot_gap_aware(axes[1], part["sample_time"], part[column], linewidth=1.0, color=color, label=column)
        axes[1].axhline(0, color="#bdc3c7", linewidth=0.8)
        axes[1].set_ylabel("robust z")
        axes[1].legend(loc="upper left", fontsize=8)
        axes[1].grid(alpha=0.2)

        for column, color in (("z_resid_env_1x_log", "#16a085"), ("z_resid_hf_ratio", "#d35400")):
            if column in part:
                _plot_gap_aware(axes[2], part["sample_time"], part[column], linewidth=1.0, color=color, label=column)
        axes[2].axhline(0, color="#bdc3c7", linewidth=0.8)
        axes[2].set_ylabel("robust z")
        axes[2].set_xlabel("time")
        axes[2].legend(loc="upper left", fontsize=8)
        axes[2].grid(alpha=0.2)

        index = config.CYLINDER_TO_INDEX[cylinder]
        handles, labels = axes[0].get_legend_handles_labels()
        handles.append(Patch(facecolor="#e67e22", alpha=0.25, label="vibration label"))
        labels.append("vibration label")
        handles.append(Patch(facecolor="#7f8c8d", alpha=0.2, hatch="//", label="temperature reference"))
        labels.append("temperature reference")
        axes[0].legend(handles, labels, loc="upper left", fontsize=8)
        axes[0].set_title(f"{event['event_id']} cylinder#{index} (gaps = no running data)")
        fig.tight_layout()
        path = config.PLOT_DIR / f"event_{event['event_id']}_cyl{index}.png"
        fig.savefig(path, dpi=130)
        plt.close(fig)
        print(f"[plot] {path}", flush=True)


def plot_heatmap(frame: pd.DataFrame) -> None:
    if frame.empty:
        return
    daily = frame.assign(day=frame["sample_time"].dt.date).pivot_table(
        index="day", columns="cylinder", values="score", aggfunc="max"
    )
    if daily.empty:
        return
    ordered = [name for name in config.CYLINDERS if name in daily.columns]
    daily = daily[ordered]
    fig, ax = plt.subplots(figsize=(13, 5))
    image = ax.imshow(daily.T.to_numpy(), aspect="auto", cmap="YlOrRd", interpolation="nearest")
    ax.set_yticks(range(len(ordered)))
    ax.set_yticklabels([f"cyl{config.CYLINDER_TO_INDEX[name]}" for name in ordered])
    ax.set_xticks(range(0, len(daily.index), max(1, len(daily.index) // 12)))
    ax.set_xticklabels([str(daily.index[i]) for i in range(0, len(daily.index), max(1, len(daily.index) // 12))], rotation=45, ha="right")
    ax.set_title("daily max vibration anomaly score")
    fig.colorbar(image, ax=ax, fraction=0.025, pad=0.02)
    fig.tight_layout()
    path = config.PLOT_DIR / "score_heatmap.png"
    fig.savefig(path, dpi=130)
    plt.close(fig)
    print(f"[plot] {path}", flush=True)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--stage", choices=("coarse", "refine", "all"), default="all")
    parser.add_argument("--coarse-minutes", type=int, default=60)
    parser.add_argument("--chunk-size", type=int, default=200)
    parser.add_argument("--limit-per-cylinder", type=int, default=0)
    parser.add_argument("--start", default=None)
    parser.add_argument("--end", default=None)
    parser.add_argument("--refresh", action="store_true")
    args = parser.parse_args()

    config.OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    start = parse_time(args.start)
    end = parse_time(args.end)
    summary: dict = {
        "unit": config.UNIT,
        "coarse_minutes": args.coarse_minutes,
        "stage": args.stage,
    }

    coarse_features_path = config.OUTPUT_DIR / "vibration_features_coarse.csv"
    coarse_labels_path = config.OUTPUT_DIR / "vibration_labels_coarse.csv"
    refined_features_path = config.OUTPUT_DIR / "vibration_features_refined.csv"
    refined_labels_path = config.OUTPUT_DIR / "vibration_labels_refined.csv"
    validation_path = config.OUTPUT_DIR / "event_validation_refined.csv"
    baseline_path = config.OUTPUT_DIR / "baseline.json"
    summary_path = config.OUTPUT_DIR / "run_summary.json"

    coarse_labels = pd.DataFrame()
    if args.stage in ("coarse", "all"):
        print("[coarse] fetch metadata ...", flush=True)
        metadata = data.fetch_metadata(start=start, end=end)
        coarse_metadata = data.coarse_metadata(metadata, args.coarse_minutes)
        if args.limit_per_cylinder > 0:
            limited = []
            for _, group in pd.DataFrame(coarse_metadata).groupby("point_name", sort=False):
                limited.extend(group.head(args.limit_per_cylinder).to_dict("records"))
            coarse_metadata = limited
        print(f"[coarse] files={len(coarse_metadata):,} (all={len(metadata):,})", flush=True)
        rows = load_or_extract(coarse_metadata, coarse_features_path, args.chunk_size, "coarse", args.refresh)
        coarse_frame = labels.prepare_frame(rows, args.coarse_minutes)
        coarse_baseline = labels.compute_baseline(coarse_frame)
        baseline_path.write_text(json.dumps(coarse_baseline, ensure_ascii=False, indent=2), encoding="utf-8")
        coarse_frame, coarse_labels, coarse_thresholds = labels.build_labels(
            coarse_frame, args.coarse_minutes, baseline=coarse_baseline
        )
        write_frame(coarse_frame, coarse_features_path)
        write_frame(coarse_labels, coarse_labels_path)
        summary.update(
            {
                "coarse_files": len(coarse_metadata),
                "coarse_feature_rows": len(coarse_frame),
                "coarse_labels": len(coarse_labels),
                "coarse_thresholds": coarse_thresholds,
                "baseline": str(baseline_path),
            }
        )
        print(f"[coarse] labels={len(coarse_labels):,}", flush=True)

    if args.stage in ("refine", "all"):
        if coarse_labels.empty and coarse_labels_path.exists():
            coarse_labels = pd.read_csv(coarse_labels_path)
        # 二次评估固定围绕 3 次事件取窗口；粗筛候选只作为阶段一结果保留。
        windows = data.refine_windows(pd.DataFrame(), config.EVENTS)
        print(f"[refine] windows: {[(name, len(items)) for name, items in windows.items()]}", flush=True)
        refine_metadata = metadata_for_windows(windows, args.chunk_size)
        print(f"[refine] files={len(refine_metadata):,}", flush=True)
        rows = load_or_extract(refine_metadata, refined_features_path, args.chunk_size, "refine", args.refresh)
        refined_frame = labels.prepare_frame(rows, 15)
        baseline = None
        if baseline_path.exists():
            baseline = json.loads(baseline_path.read_text(encoding="utf-8"))
            print(f"[refine] baseline {baseline_path}", flush=True)
        refined_frame, refined_labels, refined_thresholds = labels.build_labels(refined_frame, 15, baseline=baseline)
        write_frame(refined_frame, refined_features_path)
        write_frame(refined_labels, refined_labels_path)

        temperature = data.load_temperature_reference()
        validation = labels.validate_events(refined_labels, refined_frame, temperature)
        write_frame(validation, validation_path)
        summary.update(
            {
                "refine_files": len(refine_metadata),
                "refine_feature_rows": len(refined_frame),
                "refine_labels": len(refined_labels),
                "refine_thresholds": refined_thresholds,
                "event_hits": int(validation["vibration_hit"].sum()) if not validation.empty else 0,
                "false_alarms": int(validation.attrs.get("false_alarm_count", 0)),
                "false_alarms_per_30d": float(validation.attrs.get("false_alarm_per_30d", 0.0)),
                "event_validation": validation.to_dict("records"),
            }
        )
        plot_events(refined_frame, refined_labels, temperature)
        plot_heatmap(refined_frame)

    summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2, default=str), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False, indent=2, default=str))
    print(f"output: {config.OUTPUT_DIR}")


if __name__ == "__main__":
    main()