Compressor/models/common/modeling.py

"""Generic XGBoost training / evaluation / inference for both model types.

Ports the original generalized model pipeline's training, evaluation and
inference logic.  The only model-specific input is the feature module
(``features_pressure`` or ``features_vibration``).
"""
from __future__ import annotations

import csv
import json
from collections import defaultdict
from pathlib import Path

import joblib
import numpy as np
from sklearn.impute import SimpleImputer
from sklearn.metrics import average_precision_score, confusion_matrix, precision_recall_curve, precision_recall_fscore_support
from sklearn.neighbors import KNeighborsClassifier, NearestNeighbors
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from xgboost import XGBClassifier
from xgboost.callback import TrainingCallback


def load_samples(directory: Path) -> list[dict]:
    return [json.loads(path.read_text(encoding="utf-8")) for path in sorted(Path(directory).glob("*.json"))]


def matrix(feature_dicts: list[dict], names: list[str]) -> np.ndarray:
    return np.asarray([[item.get(name) for name in names] for item in feature_dicts], dtype=float)


class ProgressCallback(TrainingCallback):
    def __init__(self, rounds: int, learning_rate: float) -> None:
        self.rounds = rounds
        self.learning_rate = learning_rate

    def after_iteration(self, model, epoch: int, evals_log):
        losses = []
        for dataset, metrics in evals_log.items():
            for name, values in metrics.items():
                if values:
                    losses.append(f"{dataset}-{name}={values[-1]:.5f}")
        print(f"[round {epoch + 1:>4}/{self.rounds}] lr={self.learning_rate:.4f} " + " ".join(losses), flush=True)
        return False


def stratified_normal_sample(samples: list[dict], max_ratio: int, seed: int) -> list[dict]:
    positives = [sample for sample in samples if sample["label"] == 1]
    normals = [sample for sample in samples if sample["label"] == 0]
    limit = len(positives) * max_ratio
    if len(normals) <= limit:
        return samples
    buckets = defaultdict(list)
    for sample in normals:
        month = sample["start_hour"][:7]
        buckets[(sample["unit"], sample["batch"], month)].append(sample)
    rng = np.random.default_rng(seed)
    selected = []
    keys = sorted(buckets)
    quota, remainder = divmod(limit, len(keys))
    for index, key in enumerate(keys):
        bucket = buckets[key]
        take = min(len(bucket), quota + (1 if index < remainder else 0))
        selected.extend(bucket[index] for index in rng.choice(len(bucket), size=take, replace=False))
    if len(selected) < limit:
        leftovers = [sample for sample in normals if sample not in selected]
        need = min(limit - len(selected), len(leftovers))
        selected.extend(leftovers[index] for index in rng.choice(len(leftovers), size=need, replace=False))
    return positives + selected


def choose_threshold(y: np.ndarray, probabilities: np.ndarray) -> tuple[float, dict]:
    precision, recall, thresholds = precision_recall_curve(y, probabilities)
    f1 = 2 * precision * recall / np.maximum(precision + recall, 1e-12)
    index = int(np.nanargmax(f1[:-1])) if len(thresholds) else 0
    threshold = float(thresholds[index]) if len(thresholds) else 0.5
    return threshold, {
        "threshold_selection": "max_validation_f1",
        "threshold": threshold,
        "validation_precision": float(precision[index]),
        "validation_recall": float(recall[index]),
        "validation_f1": float(f1[index]),
        "validation_pr_auc": float(average_precision_score(y, probabilities)),
    }


def train(model_type: str, feature_module, dataset_dir: Path, models_dir: Path, seed: int, rounds: int,
          early_stopping_rounds: int, all_normal: bool, normal_ratio: int | None) -> dict:
    ratio = normal_ratio if normal_ratio is not None else (4 if model_type == "pressure" else 5)
    dataset = Path(dataset_dir)
    train_all = load_samples(dataset / "train")
    validation = load_samples(dataset / "validation")
    if not train_all or not validation:
        raise RuntimeError("训练或验证样本为空")
    if not any(sample["label"] == 1 for sample in train_all) or not any(sample["label"] == 1 for sample in validation):
        raise RuntimeError("训练和验证集都必须包含异常样本")
    train = train_all if all_normal else stratified_normal_sample(train_all, ratio, seed)
    baselines = feature_module.compute_unit_baselines(train_all)
    names = feature_module.feature_names()
    x_train, y_train = matrix(feature_module.feature_matrix(train, baselines), names), np.asarray([sample["label"] for sample in train])
    x_validation, y_validation = matrix(feature_module.feature_matrix(validation, baselines), names), np.asarray([sample["label"] for sample in validation])
    imputer = SimpleImputer(strategy="median")
    x_train = imputer.fit_transform(x_train)
    x_validation = imputer.transform(x_validation)
    scale_pos_weight = float(np.sum(y_train == 0) / max(np.sum(y_train == 1), 1))
    learning_rate = 0.05
    print(f"training: rounds={rounds}, learning_rate={learning_rate:.4f}, early_stopping_rounds={early_stopping_rounds}, samples={len(train):,}", flush=True)
    model = XGBClassifier(
        n_estimators=rounds, max_depth=3, learning_rate=learning_rate,
        subsample=0.8, colsample_bytree=0.8, min_child_weight=3,
        objective="binary:logistic", eval_metric="logloss", scale_pos_weight=scale_pos_weight,
        random_state=seed, n_jobs=4, early_stopping_rounds=early_stopping_rounds,
        callbacks=[ProgressCallback(rounds, learning_rate)],
    )
    model.fit(x_train, y_train, eval_set=[(x_train, y_train), (x_validation, y_validation)], verbose=False)
    rounds_run = len(model.evals_result()["validation_0"]["logloss"])
    stopped_early = rounds_run < rounds
    print(f"training complete: rounds_run={rounds_run}/{rounds}, best_iteration={model.best_iteration + 1}, best_validation_logloss={float(model.best_score):.5f}, early_stopped={'yes' if stopped_early else 'no'}", flush=True)
    validation_probabilities = model.predict_proba(x_validation)[:, 1]
    threshold, validation_metrics = choose_threshold(y_validation, validation_probabilities)
    similarity = Pipeline([("scale", StandardScaler()), ("knn", KNeighborsClassifier(n_neighbors=min(5, len(train)), weights="distance"))])
    similarity.fit(x_train, y_train)
    models = Path(models_dir)
    models.mkdir(parents=True, exist_ok=True)
    model.callbacks = []
    joblib.dump(
        {"model": model, "imputer": imputer, "knn": similarity, "feature_names": names, "unit_baselines": baselines},
        models / f"{model_type}_model.joblib",
    )
    (models / f"{model_type}_unit_baselines.json").write_text(json.dumps(baselines, ensure_ascii=False, indent=2), encoding="utf-8")
    metadata = {
        "model_type": model_type, "algorithm": "XGBoost binary classifier + KNN similarity",
        "feature_version": "generalized_cycle_relative_v2",
        "feature_names": names,
        "dropped_absolute_features": list(feature_module.ABSOLUTE_FEATURES),
        "unit_baselines": baselines,
        "normal_sampling_strategy": "all_normal" if all_normal else "stratified_max_ratio",
        "normal_to_anomaly_training_ratio_limit": None if all_normal else ratio,
        "actual_normal_to_anomaly_training_ratio": float(np.sum(y_train == 0) / max(np.sum(y_train == 1), 1)),
        "training_samples_before_sampling": len(train_all), "training_samples_after_sampling": len(train),
        "training_normal": int(np.sum(y_train == 0)), "training_anomaly": int(np.sum(y_train == 1)),
        "scale_pos_weight": scale_pos_weight, "validation_samples": len(validation),
        "test_set_read": False, **validation_metrics,
        "rounds_requested": rounds, "early_stopping_rounds_requested": early_stopping_rounds,
        "rounds_run": rounds_run, "early_stopped": stopped_early,
        "best_iteration": int(model.best_iteration), "best_score": float(model.best_score),
    }
    (models / f"{model_type}_metadata.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(metadata, ensure_ascii=False, indent=2))
    return metadata


def metrics(y: np.ndarray, predicted: np.ndarray, probabilities: np.ndarray) -> dict:
    labels = [0, 1]
    per_precision, per_recall, per_f1, per_support = precision_recall_fscore_support(y, predicted, labels=labels, zero_division=0)
    tn, fp, fn, tp = confusion_matrix(y, predicted, labels=labels).ravel()
    samples = int(len(y))
    normal_support, anomaly_support = int(per_support[0]), int(per_support[1])
    specificity = float(tn / (tn + fp)) if (tn + fp) else 0.0
    accuracy = float((tp + tn) / samples) if samples else 0.0
    return {
        "samples": samples, "normal_samples": normal_support, "anomaly_samples": anomaly_support,
        "tp": int(tp), "fp": int(fp), "fn": int(fn), "tn": int(tn),
        "precision": float(per_precision[1]), "recall": float(per_recall[1]), "f1": float(per_f1[1]),
        "pr_auc": float(average_precision_score(y, probabilities)) if len(set(y)) == 2 else None,
        "specificity": specificity, "accuracy": accuracy,
        "balanced_accuracy": float((per_recall[0] + per_recall[1]) / 2),
        "macro_f1": float((per_f1[0] + per_f1[1]) / 2),
        "weighted_f1": float((per_f1[0] * normal_support + per_f1[1] * anomaly_support) / samples) if samples else 0.0,
        "class_metrics": {
            "normal": {"precision": float(per_precision[0]), "recall": float(per_recall[0]), "f1": float(per_f1[0]), "support": normal_support},
            "anomaly": {"precision": float(per_precision[1]), "recall": float(per_recall[1]), "f1": float(per_f1[1]), "support": anomaly_support},
        },
        "confusion_matrix_counts": {"tn": int(tn), "fp": int(fp), "fn": int(fn), "tp": int(tp)},
        "confusion_matrix_normalized": {
            "normal": {"predicted_normal": float(tn / normal_support) if normal_support else 0.0, "predicted_anomaly": float(fp / normal_support) if normal_support else 0.0},
            "anomaly": {"predicted_normal": float(fn / anomaly_support) if anomaly_support else 0.0, "predicted_anomaly": float(tp / anomaly_support) if anomaly_support else 0.0},
        },
    }


def event_metrics(samples: list[dict], predicted: np.ndarray) -> dict:
    groups = defaultdict(list)
    for sample, value in zip(samples, predicted):
        groups[sample["event_group_id"]].append((sample, int(value)))
    positives = [values for values in groups.values() if values[0][0]["label"] == 1]
    negatives = [values for values in groups.values() if values[0][0]["label"] == 0]
    tp = sum(any(value for _, value in values) for values in positives)
    fn = len(positives) - tp
    fp = sum(any(value for _, value in values) for values in negatives)
    return {
        "anomaly_events": len(positives), "normal_events": len(negatives),
        "detected_anomaly_events": tp, "missed_anomaly_events": fn, "false_positive_normal_events": fp,
        "event_recall": tp / len(positives) if positives else None,
    }


def evaluate(model_type: str, feature_module, dataset_dir: Path, models_dir: Path, output_dir: Path) -> dict:
    bundle = joblib.load(Path(models_dir) / f"{model_type}_model.joblib")
    metadata = json.loads((Path(models_dir) / f"{model_type}_metadata.json").read_text(encoding="utf-8"))
    samples = load_samples(Path(dataset_dir) / "test")
    if not samples or not any(sample["label"] == 1 for sample in samples):
        raise RuntimeError("封存测试集必须包含至少一个异常样本")
    x = bundle["imputer"].transform(matrix(feature_module.feature_matrix(samples, bundle["unit_baselines"]), bundle["feature_names"]))
    y = np.asarray([sample["label"] for sample in samples])
    probabilities = bundle["model"].predict_proba(x)[:, 1]
    predicted = (probabilities >= metadata["threshold"]).astype(int)
    distances, neighbors = bundle["knn"].named_steps["knn"].kneighbors(bundle["knn"].named_steps["scale"].transform(x), n_neighbors=1)
    rows = []
    for index, sample in enumerate(samples):
        rows.append({
            "sample_id": sample["sample_id"], "unit": sample["unit"], "batch": sample["batch"],
            "event_group_id": sample["event_group_id"], "start_hour": sample["start_hour"], "end_hour": sample["end_hour"],
            "true_label": sample["label"], "anomaly_probability": probabilities[index], "predicted_label": int(predicted[index]),
            "nearest_training_sample_index": int(neighbors[index][0]), "nearest_training_distance": float(distances[index][0]),
            "nearest_training_similarity": float(1 / (1 + distances[index][0])),
        })
    by_unit = {}
    for unit in sorted({sample["unit"] for sample in samples}):
        indices = [index for index, sample in enumerate(samples) if sample["unit"] == unit]
        by_unit[unit] = metrics(y[indices], predicted[indices], probabilities[indices])
    report = {
        "model_type": model_type, "evaluation_set": "sealed_test", "threshold": metadata["threshold"],
        "window_metrics": metrics(y, predicted, probabilities), "event_metrics": event_metrics(samples, predicted),
        "by_unit": by_unit,
        "warning": "Labels originate from current scanner rules unless manually updated; metrics measure held-out rule-label agreement, not confirmed physical-fault diagnosis.",
    }
    output = Path(output_dir)
    output.mkdir(parents=True, exist_ok=True)
    with (output / f"{model_type}_test_predictions.csv").open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    (output / f"{model_type}_test_evaluation.json").write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    return report


def combined_result(rule_label: int, probability: float, anomaly_similarity: float, threshold: float) -> str:
    model_support = probability >= threshold or anomaly_similarity >= 0.70
    if rule_label == 1 and model_support:
        return "模型辅助确认异常"
    if rule_label == 1:
        return "保留规则异常，模型意见不一致，人工复核"
    if model_support:
        return "规则正常但模型疑似异常，继续观察"
    return "正常"


def infer(model_type: str, feature_module, dataset_dir: Path, models_dir: Path, output_dir: Path) -> dict:
    dataset = Path(dataset_dir)
    bundle = joblib.load(Path(models_dir) / f"{model_type}_model.joblib")
    metadata = json.loads((Path(models_dir) / f"{model_type}_metadata.json").read_text(encoding="utf-8"))
    samples = load_samples(dataset / "candidates")
    training = load_samples(dataset / "train")
    baselines = bundle["unit_baselines"]
    x = bundle["imputer"].transform(matrix(feature_module.feature_matrix(samples, baselines), bundle["feature_names"]))
    x_train = bundle["imputer"].transform(matrix(feature_module.feature_matrix(training, baselines), bundle["feature_names"]))
    probabilities = bundle["model"].predict_proba(x)[:, 1]
    scaled = bundle["knn"].named_steps["scale"].transform(x)
    train_scaled = bundle["knn"].named_steps["scale"].transform(x_train)
    anomaly_indexes = [index for index, sample in enumerate(training) if sample["label"] == 1]
    if not anomaly_indexes:
        raise RuntimeError("训练集没有异常样本，无法建立异常相似性索引")
    anomaly_neighbors = NearestNeighbors(n_neighbors=1).fit(train_scaled[anomaly_indexes])
    distances, nearest_indexes = anomaly_neighbors.kneighbors(scaled)
    rows = []
    for index, sample in enumerate(samples):
        nearest = training[anomaly_indexes[int(nearest_indexes[index][0])]]
        anomaly_similarity = float(1 / (1 + distances[index][0]))
        rule_label = int(sample["label"])
        rows.append({
            "机组": sample["unit"], "批次": sample["batch"], "窗口开始小时": sample["start_hour"], "窗口结束小时": sample["end_hour"],
            "样本编号": sample["sample_id"], "现有规则结果": "异常" if rule_label else "正常",
            "规则阶段编号": sample.get("rule_stage_id") or "", "模型异常概率": round(float(probabilities[index]), 6),
            "最相似历史异常样本": nearest["sample_id"], "历史异常相似度": round(anomaly_similarity, 6),
            "联合判断": combined_result(rule_label, probabilities[index], anomaly_similarity, metadata["threshold"]),
        })
    rows.sort(key=lambda row: (int(row["机组"][0]), row["窗口结束小时"], row["批次"]))
    output = Path(output_dir)
    output.mkdir(parents=True, exist_ok=True)
    name = "压力模型辅助确认.csv" if model_type == "pressure" else "振动模型辅助确认.csv"
    with (output / name).open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    return {"rows": len(rows), "output": str(output / name)}


__all__ = [
    "load_samples", "matrix", "train", "evaluate", "infer", "metrics", "event_metrics", "choose_threshold",
]