"""Per-class ECE / Brier / F1 + eval_metrics.json builder (D-CAL-04..07, D-MASK-04).

Pitfall 8: every numpy scalar is cast to Python float() / int() at the JSON boundary.
Pitfall 12: all numbers come from data/eval.parquet, NOT cv folds (the
    CalibratedClassifierCV's internal CV is for calibration only — we
    evaluate the resulting calibrated estimator on the disjoint eval split).
"""
from __future__ import annotations

from typing import Any

import numpy as np
from sklearn.calibration import calibration_curve
from sklearn.metrics import (
    brier_score_loss,
    classification_report,
    f1_score,
)

from model.features import CLASSES

# W-5 plan-checker resolution: removed dead `from model.inference import MASK_TABLE` import.
# The OQ-4 distinction (per_mode_macro_f1 subsets eval rows by ACTUAL `network_mode` column,
# NOT applying mask uniformly to all rows) is documented inline in per_mode_macro_f1's comment.

# D-CAL-05: Brier baseline = uniform-prior (always predict 1/10 across all 10 classes)
UNIFORM_PRIOR_PROB: float = 1.0 / len(CLASSES)


def per_class_ece(
    y_true: np.ndarray, y_proba: np.ndarray, *, n_bins: int = 10
) -> dict[str, float]:
    """Per-class ECE with `n_bins` equal-width bins (D-CAL-07: default 10).

    For each class c, treat (y_true == c) as binary target; bin the per-class
    probability column y_proba[:, c]; ECE = mean over occupied bins of
    |fraction_pos - mean_pred|. Bins with zero samples are excluded by
    sklearn's calibration_curve, so the average is over occupied bins.
    """
    out: dict[str, float] = {}
    for c, slug in enumerate(CLASSES):
        y_true_binary = (y_true == c).astype(np.int64)
        prob_true, prob_pred = calibration_curve(
            y_true_binary, y_proba[:, c], n_bins=n_bins, strategy="uniform"
        )
        if len(prob_true) == 0:
            out[slug] = 0.0
        else:
            out[slug] = float(np.mean(np.abs(prob_true - prob_pred)))
    return out


def per_class_brier(
    y_true: np.ndarray, y_proba: np.ndarray
) -> dict[str, float]:
    """Per-class Brier score (D-CAL-04). One-vs-rest binary Brier per class."""
    out: dict[str, float] = {}
    for c, slug in enumerate(CLASSES):
        y_true_binary = (y_true == c).astype(np.int64)
        out[slug] = float(brier_score_loss(y_true_binary, y_proba[:, c]))
    return out


def per_class_brier_baseline(y_true: np.ndarray) -> dict[str, float]:
    """Brier baseline: predict UNIFORM_PRIOR_PROB for every sample, every class.

    D-CAL-05: report alongside trained Brier per class. Reviewer reads
    "trained Brier 0.04 vs baseline Brier 0.18" and instantly sees the
    calibrated model is meaningfully better than chance.
    """
    n = len(y_true)
    n_classes = len(CLASSES)
    proba_uniform = np.full((n, n_classes), UNIFORM_PRIOR_PROB)
    return per_class_brier(y_true, proba_uniform)


def per_class_classification_report(
    y_true: np.ndarray, y_pred: np.ndarray
) -> dict[str, dict[str, float]]:
    """sklearn classification_report restricted to per-class P/R/F1/support.

    Returns dict[slug, {precision, recall, f1, support}] — Python types only
    (Pitfall 8). Cleanly populates eval_metrics.json `per_class[slug]`.
    """
    report = classification_report(
        y_true, y_pred,
        target_names=CLASSES, labels=list(range(len(CLASSES))),
        output_dict=True, zero_division=0,
    )
    out: dict[str, dict[str, float]] = {}
    for slug in CLASSES:
        row = report[slug]
        out[slug] = {
            "precision": float(row["precision"]),
            "recall":    float(row["recall"]),
            "f1":        float(row["f1-score"]),
            "support":   int(row["support"]),
        }
    return out


def per_mode_macro_f1(
    y_true: np.ndarray,
    y_pred_calibrated: np.ndarray,
    network_mode_per_row: np.ndarray,
) -> dict[str, float]:
    """D-MASK-04 / OQ-4 resolution: macro F1 on the SUBSET of eval rows whose
    `network_mode` column actually equals each mode.

    This is the "real" per-mode F1 — what reviewers will assume "by network_mode"
    means. Distinct from "apply mask X uniformly to all rows" diagnostic
    (rejected per OQ-4 / D-MASK-04). Note: NOT apply_mask_and_renormalize —
    subsets eval rows by actual network_mode column (W-5 plan-checker note).

    y_pred_calibrated: argmax over post-mask renormalized probs (so the
    calibrator + mask combination is what the production stack does).
    """
    out: dict[str, float] = {}
    for mode in ("enterprise", "captive", "home", "unknown"):
        mask = network_mode_per_row == mode
        if not mask.any():
            out[mode] = 0.0
            continue
        out[mode] = float(
            f1_score(y_true[mask], y_pred_calibrated[mask],
                     average="macro", zero_division=0,
                     labels=list(range(len(CLASSES)))),
        )
    return out


def build_eval_metrics(
    *,
    y_eval: np.ndarray,
    calibrated_proba: np.ndarray,
    y_pred_after_mask: np.ndarray,
    network_mode_per_row: np.ndarray,
    anomaly_threshold: float,
    per_class_lead_times: dict[str, np.ndarray],
    per_class_miss_rates: dict[str, float],
    schema_version: str = "1.0.0",
) -> dict[str, Any]:
    """Assemble the eval_metrics.json payload (Pattern 11).

    Every numeric is cast at the boundary (Pitfall 8). All inputs come
    from data/eval.parquet — never cv folds (Pitfall 12).
    """
    ece = per_class_ece(y_eval, calibrated_proba, n_bins=10)
    brier = per_class_brier(y_eval, calibrated_proba)
    brier_baseline = per_class_brier_baseline(y_eval)
    clf_report = per_class_classification_report(y_eval, y_pred_after_mask)

    per_class: dict[str, dict[str, float]] = {}
    for slug in CLASSES:
        per_class[slug] = {
            **clf_report[slug],
            "ece":   ece[slug],
            "brier": brier[slug],
            "brier_baseline_uniform": brier_baseline[slug],
        }

    macro_f1 = float(
        f1_score(
            y_eval, y_pred_after_mask, average="macro",
            zero_division=0, labels=list(range(len(CLASSES))),
        )
    )
    ece_mean = float(np.mean(list(ece.values())))

    # Lead-times: aggregate over detected windows only (OQ-3 resolution)
    all_lts = (
        np.concatenate([arr for arr in per_class_lead_times.values() if len(arr) > 0])
        if any(len(a) > 0 for a in per_class_lead_times.values())
        else np.array([0.0])
    )

    per_class_lt_median: dict[str, float] = {}
    for slug in CLASSES:
        arr = per_class_lead_times[slug]
        per_class_lt_median[slug] = float(np.median(arr)) if len(arr) > 0 else 0.0

    return {
        "schema_version": schema_version,
        "macro_f1": macro_f1,
        "ece_mean": ece_mean,
        "per_class": per_class,
        "anomaly": {
            "threshold_95p_normal": float(anomaly_threshold),
            "lead_time_aggregate_median_s": float(np.median(all_lts)),
            "per_class_lead_time_median_s": per_class_lt_median,
            "per_class_miss_rate": {
                slug: float(per_class_miss_rates.get(slug, 0.0)) for slug in CLASSES
            },
        },
        "by_network_mode_macro_f1": per_mode_macro_f1(
            y_eval, y_pred_after_mask, network_mode_per_row
        ),
    }