#!/usr/bin/env python3
"""Evaluate a fixed PKS assistant model against the sealed test directory."""
from __future__ import annotations

import argparse
import csv
import importlib
import json
from collections import Counter, defaultdict
from pathlib import Path

import joblib
import numpy as np
from sklearn.metrics import average_precision_score, confusion_matrix, precision_recall_fscore_support


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 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 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)
    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(Path(args.datasets_dir) / args.model_type / "test")
    if not samples or not any(sample["label"] == 1 for sample in samples):
        parser.error("封存测试集必须包含至少一个异常样本")
    x = bundle["imputer"].transform(matrix(gf.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": args.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(args.output_dir)
    output.mkdir(parents=True, exist_ok=True)
    with (output / f"{args.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"{args.model_type}_test_evaluation.json").write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    summary = report["window_metrics"]
    print(f"测试集：{summary['samples']} 窗口（正常 {summary['normal_samples']}、异常 {summary['anomaly_samples']}），阈值 {report['threshold']:.6f}")
    print(f"{'类别':<6}{'查准率':>10}{'召回率':>10}{'F1':>10}{'样本数':>10}")
    for name, label in (("正常", "normal"), ("异常", "anomaly")):
        item = summary["class_metrics"][label]
        print(f"{name:<6}{item['precision']:>10.4f}{item['recall']:>10.4f}{item['f1']:>10.4f}{item['support']:>10}")
    print(
        f"准确率 {summary['accuracy']:.4f}  平衡准确率 {summary['balanced_accuracy']:.4f}  "
        f"macro-F1 {summary['macro_f1']:.4f}  加权F1 {summary['weighted_f1']:.4f}  "
        f"特异度 {summary['specificity']:.4f}  PR-AUC {summary['pr_auc']:.4f}"
    )
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()