#!/usr/bin/env python3
"""Unit-relative (generalized) vibration features for cross-unit PKS models.

Absolute vibration levels differ between units, so a shared model trained on raw
values treats one unit's normal level as another unit's anomaly.  These features
express each window relative to its own unit's recent robust baseline (the same
7*24-hour baseline the scanner rule uses) and to its own unit's normal
distribution (median/MAD), so one model can be shared across units.

Baselines must be computed from training-split normal windows only.
"""
from __future__ import annotations

from collections import defaultdict
from statistics import median

MIN_BASELINE_SAMPLES = 24
SIDES = ("coupling", "chain")

# Absolute vibration features dropped from the generalized feature set.
ABSOLUTE_FEATURES = tuple(
    f"{side}_{kind}_{stat}"
    for side in SIDES
    for kind in ("avg", "peak")
    for stat in ("mean", "max", "range")
) + tuple(f"{side}_{kind}_slope" for side in SIDES for kind in ("avg", "peak"))

# Fixed feature order for the shared model.
FEATURE_NAMES = tuple(
    name
    for side in SIDES
    for name in (
        f"{side}_avg_dev",
        f"{side}_peak_dev",
        f"{side}_avg_rel_last",
        f"{side}_peak_rel_last",
        f"{side}_avg_rel_mean",
        f"{side}_peak_rel_max",
        f"{side}_avg_z",
        f"{side}_peak_z",
        f"{side}_avg_slope_z",
        f"{side}_high_ratio_mean",
        f"{side}_high_high_ratio_mean",
    )
) + ("both_sides_rising", "running_coverage_mean", "running_coverage_min")


def feature_names() -> list[str]:
    return list(FEATURE_NAMES)


def compute_unit_baselines(samples: list[dict]) -> dict[str, dict]:
    """Compute per-unit normal median/MAD for each side from training normal windows.

    Raises if a unit has fewer than MIN_BASELINE_SAMPLES training normal windows
    or a degenerate MAD, because the normalization would be unstable.
    """
    by_unit = defaultdict(list)
    for sample in samples:
        if sample["label"] != 0:
            continue
        by_unit[sample["unit"]].append(sample["features"])
    if not by_unit:
        raise RuntimeError("训练集没有正常窗口，无法计算机组基准")
    baselines = {}
    for unit, features in sorted(by_unit.items()):
        if len(features) < MIN_BASELINE_SAMPLES:
            raise RuntimeError(
                f"{unit} 训练正常窗口仅 {len(features)} 个，少于 {MIN_BASELINE_SAMPLES} 个，无法计算稳定基准"
            )
        unit_baseline = {"normal_samples": len(features)}
        for side in SIDES:
            for kind in ("avg", "peak"):
                key = f"{side}_{kind}_mean"
                values = [item[key] for item in features if item.get(key) is not None]
                if len(values) < MIN_BASELINE_SAMPLES:
                    raise RuntimeError(f"{unit} {key} 有效正常窗口不足，无法计算稳定基准")
                center = float(median(values))
                mad = float(median([abs(value - center) for value in values]))
                if mad <= 0:
                    raise RuntimeError(f"{unit} {key} 的 MAD=0，无法归一化")
                unit_baseline[f"{side}_{kind}"] = {"median": center, "mad": mad}
        baselines[unit] = unit_baseline
    return baselines


def generalized_features(sample: dict, baselines: dict[str, dict]) -> dict:
    features = sample["features"]
    unit = sample["unit"]
    if unit not in baselines:
        raise KeyError(f"缺少 {unit} 的机组基准")
    unit_baseline = baselines[unit]
    result = {}
    for side in SIDES:
        avg = unit_baseline[f"{side}_avg"]
        peak = unit_baseline[f"{side}_peak"]
        avg_mean = features.get(f"{side}_avg_mean")
        peak_mean = features.get(f"{side}_peak_mean")
        avg_slope = features.get(f"{side}_avg_slope")
        result[f"{side}_avg_dev"] = features.get(f"{side}_avg_baseline_deviation")
        result[f"{side}_peak_dev"] = features.get(f"{side}_peak_baseline_deviation")
        result[f"{side}_avg_rel_last"] = features.get(f"{side}_avg_rel_last")
        result[f"{side}_peak_rel_last"] = features.get(f"{side}_peak_rel_last")
        result[f"{side}_avg_rel_mean"] = features.get(f"{side}_avg_rel_mean")
        result[f"{side}_peak_rel_max"] = features.get(f"{side}_peak_rel_max")
        result[f"{side}_avg_z"] = (avg_mean - avg["median"]) / avg["mad"] if avg_mean is not None else None
        result[f"{side}_peak_z"] = (peak_mean - peak["median"]) / peak["mad"] if peak_mean is not None else None
        result[f"{side}_avg_slope_z"] = avg_slope / avg["mad"] if avg_slope is not None else None
        result[f"{side}_high_ratio_mean"] = features.get(f"{side}_high_ratio_mean")
        result[f"{side}_high_high_ratio_mean"] = features.get(f"{side}_high_high_ratio_mean")
    result["both_sides_rising"] = features.get("both_sides_rising")
    result["running_coverage_mean"] = features.get("running_coverage_mean")
    result["running_coverage_min"] = features.get("running_coverage_min")
    return result


def feature_matrix(samples: list[dict], baselines: dict[str, dict]) -> list[dict]:
    return [generalized_features(sample, baselines) for sample in samples]