Blame for predictions/predict_single.py

"""
Single-person inference for the frozen CKB 17-variable model package.

The package has two separate layers:

1. The cluster label is produced by the frozen CKB nearest-centroid model.
   Its SHAP files explain a multiclass random-forest surrogate trained to
   reproduce deployed C1/C2/C3 labels; they do not change the cluster label.
2. Each incident outcome is predicted by the selected first-stage classifier.
   The outcome SHAP files explain the corresponding second-stage regression
   surrogate for that classifier's probability score.  The classifier's
   probability and the surrogate score are therefore reported separately.

The command-line entry point is intentionally small, but the same file can
be imported as a library:

    from predict_single import predict_one

    result = predict_one({
        "sex": 1, "age": 55, "edu_level": 2,
        "marital_status": 1, "work": 1, "retire": 0,
        "hh_size": 3, "smoking": 0, "alcohol": 1,
        "height_cm": 165, "weight_kg": 65, "waist_cm": 82,
        "sbp_mmhg": 125, "dbp_mmhg": 78, "bp_drugs": 0,
        "self_health": 2, "chronic_pain": 0,
    })

The outcome high-risk flag is defined as probability >= the frozen full-CKB
Youden threshold recorded in manifest.json.  No missing-value imputation,
re-standardisation, feature selection, or cohort-specific refitting is done.
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import warnings
from pathlib import Path
from typing import Any, Mapping

import joblib
import numpy as np
import pandas as pd
import shap


PACKAGE_ROOT = Path(__file__).resolve().parent
MANIFEST_PATH = PACKAGE_ROOT / "manifest.json"


def _json_default(value: Any) -> Any:
    """Convert common NumPy/path values into JSON-safe values."""
    if isinstance(value, (np.integer,)):
        return int(value)
    if isinstance(value, (np.floating,)):
        return float(value)
    if isinstance(value, (np.bool_,)):
        return bool(value)
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, np.ndarray):
        return value.tolist()
    raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")


def _read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def _safe_float(value: Any) -> float | None:
    try:
        number = float(value)
    except (TypeError, ValueError):
        return None
    return number if np.isfinite(number) else None


class CKBSingleSamplePredictor:
    """Load the packaged models and make one complete-case prediction."""

    def __init__(self, package_root: str | Path | None = None) -> None:
        self.package_root = Path(package_root or PACKAGE_ROOT).resolve()
        self.manifest = _read_json(self.package_root / "manifest.json")
        self.features = list(self.manifest["features"])
        self.feature_labels = dict(self.manifest.get("feature_labels", {}))
        self.outcome_meta = dict(self.manifest["outcomes"])

        cluster_dir = self.package_root / "cluster"
        with np.load(cluster_dir / "frozen_model.npz", allow_pickle=True) as data:
            self.cluster_centers = np.asarray(data["centers"], dtype=np.float64)
            self.cluster_mu = np.asarray(data["mu"], dtype=np.float64)
            self.cluster_sigma = np.asarray(data["sigma"], dtype=np.float64)
            self.cluster_log_applied = np.asarray(data["log_applied"], dtype=bool)
            frozen_features = [str(x) for x in data["features"].tolist()]
            self.cluster_k = int(np.asarray(data["K"]).ravel()[0])

        if frozen_features != self.features:
            raise ValueError(
                "The frozen cluster feature order does not match manifest.json: "
                f"{frozen_features} != {self.features}"
            )
        if self.cluster_centers.shape != (self.cluster_k, len(self.features)):
            raise ValueError("Unexpected frozen cluster-centre shape")

        self.cluster_thresholds = _read_json(cluster_dir / "ood_thresholds.json")
        calibration_path = cluster_dir / "calibration_model.npz"
        self.cluster_calibration = None
        if calibration_path.exists():
            with np.load(calibration_path, allow_pickle=True) as data:
                self.cluster_calibration = {
                    "coef": np.asarray(data["coef"], dtype=np.float64),
                    "intercept": np.asarray(data["intercept"], dtype=np.float64),
                    "d_std": float(np.asarray(data["d_std"]).ravel()[0]),
                }

        self._cluster_shap_model = None
        self._cluster_shap_explainer = None
        self._outcome_cache: dict[str, dict[str, Any]] = {}

    def _validate_input(self, values: Mapping[str, Any]) -> tuple[pd.DataFrame, list[str]]:
        if not isinstance(values, Mapping):
            raise TypeError("Input must be a mapping with the 17 required feature names")

        missing = [feature for feature in self.features if feature not in values]
        if missing:
            raise ValueError(f"Missing required input variables: {missing}")

        ignored = sorted(set(values) - set(self.features))
        frame = pd.DataFrame([{feature: values[feature] for feature in self.features}])
        for feature in self.features:
            frame[feature] = pd.to_numeric(frame[feature], errors="coerce")
        if frame.isna().any(axis=None):
            bad = frame.columns[frame.isna().any()].tolist()
            raise ValueError(f"All 17 variables must be finite numeric values; invalid: {bad}")
        array = frame.to_numpy(dtype=np.float64)
        if not np.isfinite(array).all():
            raise ValueError("All 17 variables must be finite numeric values")

        log_features = [
            feature for feature, use_log in zip(self.features, self.cluster_log_applied)
            if use_log
        ]
        negative_log = [feature for feature in log_features if float(frame.iloc[0][feature]) < 0]
        if negative_log:
            raise ValueError(
                "The following variables are log1p-transformed by the frozen cluster "
                f"model and cannot be negative: {negative_log}"
            )
        return frame, ignored

    def _cluster_preprocess(self, frame: pd.DataFrame) -> np.ndarray:
        raw = frame[self.features].to_numpy(dtype=np.float64, copy=True)
        for index, use_log in enumerate(self.cluster_log_applied):
            if use_log:
                # This is the locked external-assignment rule from the CKB
                # validation code. Input validation above prevents negatives.
                raw[:, index] = np.log1p(np.maximum(raw[:, index], 0.0))
        denominator = np.where(self.cluster_sigma > 1e-8, self.cluster_sigma, 1.0)
        return (raw - self.cluster_mu) / denominator

    def _cluster_confidence(self, distances: np.ndarray) -> np.ndarray:
        calibration = self.cluster_calibration
        if calibration is None:
            return np.full(distances.shape[0], np.nan, dtype=np.float64)
        centered = distances - distances.min(axis=1, keepdims=True)
        centered = np.clip(centered / (calibration["d_std"] + 1e-8), -10, 10)
        logits = centered @ calibration["coef"].T + calibration["intercept"]
        logits -= logits.max(axis=1, keepdims=True)
        probabilities = np.exp(logits)
        probabilities /= probabilities.sum(axis=1, keepdims=True)
        return probabilities.max(axis=1)

    def _predict_cluster(self, frame: pd.DataFrame) -> dict[str, Any]:
        z = self._cluster_preprocess(frame)
        distances = np.linalg.norm(
            z[:, None, :] - self.cluster_centers[None, :, :], axis=2
        )
        nearest = distances.argmin(axis=1).astype(int)
        ordered = np.sort(distances, axis=1)
        d1 = ordered[:, 0]
        d2 = ordered[:, 1]
        margin = np.where(d2 > 0, (d2 - d1) / d2, 1.0)
        confidence = self._cluster_confidence(distances)

        index = int(nearest[0])
        nearest_distance = float(d1[0])
        second_distance = float(d2[0])
        nearest_margin = float(margin[0])
        nearest_confidence = _safe_float(confidence[0])

        ood = int(nearest_distance > float(self.cluster_thresholds["p99_9_overall"]))
        per_cluster = self.cluster_thresholds.get("p99_5_per_cluster", {})
        uncertain = int(
            nearest_distance > float(per_cluster.get(str(index), np.inf))
            or nearest_margin < float(self.cluster_thresholds["margin_p5"])
        )
        confidence_p10 = self.cluster_thresholds.get("confidence_p10")
        low_conf = int(
            (confidence_p10 is not None and nearest_confidence is not None
             and nearest_confidence < float(confidence_p10))
            or nearest_margin < float(self.cluster_thresholds["margin_p5"])
        )

        if ood:
            status = "unclassified_ood"
        elif uncertain:
            status = "uncertain"
        else:
            status = "assigned"
        label = f"C{index + 1}"

        result = {
            # provisional_cluster_label is always available; cluster_label is
            # only definitive when the frozen protocol calls it assigned.
            "provisional_cluster_label": label,
            "cluster_label": label if status == "assigned" else None,
            "cluster_index": index,
            "classification_status": status,
            "distance_nearest": nearest_distance,
            "distance_second": second_distance,
            "margin": nearest_margin,
            "calibrated_confidence": nearest_confidence,
            "ood_flag": bool(ood),
            "uncertain_flag": bool(uncertain),
            "low_confidence_flag": bool(low_conf),
        }

        self._ensure_cluster_shap()
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            surrogate_probability = self._cluster_shap_model.predict_proba(
                frame.to_numpy(dtype=np.float64)
            )[0]
        result["shap_surrogate_cluster_probability"] = {
            f"C{int(class_index) + 1}": float(probability)
            for class_index, probability in zip(
                self._cluster_shap_model.classes_, surrogate_probability
            )
        }
        return result

    def _ensure_cluster_shap(self) -> None:
        if self._cluster_shap_model is not None:
            return
        cluster_dir = self.package_root / "cluster"
        self._cluster_shap_model = joblib.load(
            cluster_dir / "cluster_shap_surrogate_random_forest.joblib"
        )
        self._cluster_shap_explainer = shap.TreeExplainer(self._cluster_shap_model)

    @staticmethod
    def _normalise_cluster_shap_values(
        values: Any, n_classes: int, n_features: int
    ) -> np.ndarray:
        """Return an array with shape (class, sample, feature)."""
        if isinstance(values, list):
            arrays = [np.asarray(item, dtype=np.float64) for item in values]
            return np.stack(arrays, axis=0)
        array = np.asarray(values, dtype=np.float64)
        if array.ndim == 2:
            return array[None, :, :]
        if array.ndim != 3:
            raise ValueError(f"Unexpected multiclass SHAP shape: {array.shape}")
        if array.shape[0] == 1 and array.shape[1] == n_features:
            return np.transpose(array, (2, 0, 1))
        if array.shape[0] == n_classes and array.shape[2] == n_features:
            return array
        if array.shape[0] == n_classes and array.shape[1] == n_features:
            return array
        raise ValueError(f"Could not normalise multiclass SHAP shape: {array.shape}")

    def _cluster_force_plots(
        self, frame: pd.DataFrame, output_dir: Path
    ) -> dict[str, str]:
        self._ensure_cluster_shap()
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            raw_values = self._cluster_shap_explainer.shap_values(
                frame.to_numpy(dtype=np.float64)
            )
        n_classes = len(self._cluster_shap_model.classes_)
        values = self._normalise_cluster_shap_values(
            raw_values, n_classes=n_classes, n_features=len(self.features)
        )
        expected = np.asarray(self._cluster_shap_explainer.expected_value, dtype=float).reshape(-1)
        plot_paths: dict[str, str] = {}
        feature_values = frame.iloc[0].to_numpy(dtype=np.float64)
        for class_position, class_index in enumerate(self._cluster_shap_model.classes_):
            class_label = f"C{int(class_index) + 1}"
            path = output_dir / f"single_sample_cluster_{class_label}_force.html"
            force = shap.force_plot(
                float(expected[class_position]),
                values[class_position, 0, :],
                feature_values,
                feature_names=[self.feature_labels.get(f, f) for f in self.features],
                matplotlib=False,
            )
            shap.save_html(str(path), force)
            plot_paths[class_label] = str(path)
        return plot_paths

    def _load_outcome_bundle(self, outcome: str) -> dict[str, Any]:
        if outcome in self._outcome_cache:
            return self._outcome_cache[outcome]
        meta = self.outcome_meta[outcome]
        outcome_dir = self.package_root / "outcomes" / outcome
        bundle = {
            "classifier": joblib.load(outcome_dir / meta["classifier_file"]),
            "shap_surrogate": joblib.load(outcome_dir / meta["shap_surrogate_file"]),
        }
        self._outcome_cache[outcome] = bundle
        return bundle

    def _outcome_force_plot(
        self,
        outcome: str,
        frame: pd.DataFrame,
        bundle: dict[str, Any],
        output_dir: Path,
    ) -> tuple[str, float]:
        # The bundled explanation pickles were created on a different Python
        # runtime. In particular, PermutationExplainer serialises a Numba
        # dispatcher whose bytecode is not portable between Python versions.
        # Rebuild the explainer in the active runtime from the frozen surrogate
        # and its stored background sample instead of invoking that dispatcher.
        # The prediction-only API still never pays this cost.
        if "shap_explainer" not in bundle:
            meta = self.outcome_meta[outcome]
            outcome_dir = self.package_root / "outcomes" / outcome
            surrogate = bundle["shap_surrogate"]
            if hasattr(surrogate, "estimators_"):
                # Tree SHAP is deterministic for the frozen RandomForest
                # surrogates and avoids another cross-version pickle edge case.
                bundle["shap_explainer"] = shap.TreeExplainer(surrogate)
            else:
                stored_explainer = joblib.load(
                    outcome_dir / meta["shap_explainer_file"]
                )
                background = getattr(
                    getattr(stored_explainer, "masker", None), "data", None
                )
                if background is None:
                    raise RuntimeError(
                        f"Missing SHAP background sample for outcome {outcome}."
                    )
                bundle["shap_explainer"] = shap.PermutationExplainer(
                    surrogate.predict,
                    np.asarray(background, dtype=np.float64),
                )
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            explanation = bundle["shap_explainer"](
                frame.to_numpy(dtype=np.float64)
            )
        values = np.asarray(explanation.values, dtype=np.float64)
        if values.ndim == 3:
            # The outcome SHAP artifacts are univariate. This branch makes a
            # malformed or future artifact fail loudly rather than silently
            # choosing an arbitrary class.
            if values.shape[-1] != 1:
                raise ValueError(f"Unexpected SHAP output shape for {outcome}: {values.shape}")
            values = values[:, :, 0]
        base_values = np.asarray(explanation.base_values, dtype=np.float64).reshape(-1)
        if values.shape != (1, len(self.features)) or base_values.size < 1:
            raise ValueError(
                f"Unexpected SHAP explanation shape for {outcome}: "
                f"values={values.shape}, base={base_values.shape}"
            )
        surrogate_score = float(bundle["shap_surrogate"].predict(frame)[0])
        path = output_dir / f"single_sample_{outcome}_force.html"
        force = shap.force_plot(
            float(base_values[0]),
            values[0],
            frame.iloc[0].to_numpy(dtype=np.float64),
            feature_names=[self.feature_labels.get(f, f) for f in self.features],
            matplotlib=False,
        )
        shap.save_html(str(path), force)
        return str(path), surrogate_score

    def predict_one(
        self,
        values: Mapping[str, Any],
        output_dir: str | Path | None = None,
        make_force_plots: bool = True,
        sample_id: str | None = None,
    ) -> dict[str, Any]:
        """Predict the cluster and all 13 outcomes for one complete case.

        Parameters
        ----------
        values:
            Mapping containing the 17 required variables. Extra keys are
            ignored and reported in ``input.ignored_fields``.
        output_dir:
            Directory for JSON and HTML force plots. If omitted, a timestamped
            folder is created under ``1-output/model/predictions``.
        make_force_plots:
            Save one HTML force plot for each cluster class and outcome when
            true. Setting false only skips plots; prediction probabilities and
            labels are unchanged.
        sample_id:
            Optional identifier written to the result JSON. It is never used
            as a model feature.
        """
        frame, ignored = self._validate_input(values)
        if output_dir is None:
            stamp = dt.datetime.now().strftime("%Y%m%d_%H%M%S")
            output_path = self.package_root / "predictions" / stamp
        else:
            output_path = Path(output_dir).expanduser().resolve()
        output_path.mkdir(parents=True, exist_ok=True)
        force_dir = output_path / "force_plots"
        if make_force_plots:
            force_dir.mkdir(parents=True, exist_ok=True)

        result: dict[str, Any] = {
            "package_version": self.manifest.get("package_version"),
            "generated_at": dt.datetime.now(dt.timezone.utc).isoformat(),
            "sample_id": sample_id,
            "input": {
                "features": {
                    feature: float(frame.iloc[0][feature]) for feature in self.features
                },
                "ignored_fields": ignored,
            },
            "cluster": self._predict_cluster(frame),
            "outcomes": {},
            "force_plots": {},
            "notes": [
                "The definitive cluster is the frozen nearest-centroid assignment.",
                "Cluster SHAP explains a multiclass RF surrogate of deployed cluster labels.",
                "Outcome SHAP explains a regression surrogate of the selected classifier probability.",
                "High risk is probability >= the frozen full-CKB Youden threshold.",
            ],
        }

        if make_force_plots:
            result["force_plots"]["cluster"] = self._cluster_force_plots(frame, force_dir)

        X = frame.to_numpy(dtype=np.float64)
        for outcome, meta in self.outcome_meta.items():
            bundle = self._load_outcome_bundle(outcome)
            classifier = bundle["classifier"]
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                probabilities = np.asarray(classifier.predict_proba(X), dtype=np.float64)
            classes = np.asarray(getattr(classifier, "classes_", [0, 1]))
            positive = np.where(classes == 1)[0]
            if positive.size != 1:
                raise ValueError(f"Outcome {outcome} does not have a unique positive class")
            probability = float(probabilities[0, int(positive[0])])
            threshold = float(meta["youden_threshold"])
            risk = "high_risk" if probability >= threshold else "low_risk"

            row: dict[str, Any] = {
                "outcome": meta.get("display_name", outcome),
                "model": meta["model"],
                "probability": probability,
                "youden_threshold": threshold,
                "risk_class": risk,
                "threshold_rule": "high_risk if probability >= Youden threshold",
                "shap_explanation_target": "selected_classifier_probability_surrogate",
            }
            if make_force_plots:
                path, surrogate_score = self._outcome_force_plot(
                    outcome, frame, bundle, force_dir
                )
                row["shap_surrogate_score"] = surrogate_score
                result["force_plots"][outcome] = path
            else:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    row["shap_surrogate_score"] = float(
                        bundle["shap_surrogate"].predict(frame)[0]
                    )
            result["outcomes"][outcome] = row

        result_path = output_path / "single_sample_prediction.json"
        result["result_json"] = str(result_path)
        result_path.write_text(
            json.dumps(result, ensure_ascii=False, indent=2, default=_json_default),
            encoding="utf-8",
        )
        return result


def predict_one(
    values: Mapping[str, Any],
    package_dir: str | Path | None = None,
    output_dir: str | Path | None = None,
    make_force_plots: bool = True,
    sample_id: str | None = None,
) -> dict[str, Any]:
    """Convenience wrapper around :class:`CKBSingleSamplePredictor`."""
    predictor = CKBSingleSamplePredictor(package_dir)
    return predictor.predict_one(
        values,
        output_dir=output_dir,
        make_force_plots=make_force_plots,
        sample_id=sample_id,
    )


def _load_input_json(path: Path) -> tuple[dict[str, Any], str | None]:
    data = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise ValueError("Input JSON must contain an object")
    sample_id = data.get("sample_id")
    if isinstance(data.get("features"), dict):
        data = data["features"]
    return data, str(sample_id) if sample_id is not None else None


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run CKB frozen cluster and 13-outcome single-sample inference."
    )
    parser.add_argument(
        "--input-json", required=True, type=Path,
        help="JSON object containing the 17 variables; example_input.json is provided.",
    )
    parser.add_argument(
        "--output-dir", type=Path, default=None,
        help="Output folder; default is a timestamped folder under predictions/.",
    )
    parser.add_argument(
        "--sample-id", default=None,
        help="Optional identifier; it is written to output but not used for prediction.",
    )
    parser.add_argument(
        "--no-force-plots", action="store_true",
        help="Skip HTML force plots while retaining predictions and SHAP surrogate scores.",
    )
    args = parser.parse_args()
    values, input_sample_id = _load_input_json(args.input_json)
    sample_id = args.sample_id if args.sample_id is not None else input_sample_id
    result = predict_one(
        values,
        output_dir=args.output_dir,
        make_force_plots=not args.no_force_plots,
        sample_id=sample_id,
    )
    print(json.dumps({
        "result_json": result["result_json"],
        "cluster": result["cluster"],
        "outcomes": result["outcomes"],
        "force_plots": result["force_plots"],
    }, ensure_ascii=False, indent=2, default=_json_default))


if __name__ == "__main__":
    main()