Compressor/_tmp_cylinder_vibration/intervals.py

"""第二轮：把单点候选扩展成连续时间段，并做事件验证。"""
from __future__ import annotations

from datetime import timedelta

import numpy as np
import pandas as pd

from . import config


def _level(max_score: float, active_ratio: float, duration_hours: float) -> str:
    if max_score >= 6.0 and active_ratio >= 0.5 and duration_hours >= 2.0:
        return "严重"
    if max_score >= 4.5 and active_ratio >= 0.3:
        return "异常"
    return "轻微"


def _confidence(max_score: float, active_ratio: float, duration_hours: float) -> str:
    if max_score >= 6.0 and active_ratio >= 0.5 and duration_hours >= 2.0:
        return "高"
    if max_score >= 4.5 and active_ratio >= 0.3:
        return "中"
    return "低"


def build_intervals(
    frame: pd.DataFrame,
    point_labels: pd.DataFrame,
    before_hours: float = 1.0,
    after_hours: float = 1.0,
    merge_gap_hours: float = 2.0,
    min_duration_hours: float = 0.75,
) -> pd.DataFrame:
    """每个单点候选向左右扩展 before/after 小时，再按缸合并相邻段。"""
    if point_labels.empty or frame.empty:
        return pd.DataFrame()

    frame = frame.copy()
    frame["sample_time"] = pd.to_datetime(frame["sample_time"])
    rows: list[dict] = []
    for cylinder, group in frame.groupby("cylinder", sort=False):
        group = group.sort_values("sample_time").reset_index(drop=True)
        seeds = point_labels[point_labels["cylinder"] == cylinder]
        if seeds.empty:
            continue
        candidates: list[tuple[pd.Timestamp, pd.Timestamp, pd.Series]] = []
        for seed in seeds.itertuples(index=False):
            seed_time = pd.Timestamp(getattr(seed, "max_time"))
            start = seed_time - timedelta(hours=before_hours)
            end = seed_time + timedelta(hours=after_hours)
            part = group[(group["sample_time"] >= start) & (group["sample_time"] <= end)]
            if part.empty:
                part = group.iloc[[int((group["sample_time"] - seed_time).abs().argmin())]]
            # 不足 1 小时时，扩大到前后 1.5 小时，保证形成可解释的时间段。
            if (pd.Timestamp(part["sample_time"].iloc[-1]) - pd.Timestamp(part["sample_time"].iloc[0])).total_seconds() / 3600.0 < 1.0:
                part = group[
                    (group["sample_time"] >= seed_time - timedelta(hours=1.5))
                    & (group["sample_time"] <= seed_time + timedelta(hours=1.5))
                ]
            if part.empty:
                continue
            candidates.append((pd.Timestamp(part["sample_time"].iloc[0]), pd.Timestamp(part["sample_time"].iloc[-1]), seed))

        if not candidates:
            continue
        candidates.sort(key=lambda item: item[0])
        merged: list[tuple[pd.Timestamp, pd.Timestamp, list[pd.Series]]] = []
        for start, end, seed in candidates:
            if merged and start <= merged[-1][1] + timedelta(hours=merge_gap_hours):
                merged[-1] = (merged[-1][0], max(merged[-1][1], end), merged[-1][2] + [seed])
            else:
                merged.append((start, end, [seed]))

        for start, end, seed_rows in merged:
            part = group[(group["sample_time"] >= start) & (group["sample_time"] <= end)]
            if part.empty:
                continue
            actual_start = pd.Timestamp(part["sample_time"].iloc[0])
            actual_end = pd.Timestamp(part["sample_time"].iloc[-1])
            duration = (actual_end - actual_start).total_seconds() / 3600.0
            if duration < min_duration_hours:
                continue
            max_row = part.loc[part["score"].idxmax()]
            max_score = float(max_row["score"])
            mean_score = float(part["score"].mean())
            active_points = int((part["score"] >= 3.0).sum())
            active_ratio = active_points / max(len(part), 1)
            seed_best = max(seed_rows, key=lambda item: float(item.max_score))
            rows.append(
                {
                    "cylinder": cylinder,
                    "stage": config.CYLINDERS[cylinder]["stage"],
                    "start_time": actual_start,
                    "end_time": actual_end,
                    "duration_hours": round(duration, 2),
                    "points": int(len(part)),
                    "mean_score": round(mean_score, 3),
                    "max_score": round(max_score, 3),
                    "active_points": active_points,
                    "active_ratio": round(active_ratio, 2),
                    "max_time": pd.Timestamp(max_row["sample_time"]),
                    "driver": str(max_row["driver"]),
                    "seed_max_score": round(float(seed_best.max_score), 3),
                    "seed_max_time": pd.Timestamp(seed_best.max_time),
                    "level": _level(max_score, active_ratio, duration),
                    "confidence": _confidence(max_score, active_ratio, duration),
                },
            )

    if not rows:
        return pd.DataFrame()
    return pd.DataFrame(rows).sort_values(["start_time", "cylinder"]).reset_index(drop=True)


def validate_events(intervals: 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 = intervals[
            (intervals["cylinder"] == cylinder)
            & (intervals["start_time"] <= window_end)
            & (intervals["end_time"] >= window_start)
        ] if not intervals.empty else pd.DataFrame()

        if candidate.empty:
            hit = False
            interval = None
        else:
            interval = candidate.loc[candidate["max_score"].idxmax()]
            hit = True

        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")

        if interval is None:
            rows.append(
                {
                    "event_id": event["event_id"],
                    "event_date": event["date"],
                    "cylinder": cylinder,
                    "description": event["description"],
                    "vibration_hit": False,
                    "interval_start": None,
                    "interval_end": None,
                    "duration_hours": None,
                    "mean_score": None,
                    "max_score": None,
                    "driver": None,
                    "level": None,
                    "confidence": None,
                    "lead_days": None,
                    "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()),
                },
            )
            continue

        if pd.Timestamp(interval["start_time"]) <= event_time:
            lead_days = round((event_time - pd.Timestamp(interval["start_time"])).total_seconds() / 86400.0, 2)
        else:
            lead_days = -round((pd.Timestamp(interval["start_time"]) - event_time).total_seconds() / 86400.0, 2)

        rows.append(
            {
                "event_id": event["event_id"],
                "event_date": event["date"],
                "cylinder": cylinder,
                "description": event["description"],
                "vibration_hit": True,
                "interval_start": interval["start_time"],
                "interval_end": interval["end_time"],
                "duration_hours": interval["duration_hours"],
                "mean_score": interval["mean_score"],
                "max_score": interval["max_score"],
                "driver": interval["driver"],
                "level": interval["level"],
                "confidence": interval["confidence"],
                "lead_days": lead_days,
                "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()),
            },
        )

    result = pd.DataFrame(rows)
    if not intervals.empty:
        related = intervals.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 = intervals[~related]
        false_hours = float(false_alarms["duration_hours"].sum()) if not false_alarms.empty else 0.0
    else:
        false_alarms = pd.DataFrame()
        false_hours = 0.0
    result.attrs["false_alarm_intervals"] = int(len(false_alarms))
    result.attrs["false_alarm_hours"] = round(false_hours, 2)
    return result