#!/usr/bin/env python3
"""Generate retrospective model-assistance outputs without changing rule CSVs."""
from __future__ import annotations

import argparse
import csv
import importlib
import json
from pathlib import Path

import joblib
import numpy as np
from sklearn.neighbors import NearestNeighbors


def feature_module(model_type: str):
    """Pressure and vibration use independent feature modules."""
    return importlib.import_module(
        "generalized_features" if model_type == "pressure" else "generalized_features_vibration"
    )


def load_samples(directory: Path) -> list[dict]:
    return [json.loads(path.read_text(encoding="utf-8")) for path in sorted(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)


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 main() -> None:
    parser = argparse.ArgumentParser(description="生成PKS规则与模型的历史联合辅助确认结果")
    parser.add_argument("model_type", choices=("pressure", "vibration"))
    parser.add_argument("--datasets-dir", default=str(Path(__file__).resolve().parent / "datasets"))
    parser.add_argument("--models-dir", default=str(Path(__file__).resolve().parent / "models"))
    parser.add_argument("--output-dir", default=str(Path(__file__).resolve().parents[2] / "cache" / "pks_scanner" / "model_assistant_generalized"))
    args = parser.parse_args()
    gf = feature_module(args.model_type)
    dataset = Path(args.datasets_dir) / args.model_type
    bundle = joblib.load(Path(args.models_dir) / f"{args.model_type}_model.joblib")
    metadata = json.loads((Path(args.models_dir) / f"{args.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(gf.feature_matrix(samples, baselines), bundle["feature_names"]))
    x_train = bundle["imputer"].transform(matrix(gf.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:
        parser.error("训练集没有异常样本，无法建立异常相似性索引")
    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(args.output_dir)
    output.mkdir(parents=True, exist_ok=True)
    name = "压力模型辅助确认.csv" if args.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)
    print(f"{args.model_type} 联合辅助结果：{len(rows):,} 行，输出：{output / name}")


if __name__ == "__main__":
    main()