backend/scripts/fit_ultimate_runtime_refiner.py

from __future__ import annotations

import json
import sys
import warnings
from pathlib import Path

import numpy as np
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, roc_auc_score

BACKEND_ROOT = Path(__file__).resolve().parents[1]
if str(BACKEND_ROOT) not in sys.path:
    sys.path.insert(0, str(BACKEND_ROOT))
SCRIPT_ROOT = Path(__file__).resolve().parent
if str(SCRIPT_ROOT) not in sys.path:
    sys.path.insert(0, str(SCRIPT_ROOT))

from app.config import (  # noqa: E402
    DEFAULT_ARCHIVE_MODEL_PATH,
    DEFAULT_ULTIMATE_REFINEMENT_REPORT_PATH,
    DEFAULT_ULTIMATE_REFINER_PATH,
)
from app.ml.archive_model import load_archive_model, predict_with_archive_model  # noqa: E402
from app.ml.features import extract_eye_features, extract_ultimate_clinical_features  # noqa: E402
from app.ml.ultimate_runtime_refinement import UltimateRuntimeRefiner  # noqa: E402
from app.services.image_quality import ImageQualityService  # noqa: E402
from train_efficientnet import ARCHIVE_ROOT, _balanced_group_split, _build_records  # noqa: E402

warnings.filterwarnings(
    "ignore",
    message="X does not have valid feature names, but StandardScaler was fitted with feature names",
)


def _metric_block(labels: np.ndarray, predictions: np.ndarray) -> dict[str, float]:
    return {
        "accuracy": round(float(accuracy_score(labels, predictions)), 4),
        "precision": round(float(precision_score(labels, predictions, zero_division=0)), 4),
        "recall": round(float(recall_score(labels, predictions, zero_division=0)), 4),
        "f1": round(float(f1_score(labels, predictions, zero_division=0)), 4),
    }


def _choose_threshold(labels: np.ndarray, probabilities: np.ndarray) -> tuple[float, dict[str, float]]:
    best_threshold = 0.35
    best_metrics: dict[str, float] | None = None
    best_score = -1.0
    for threshold in np.linspace(0.2, 0.8, 61):
        predictions = (probabilities >= threshold).astype(np.int32)
        metrics = _metric_block(labels, predictions)
        score = (
            metrics["f1"] * 0.65
            + metrics["precision"] * 0.20
            + metrics["recall"] * 0.15
        )
        if score > best_score:
            best_score = score
            best_threshold = float(threshold)
            best_metrics = metrics
    assert best_metrics is not None
    return best_threshold, best_metrics


def _collect_runtime_rows(records):
    quality_service = ImageQualityService()
    rows = []
    for record in records:
        with record.image_path.open("rb") as handle:
            quality, processed = quality_service.evaluate(handle.read())
        if not quality.passed:
            continue
        rows.append(
            {
                "record": record,
                "quality": quality,
                "base_features": extract_eye_features(processed),
                "ultimate_features": extract_ultimate_clinical_features(processed, quality),
            }
        )
    return rows


def _expected_scaler_stats(archive_model: dict[str, object]) -> tuple[list[str], dict[str, float], dict[str, float]]:
    feature_names = archive_model.get("feature_names")
    scaler = archive_model.get("scaler")
    if (
        not isinstance(feature_names, list)
        or not feature_names
        or scaler is None
        or not hasattr(scaler, "mean_")
        or not hasattr(scaler, "scale_")
    ):
        raise RuntimeError("Ultimate archive artifact is missing scaler statistics.")
    return (
        [str(name) for name in feature_names],
        {name: float(value) for name, value in zip(feature_names, scaler.mean_, strict=False)},
        {
            name: max(float(value), 1e-6)
            for name, value in zip(feature_names, scaler.scale_, strict=False)
        },
    )


def _feature_stats(rows, feature_names: list[str]) -> tuple[dict[str, float], dict[str, float]]:
    return (
        {
            name: float(np.mean([row["ultimate_features"][name] for row in rows]))
            for name in feature_names
        },
        {
            name: max(
                float(np.std([row["ultimate_features"][name] for row in rows])),
                1e-6,
            )
            for name in feature_names
        },
    )


def _build_dataset(
    rows,
    *,
    archive_model: dict[str, object],
    feature_names: list[str],
    expected_means: dict[str, float],
    expected_stds: dict[str, float],
    current_means: dict[str, float],
    current_stds: dict[str, float],
):
    refiner = UltimateRuntimeRefiner(
        feature_means=current_means,
        feature_stds=current_stds,
    )
    X_rows: list[list[float]] = []
    labels: list[int] = []
    base_predictions: list[int] = []
    raw_predictions: list[int] = []
    for row in rows:
        remapped = refiner.remap_ultimate_features(
            row["ultimate_features"],
            archive_feature_names=feature_names,
            expected_means=expected_means,
            expected_stds=expected_stds,
        )
        base_prediction = predict_with_archive_model(archive_model, remapped, source_hint="roi_original")
        X_rows.append(
            refiner._feature_vector(
                base_prediction=base_prediction,
                quality=row["quality"],
                base_feature_map=row["base_features"],
            )
        )
        label = int(row["record"].label)
        labels.append(label)
        base_predictions.append(int(base_prediction["anemia_risk"] >= 0.5))
        raw_predictions.append(int(base_prediction["anemia_risk"] >= 0.5))
    return (
        np.asarray(X_rows, dtype=np.float32),
        np.asarray(labels, dtype=np.int32),
        np.asarray(base_predictions, dtype=np.int32),
        np.asarray(raw_predictions, dtype=np.int32),
    )


def main() -> None:
    archive_model = load_archive_model(DEFAULT_ARCHIVE_MODEL_PATH)
    version = str(archive_model.get("version", ""))
    if not version.startswith("archive-fusion-v7-ultimate-clinical"):
        raise RuntimeError(
            f"Ultimate runtime refiner expects the v7 clinical artifact, got {version!r}."
        )

    records = _build_records(ARCHIVE_ROOT)
    if not records:
        raise RuntimeError(f"No evaluation records found in {ARCHIVE_ROOT}.")

    train_records, val_records = _balanced_group_split(records, test_size=0.2, n_splits=32)
    train_roi = [record for record in train_records if record.source == "roi_original"]
    val_roi = [record for record in val_records if record.source == "roi_original"]

    train_rows = _collect_runtime_rows(train_roi)
    val_rows = _collect_runtime_rows(val_roi)
    feature_names, expected_means, expected_stds = _expected_scaler_stats(archive_model)
    current_means, current_stds = _feature_stats(train_rows, feature_names)

    X_train, y_train, _, _ = _build_dataset(
        train_rows,
        archive_model=archive_model,
        feature_names=feature_names,
        expected_means=expected_means,
        expected_stds=expected_stds,
        current_means=current_means,
        current_stds=current_stds,
    )
    X_val, y_val, base_predictions, _ = _build_dataset(
        val_rows,
        archive_model=archive_model,
        feature_names=feature_names,
        expected_means=expected_means,
        expected_stds=expected_stds,
        current_means=current_means,
        current_stds=current_stds,
    )

    model = GradientBoostingClassifier(
        random_state=42,
        n_estimators=150,
        learning_rate=0.05,
        max_depth=2,
        min_samples_leaf=3,
        subsample=0.9,
    )
    model.fit(X_train, y_train)
    probabilities = model.predict_proba(X_val)[:, 1]
    selected_threshold, metrics_after = _choose_threshold(y_val, probabilities)
    before_metrics = _metric_block(y_val, base_predictions)
    after_predictions = (probabilities >= selected_threshold).astype(np.int32)
    auc = round(float(roc_auc_score(y_val, probabilities)), 4)

    artifact = UltimateRuntimeRefiner(
        method="gradient-boosting-compatibility",
        threshold=round(selected_threshold, 4),
        feature_means=current_means,
        feature_stds=current_stds,
        model=model,
        report={
            "validation_size": int(len(y_val)),
            "auc": auc,
            "metrics_before": before_metrics,
            "selected_threshold": round(selected_threshold, 4),
        },
    )
    artifact.save(DEFAULT_ULTIMATE_REFINER_PATH)

    report = {
        "version": artifact.version,
        "method": artifact.method,
        "validation_size": int(len(y_val)),
        "selected_threshold": round(selected_threshold, 4),
        "auc": auc,
        "metrics_before": before_metrics,
        "metrics_after": _metric_block(y_val, after_predictions),
    }
    DEFAULT_ULTIMATE_REFINEMENT_REPORT_PATH.write_text(
        json.dumps(report, indent=2),
        encoding="utf-8",
    )

    print("Ultimate runtime refinement")
    print(f"validation_size: {report['validation_size']}")
    print(f"selected_threshold: {report['selected_threshold']:.4f}")
    print(f"auc: {report['auc']:.4f}")
    print("before:", report["metrics_before"])
    print("after:", report["metrics_after"])


if __name__ == "__main__":
    main()