Compressor/_tmp_cylinder_vibration/adaptive.py

"""第三轮：自适应扩段 + 多特征确认。"""
from __future__ import annotations

from datetime import timedelta

import numpy as np
import pandas as pd

from . import config

RMS_HIGH = 1.10
KURT_HIGH = 1.50
ACTIVE_RATIO = 0.30
MAX_GROW_HOURS = 6.0
MAX_GAP_MINUTES = 60


def add_ratios(frame: pd.DataFrame, baseline: dict) -> pd.DataFrame:
    frame = frame.copy()
    frame["rms_ratio"] = np.nan
    frame["kurt_ratio"] = np.nan
    for cylinder, group in frame.groupby("cylinder", sort=False):
        meta = baseline.get(cylinder, {})
        rms_median = float(meta.get("rms_log", [0.0, 1.0])[0])
        kurt_median = float(meta.get("kurt", [1.0, 1.0])[0])
        frame.loc[group.index, "rms_ratio"] = np.power(10.0, group["rms_log"] - rms_median)
        if kurt_median > 0:
            frame.loc[group.index, "kurt_ratio"] = group["kurt"] / kurt_median
    return frame


def _expand(group: pd.DataFrame, seed_index: int) -> tuple[int, int]:
    seed_time = pd.Timestamp(group["sample_time"].iloc[seed_index])
    left = right = seed_index

    def allowed(index: int) -> bool:
        row = group.iloc[index]
        gap = abs((pd.Timestamp(row["sample_time"]) - seed_time).total_seconds()) / 60.0
        if gap > MAX_GROW_HOURS * 60:
            return False
        neighbor_gap = abs((pd.Timestamp(row["sample_time"]) - pd.Timestamp(group["sample_time"].iloc[right if index > right else left])).total_seconds()) / 60.0
        if index > right:
            neighbor_gap = (pd.Timestamp(row["sample_time"]) - pd.Timestamp(group["sample_time"].iloc[right])).total_seconds() / 60.0
        elif index < left:
            neighbor_gap = (pd.Timestamp(group["sample_time"].iloc[left]) - pd.Timestamp(row["sample_time"])).total_seconds() / 60.0
        if neighbor_gap > MAX_GAP_MINUTES:
            return False
        near_seed = gap <= 60.0
        return (
            near_seed
            or float(row["score"]) >= 3.0
            or float(row["kurt_ratio"]) >= KURT_HIGH
            or float(row["rms_ratio"]) >= RMS_HIGH
        )

    while left - 1 >= 0 and allowed(left - 1):
        left -= 1
    while right + 1 < len(group) and allowed(right + 1):
        right += 1

    start = pd.Timestamp(group["sample_time"].iloc[left])
    end = pd.Timestamp(group["sample_time"].iloc[right])
    if (end - start).total_seconds() / 3600.0 < 1.0:
        start_time = seed_time - timedelta(hours=1)
        end_time = seed_time + timedelta(hours=1)
        part = group[(group["sample_time"] >= start_time) & (group["sample_time"] <= end_time)]
        if not part.empty:
            left = int(part.index[0])
            right = int(part.index[-1])
    return left, right


def build_intervals(frame: pd.DataFrame, point_labels: pd.DataFrame, baseline: dict) -> pd.DataFrame:
    if frame.empty or point_labels.empty:
        return pd.DataFrame()
    if "rms_ratio" not in frame.columns or "kurt_ratio" not in frame.columns:
        frame = add_ratios(frame, baseline)
    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[int, int]] = []
        for seed in seeds.itertuples(index=False):
            seed_time = pd.Timestamp(getattr(seed, "max_time"))
            index = int((group["sample_time"] - seed_time).abs().argmin())
            candidates.append(_expand(group, index))

        candidates.sort()
        merged: list[tuple[int, int]] = []
        for left, right in candidates:
            if merged and pd.Timestamp(group["sample_time"].iloc[left]) <= pd.Timestamp(group["sample_time"].iloc[merged[-1][1]]) + timedelta(hours=2):
                merged[-1] = (merged[-1][0], max(merged[-1][1], right))
            else:
                merged.append((left, right))

        for left, right in merged:
            part = group.iloc[left:right + 1]
            if part.empty:
                continue
            start = pd.Timestamp(part["sample_time"].iloc[0])
            end = pd.Timestamp(part["sample_time"].iloc[-1])
            duration = (end - start).total_seconds() / 3600.0
            if duration < 0.75:
                continue
            max_row = part.loc[part["score"].idxmax()]
            active_rms = float((part["rms_ratio"] >= RMS_HIGH).mean())
            active_kurt = float((part["kurt_ratio"] >= KURT_HIGH).mean())
            max_rms_ratio = float(part["rms_ratio"].max())
            max_kurt_ratio = float(part["kurt_ratio"].max())
            mean_rms_ratio = float(part["rms_ratio"].mean())
            mean_kurt_ratio = float(part["kurt_ratio"].mean())
            high_rms = max_rms_ratio >= RMS_HIGH
            high_kurt = max_kurt_ratio >= KURT_HIGH
            confirmed = high_rms and high_kurt
            if confirmed:
                direction = "高幅高冲击"
            elif high_kurt and not high_rms:
                direction = "低幅高冲击"
            elif high_rms and not high_kurt:
                direction = "高幅"
            else:
                direction = "未确认"
            score = float(max_row["score"])
            candidate_threshold = float(part["threshold"].iloc[0])
            anomaly_threshold = float(part["anomaly_threshold"].iloc[0])
            severe_threshold = float(part["severe_threshold"].iloc[0])
            if score >= severe_threshold:
                level = "严重异常"
                confidence = "高"
            elif score >= anomaly_threshold:
                level = "异常"
                confidence = "中"
            elif score >= candidate_threshold:
                level = "轻微"
                confidence = "低"
            else:
                level = "观察"
                confidence = "低"

            rows.append(
                {
                    "cylinder": cylinder,
                    "stage": config.CYLINDERS[cylinder]["stage"],
                    "start_time": start,
                    "end_time": end,
                    "duration_hours": round(duration, 2),
                    "points": int(len(part)),
                    "max_score": round(float(max_row["score"]), 3),
                    "max_time": pd.Timestamp(max_row["sample_time"]),
                    "driver": str(max_row["driver"]),
                    "active_rms_ratio": round(active_rms, 2),
                    "active_kurt_ratio": round(active_kurt, 2),
                    "max_rms_ratio": round(max_rms_ratio, 3),
                    "max_kurt_ratio": round(max_kurt_ratio, 3),
                    "mean_rms_ratio": round(mean_rms_ratio, 3),
                    "mean_kurt_ratio": round(mean_kurt_ratio, 3),
                    "multi_feature_confirmed": bool(confirmed),
                    "direction": direction,
                    "level": level,
                    "confidence": confidence,
                },
            )

    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) -> pd.DataFrame:
    if intervals.empty:
        return pd.DataFrame()
    rows: list[dict] = []
    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)
        cylinder = event["cylinder"]
        candidates = intervals[
            (intervals["cylinder"] == cylinder)
            & (intervals["start_time"] <= window_end)
            & (intervals["end_time"] >= window_start)
        ]
        if candidates.empty:
            rows.append(
                {
                    "event_id": event["event_id"],
                    "event_date": event["date"],
                    "cylinder": cylinder,
                    "description": event["description"],
                    "hit": False,
                    "confirmed": False,
                    "start_time": None,
                    "end_time": None,
                    "duration_hours": None,
                    "max_score": None,
                    "direction": None,
                    "level": None,
                    "confidence": None,
                    "lead_days": None,
                },
            )
            continue
        confirmed = candidates[candidates["multi_feature_confirmed"]]
        pool = confirmed if not confirmed.empty else candidates
        row = pool.loc[pool["max_score"].idxmax()]
        lead_days = round((event_time - pd.Timestamp(row["start_time"])).total_seconds() / 86400.0, 2)
        rows.append(
            {
                "event_id": event["event_id"],
                "event_date": event["date"],
                "cylinder": cylinder,
                "description": event["description"],
                "hit": True,
                "confirmed": bool(row["multi_feature_confirmed"]),
                "start_time": row["start_time"],
                "end_time": row["end_time"],
                "duration_hours": row["duration_hours"],
                "max_score": row["max_score"],
                "direction": row["direction"],
                "level": row["level"],
                "confidence": row["confidence"],
                "lead_days": lead_days,
            },
        )
    return pd.DataFrame(rows)