"""Model prediction as a *reference* for fault segments (no verdict).

For each fault segment produced by the rule algorithm, slide a 3-hour valid-run
window ending at every valid hour of the segment and score it with the existing
XGBoost model.  Only aggregate probabilities are reported (mean / max / min);
the model never decides whether a fault exists.

Training is untouched: this module only reuses the same window builder and
feature functions so the scored features match training exactly.
"""
from __future__ import annotations

import json
from pathlib import Path

import joblib
import numpy as np

from . import paths, windowing

WINDOW_HOURS = 3
MODEL_FIELDS = ["模型异常概率均值", "模型异常概率最大值", "模型异常概率最小值"]
NOTE_FIELD = "模型预测参考"

_BUNDLE_CACHE: dict[str, tuple] = {}


def _assets(kind: str):
    if kind == "pressure":
        from models.generalized import features_pressure
        from models.generalized.build_datasets import pressure_features

        return features_pressure, pressure_features
    from models.vibration import features_vibration
    from models.vibration.build_datasets import vibration_features

    return features_vibration, vibration_features


def load_bundle(kind: str):
    if kind in _BUNDLE_CACHE:
        return _BUNDLE_CACHE[kind]
    model_dir = paths.PRESSURE_MODELS if kind == "pressure" else paths.VIBRATION_MODELS
    model_path = model_dir / f"{kind}_model.joblib"
    meta_path = model_dir / f"{kind}_metadata.json"
    if not model_path.is_file() or not meta_path.is_file():
        pipeline = "generalized" if kind == "pressure" else "vibration"
        raise RuntimeError(
            f"未找到 {kind} 模型产物：{model_dir}。请先运行 "
            f"python models/{pipeline}/run_pipeline.py --window-hours 3"
        )
    bundle = joblib.load(model_path)
    metadata = json.loads(meta_path.read_text(encoding="utf-8"))
    _BUNDLE_CACHE[kind] = (bundle, metadata)
    return bundle, metadata


def _window_index(hourly_rows: list[dict], kind: str) -> dict[str, list[tuple[str, dict]]]:
    _, feature_fn = _assets(kind)
    samples = windowing.build_samples(hourly_rows, kind, [], WINDOW_HOURS, feature_fn)
    by_unit: dict[str, list[tuple[str, dict]]] = {}
    for sample in samples:
        by_unit.setdefault(sample["unit"], []).append((sample["end_hour"], sample))
    for series in by_unit.values():
        series.sort()
    return by_unit


def _probability(bundle, feature_module, sample) -> float:
    names = bundle["feature_names"]
    features = feature_module.generalized_features(sample, bundle["unit_baselines"])
    matrix = np.asarray([[features.get(name) for name in names]], dtype=float)
    matrix = bundle["imputer"].transform(matrix)
    return float(bundle["model"].predict_proba(matrix)[0, 1])


def annotate(kind: str, hourly_rows: list[dict], segments: list[dict]) -> dict:
    """Add model probability aggregates + note to each segment in place."""
    bundle, _metadata = load_bundle(kind)
    feature_module, _feature_fn = _assets(kind)
    by_unit = _window_index(hourly_rows, kind)
    scored = 0
    for segment in segments:
        unit = segment.get("机组")
        start, end = segment.get("开始小时"), segment.get("结束小时")
        probabilities = []
        if start and end:
            for end_hour, sample in by_unit.get(unit, []):
                if start <= end_hour <= end:
                    probabilities.append(_probability(bundle, feature_module, sample))
        if probabilities:
            mean = float(np.mean(probabilities))
            high = float(np.max(probabilities))
            low = float(np.min(probabilities))
            segment["模型异常概率均值"] = round(mean, 6)
            segment["模型异常概率最大值"] = round(high, 6)
            segment["模型异常概率最小值"] = round(low, 6)
            segment[NOTE_FIELD] = (
                f"模型预测参考：{WINDOW_HOURS}小时窗口异常概率 "
                f"均值{mean:.4f}/最大{high:.4f}/最小{low:.4f}（{len(probabilities)}次）"
            )
            scored += 1
        else:
            segment[NOTE_FIELD] = "模型预测参考：无可用3小时窗口，数据不足"
    return {"kind": kind, "segments": len(segments), "scored": scored}


__all__ = ["WINDOW_HOURS", "MODEL_FIELDS", "NOTE_FIELD", "load_bundle", "annotate"]