Compressor/_tmp_cylinder_vibration/labels.py

"""缸体振动标签：只用振动特征，不看温度、不看故障表。"""
from __future__ import annotations

from datetime import timedelta

import numpy as np
import pandas as pd

from . import config


def prepare_frame(rows: list[dict], interval_minutes: int) -> pd.DataFrame:
    frame = pd.DataFrame(rows)
    if frame.empty:
        return frame
    frame["sample_time"] = pd.to_datetime(frame["sample_time"])
    frame["cylinder"] = frame["point_name"].map(config.POINT_TO_CYLINDER)
    frame["stage"] = frame["cylinder"].map(lambda name: config.CYLINDERS[name]["stage"])
    frame["slot"] = frame["sample_time"].dt.floor(f"{interval_minutes}min")
    frame = frame.dropna(subset=["cylinder"]).sort_values(["cylinder", "sample_time"]).reset_index(drop=True)
    return frame


def event_window_mask(times: pd.Series) -> pd.Series:
    mask = pd.Series(False, index=times.index)
    for event in config.EVENTS:
        center = pd.Timestamp(event["date"])
        start = center - pd.Timedelta(days=config.BASELINE_EXCLUDE_BEFORE_DAYS)
        end = center + pd.Timedelta(days=config.BASELINE_EXCLUDE_AFTER_DAYS)
        mask |= (times >= start) & (times <= end)
    return mask


def _add_residuals(frame: pd.DataFrame) -> pd.DataFrame:
    residuals: pd.DataFrame | None = None
    for feature in config.RESIDUAL_FEATURES:
        pivot = frame.pivot_table(
            index="slot", columns="cylinder", values=feature, aggfunc="median"
        )
        if pivot.empty:
            continue
        pieces: list[pd.DataFrame] = []
        for stage, peers in config.STAGE_PEERS.items():
            available = [name for name in peers if name in pivot.columns]
            if len(available) < 2:
                continue
            for cylinder in available:
                peer = pivot[available].drop(columns=[cylinder]).median(axis=1)
                residual = pivot[cylinder] - peer
                part = residual.rename(f"resid_{feature}").reset_index()
                part["cylinder"] = cylinder
                pieces.append(part)
        if not pieces:
            continue
        feature_frame = pd.concat(pieces, ignore_index=True)
        if residuals is None:
            residuals = feature_frame
        else:
            residuals = residuals.merge(feature_frame, on=["slot", "cylinder"], how="outer")
    if residuals is None:
        return frame
    return frame.merge(residuals, on=["slot", "cylinder"], how="left")


def _normal_mask_for(normal_mask: pd.Series, index: pd.Index) -> pd.Series:
    mask = normal_mask.reindex(index)
    return mask.where(mask.notna(), False).astype(bool)


def _prepare_raw(frame: pd.DataFrame) -> pd.DataFrame:
    frame = frame.copy().sort_values(["cylinder", "sample_time"]).reset_index(drop=True)
    derived = [
        column for column in frame.columns
        if column.startswith("resid_") or column.startswith("z_")
        or column in {"score", "driver", "threshold", "contrast", "hit"}
    ]
    frame = frame.drop(columns=derived, errors="ignore")
    frame = _add_residuals(frame)
    return frame.sort_values(["cylinder", "sample_time"]).reset_index(drop=True)


def _feature_columns() -> list[str]:
    return list(config.SCORE_FEATURES) + [f"resid_{name}" for name in config.RESIDUAL_FEATURES]


def compute_baseline(frame: pd.DataFrame) -> dict[str, dict[str, list[float]]]:
    """按缸、按特征计算正常时段的中位数/MAD，供二次评估复用。"""
    prepared = _prepare_raw(frame)
    normal_mask = ~event_window_mask(prepared["sample_time"])
    columns = list(dict.fromkeys(_feature_columns() + config.PLOT_FEATURES))
    baseline: dict[str, dict[str, list[float]]] = {}
    for cylinder, group in prepared.groupby("cylinder", sort=False):
        normal = group.loc[_normal_mask_for(normal_mask, group.index)]
        if normal.empty:
            normal = group
        baseline[cylinder] = {}
        for column in columns:
            values = normal[column].dropna() if column in normal else pd.Series(dtype=float)
            if values.empty:
                values = group[column].dropna() if column in group else pd.Series(dtype=float)
            if values.empty:
                baseline[cylinder][column] = [0.0, 1.0]
                continue
            median = float(values.median())
            mad = float((values - median).abs().median())
            scale = 1.4826 * mad
            if not np.isfinite(scale) or scale <= 1e-9:
                scale = float(values.std())
            if not np.isfinite(scale) or scale <= 1e-9:
                scale = 1.0
            baseline[cylinder][column] = [median, scale]
    return baseline


def _apply_baseline_z(
    frame: pd.DataFrame,
    columns: list[str],
    baseline: dict[str, dict[str, list[float]]],
) -> None:
    for column in columns:
        if column not in frame.columns:
            frame[f"z_{column}"] = np.nan
            continue
        frame[f"z_{column}"] = np.nan
        for _, group in frame.groupby("cylinder", sort=False):
            median, scale = baseline.get(group["cylinder"].iloc[0], {}).get(column, [0.0, 1.0])
            scale = float(scale) if np.isfinite(scale) and scale > 0 else 1.0
            z = (group[column] - float(median)) / scale
            frame.loc[group.index, f"z_{column}"] = z.clip(-20.0, 20.0)


def _add_rolling_z(frame: pd.DataFrame, columns: list[str], interval_minutes: int) -> None:
    window = max(4, int(round(config.ROLLING_WINDOW_HOURS * 60 / max(interval_minutes, 1))))
    min_periods = max(2, int(round(config.ROLLING_MIN_HOURS * 60 / max(interval_minutes, 1))))
    for column in columns:
        if column not in frame.columns:
            frame[f"z_{column}"] = np.nan
            continue
        frame[f"z_{column}"] = np.nan
        for _, group in frame.groupby("cylinder", sort=False):
            series = group[column]
            median = series.rolling(window, min_periods=min_periods).median().shift(1)
            std = series.rolling(window, min_periods=min_periods).std().shift(1)
            floor = max(float(series.std()), 1e-6) * 0.2
            z = (series - median) / std.clip(lower=floor)
            frame.loc[group.index, f"z_{column}"] = z.clip(-20.0, 20.0)

def build_labels(frame: pd.DataFrame, interval_minutes: int, baseline: dict | None = None) -> tuple[pd.DataFrame, pd.DataFrame, dict[str, float]]:
    """返回带分数的明细、标签段、每缸阈值。"""
    if frame.empty:
        return frame, pd.DataFrame(), {}

    frame = _prepare_raw(frame)
    frame["score"] = 0.0
    frame["driver"] = "none"
    frame["threshold"] = config.MIN_SCORE_Z
    normal_mask = ~event_window_mask(frame["sample_time"])

    columns = _feature_columns()
    z_columns = [f"z_{name}" for name in columns]
    if baseline is None:
        _add_rolling_z(frame, columns, interval_minutes)
    else:
        _apply_baseline_z(frame, columns, baseline)

    available = [column for column in z_columns if column in frame.columns]
    if not available:
        return frame, pd.DataFrame(), {}
    z_values = frame[available].replace([np.inf, -np.inf], np.nan).abs()
    max_scores = z_values.max(axis=1).fillna(0.0)
    frame["score"] = max_scores
    drivers = z_values.fillna(0.0).idxmax(axis=1).str.replace("z_", "", regex=False)
    frame["driver"] = drivers.where(max_scores > 0, "none")

    contrast_window = max(2, int(round(config.CONTRAST_WINDOW_HOURS * 60 / max(interval_minutes, 1))))
    frame["contrast"] = 0.0
    for _, group in frame.groupby("cylinder", sort=False):
        score = group["score"]
        local = score.rolling(contrast_window, min_periods=max(2, contrast_window // 4)).median().shift(1)
        frame.loc[group.index, "contrast"] = (score - local).fillna(0.0)

    thresholds: dict[str, float] = {}
    for cylinder, group in frame.groupby("cylinder", sort=False):
        normal_scores = group.loc[_normal_mask_for(normal_mask, group.index), "score"].dropna()
        quantile = float(normal_scores.quantile(config.SCORE_QUANTILE)) if len(normal_scores) else config.MIN_SCORE_Z
        thresholds[cylinder] = max(config.MIN_SCORE_Z, quantile)
    frame["threshold"] = frame["cylinder"].map(thresholds)
    frame["hit"] = (frame["score"] >= frame["threshold"]) & (frame["contrast"] >= config.MIN_CONTRAST)

    labels = _extract_segments(frame, interval_minutes)
    return frame, labels, thresholds


def _extract_segments(frame: pd.DataFrame, interval_minutes: int) -> pd.DataFrame:
    merge_gap = max(1, int(round(120 / max(interval_minutes, 1))))
    rows: list[dict] = []
    for cylinder, group in frame.groupby("cylinder", sort=False):
        group = group.sort_values("sample_time").reset_index(drop=True)
        threshold = float(group["threshold"].iloc[0])
        stage = config.CYLINDERS[cylinder]["stage"]
        peers = [name for name in config.STAGE_PEERS[stage] if name != cylinder]
        hit_positions = np.flatnonzero(group["hit"].fillna(False).to_numpy())
        if hit_positions.size == 0:
            continue

        groups: list[tuple[int, int]] = []
        start = previous = int(hit_positions[0])
        for position in hit_positions[1:]:
            position = int(position)
            if position - previous <= merge_gap + 1:
                previous = position
            else:
                groups.append((start, previous))
                start = previous = position
        groups.append((start, previous))

        for start, end in groups:
            segment = group.iloc[start:end + 1]
            hit_segment = segment[segment["hit"]]
            if hit_segment.empty:
                continue
            max_index = segment["score"].idxmax()
            max_row = segment.loc[max_index]
            start_time = pd.Timestamp(segment["sample_time"].iloc[0])
            end_time = pd.Timestamp(segment["sample_time"].iloc[-1])
            max_score = float(max_row["score"])
            max_time = pd.Timestamp(max_row["sample_time"])
            peer_rows = frame[
                (frame["cylinder"].isin(peers))
                & (frame["sample_time"] >= max_time - pd.Timedelta(minutes=interval_minutes))
                & (frame["sample_time"] <= max_time + pd.Timedelta(minutes=interval_minutes))
            ]
            peer_score = float(peer_rows["score"].max()) if not peer_rows.empty else None
            specificity = (max_score / peer_score) if peer_score and peer_score > 0 else None
            hit_points = int(len(hit_segment))
            strong_single = (
                hit_points == 1
                and max_score >= config.STRONG_SINGLE_SCORE
                and specificity is not None
                and specificity >= config.STRONG_SINGLE_SPECIFICITY
            )
            if hit_points < config.MIN_CONSECUTIVE_HITS and not strong_single:
                continue

            lookback = start_time - timedelta(days=config.ONSET_LOOKBACK_DAYS)
            onset_candidates = group[
                (group["sample_time"] >= lookback)
                & (group["sample_time"] <= start_time)
                & (group["score"] >= threshold * config.ONSET_FACTOR)
            ]
            onset_time = pd.Timestamp(onset_candidates["sample_time"].iloc[0]) if not onset_candidates.empty else start_time
            lead_days = max(0.0, (start_time - onset_time).total_seconds() / 86400.0)

            if max_score >= max(threshold + 2.0, 6.0) and hit_points >= 4:
                confidence = "高"
            elif hit_points >= 2 or strong_single:
                confidence = "中"
            else:
                confidence = "低"

            rows.append(
                {
                    "cylinder": cylinder,
                    "stage": config.CYLINDERS[cylinder]["stage"],
                    "start_time": start_time,
                    "end_time": end_time,
                    "duration_hours": round((end_time - start_time).total_seconds() / 3600.0, 2),
                    "points": int(len(segment)),
                    "hit_points": hit_points,
                    "max_score": round(max_score, 3),
                    "max_time": max_time,
                    "peer_max_score": None if peer_score is None else round(peer_score, 3),
                    "specificity": None if specificity is None else round(specificity, 2),
                    "driver": str(max_row["driver"]),
                    "onset_time": onset_time,
                    "lead_days": round(lead_days, 2),
                    "threshold": round(threshold, 3),
                    "confidence": confidence,
                },
            )
    if not rows:
        return pd.DataFrame()
    return pd.DataFrame(rows).sort_values(["start_time", "cylinder"]).reset_index(drop=True)


def validate_events(labels: pd.DataFrame, frame: pd.DataFrame, temperature: pd.DataFrame) -> pd.DataFrame:
    rows: list[dict] = []
    event_windows: list[tuple[pd.Timestamp, pd.Timestamp]] = []
    for event in config.EVENTS:
        event_time = pd.Timestamp(event["date"])
        window_start = event_time - pd.Timedelta(days=config.VALIDATION_BEFORE_DAYS)
        window_end = event_time + pd.Timedelta(days=config.VALIDATION_AFTER_DAYS)
        event_windows.append((window_start, window_end))

        cylinder = event["cylinder"]
        candidate = labels[
            (labels["cylinder"] == cylinder)
            & (labels["start_time"] <= window_end)
            & (labels["end_time"] >= window_start)
        ] if not labels.empty else pd.DataFrame()

        if candidate.empty:
            hit = False
            segment = None
            lead_days = None
        else:
            segment = candidate.loc[candidate["max_score"].idxmax()]
            hit = True
            if pd.Timestamp(segment["start_time"]) <= event_time:
                lead_days = round((event_time - pd.Timestamp(segment["start_time"])).total_seconds() / 86400.0, 2)
            else:
                lead_days = -round((pd.Timestamp(segment["start_time"]) - event_time).total_seconds() / 86400.0, 2)

        peer_score = None
        specificity = None
        if segment is not None and not frame.empty:
            stage = config.CYLINDERS[cylinder]["stage"]
            peers = [name for name in config.STAGE_PEERS[stage] if name != cylinder]
            max_time = pd.Timestamp(segment["max_time"])
            nearby = frame[
                (frame["cylinder"].isin(peers))
                & (frame["sample_time"] >= max_time - pd.Timedelta(minutes=15))
                & (frame["sample_time"] <= max_time + pd.Timedelta(minutes=15))
            ]
            if not nearby.empty:
                peer_score = round(float(nearby["score"].max()), 3)
                if peer_score > 0:
                    specificity = round(float(segment["max_score"]) / peer_score, 2)

        temp_row = pd.DataFrame()
        if not temperature.empty:
            temp_row = temperature[
                (temperature["cylinder"] == cylinder)
                & (temperature["start_time"] <= window_end)
                & (temperature["end_time"] >= window_start)
            ].sort_values("start_time")

        rows.append(
            {
                "event_id": event["event_id"],
                "event_date": event["date"],
                "cylinder": cylinder,
                "description": event["description"],
                "vibration_hit": hit,
                "segment_start": None if segment is None else segment["start_time"],
                "segment_end": None if segment is None else segment["end_time"],
                "segment_max_score": None if segment is None else segment["max_score"],
                "segment_driver": None if segment is None else segment["driver"],
                "segment_confidence": None if segment is None else segment["confidence"],
                "lead_days": lead_days,
                "peer_max_score": peer_score,
                "specificity": specificity,
                "temp_reference_start": None if temp_row.empty else temp_row["start_time"].iloc[0],
                "temp_reference_end": None if temp_row.empty else temp_row["end_time"].iloc[0],
                "temp_reference_stage": None if temp_row.empty else " / ".join(temp_row["温度阶段"].astype(str).unique()),
                "temp_reference_confirmed": None if temp_row.empty else " / ".join(temp_row["是否形成确认等级"].astype(str).unique()),
            },
        )

    total_days = max((frame["sample_time"].max() - frame["sample_time"].min()).total_seconds() / 86400.0, 1.0) if not frame.empty else 1.0
    if not labels.empty:
        related = labels.apply(
            lambda row: any(
                pd.Timestamp(row["start_time"]) <= end and pd.Timestamp(row["end_time"]) >= start
                for start, end in event_windows
            ),
            axis=1,
        )
        false_alarms = labels[~related]
    else:
        false_alarms = pd.DataFrame()

    result = pd.DataFrame(rows)
    result.attrs["false_alarm_count"] = int(len(false_alarms))
    result.attrs["false_alarm_per_30d"] = round(len(false_alarms) / total_days * 30.0, 2)
    result.attrs["total_days"] = round(total_days, 1)
    return result