backend/fusion/meta_learner.py

"""
Meta-learner fusion — XGBoost (preferred) or Logistic Regression fallback.

XGBoost model:  models/meta_learner/weights_xgb.json  (XGBoost native format)
LR fallback:    models/meta_learner/weights.json       (custom JSON)

XGBoost is loaded when weights_xgb.json exists; otherwise falls back to LR.
Both expose the same .predict(results) → (p_fake, details) interface.
"""

from __future__ import annotations

import json
import logging
import math
from pathlib import Path
from typing import Sequence

from backend.core.schema import DetectionResult

log = logging.getLogger(__name__)

_ROOT         = Path(__file__).parent.parent.parent / "models" / "meta_learner"
_XGB_MODEL    = _ROOT / "weights_xgb.json"
_XGB_META     = _ROOT / "weights_xgb_meta.json"
_LR_WEIGHTS   = _ROOT / "weights.json"


class MetaLearnerFusion:
    """
    Unified meta-learner: loads XGBoost if available, else Logistic Regression.
    Missing detector scores are imputed with 0.5 (neutral prior).
    """

    def __init__(self, weights_path: str | Path | None = None) -> None:
        self._loaded   = False
        self._mode     = "none"       # "xgboost" | "lr" | "none"
        self._features: list[str] = []
        self.loo_auc   = 0.0

        # XGBoost state
        self._booster  = None

        # LR state
        self._coef:      list[float] = []
        self._intercept: float = 0.0
        self._mean:      list[float] = []
        self._std:       list[float] = []

        if weights_path:
            # Explicit path → assume LR format
            self._load_lr(Path(weights_path))
        else:
            # Auto-select: prefer XGBoost
            if _XGB_MODEL.exists():
                self._load_xgb()
            if not self._loaded and _LR_WEIGHTS.exists():
                self._load_lr(_LR_WEIGHTS)

    # ── Loaders ──────────────────────────────────────────────────────────────

    def _load_xgb(self) -> None:
        try:
            import xgboost as xgb
        except ImportError:
            log.warning("meta_learner: xgboost not installed — falling back to LR")
            return

        try:
            booster = xgb.Booster()
            booster.load_model(str(_XGB_MODEL))
            self._booster = booster

            # Read feature list + AUC from metadata sidecar
            if _XGB_META.exists():
                meta = json.loads(_XGB_META.read_text())
                self._features = meta.get("features", [])
                self.loo_auc   = float(meta.get("cv_auc", 0.0))
            else:
                # Fallback: read feature names from booster
                self._features = booster.feature_names or []

            self._mode   = "xgboost"
            self._loaded = True
            log.info(
                "meta_learner: XGBoost loaded  features=%s  cv_auc=%.3f",
                self._features, self.loo_auc,
            )
        except Exception as exc:
            log.warning("meta_learner: XGBoost load failed (%s) — trying LR", exc)

    def _load_lr(self, path: Path) -> None:
        try:
            w = json.loads(path.read_text())
            self._features  = w["features"]
            self._coef      = [float(c) for c in w["coef"]]
            self._intercept = float(w["intercept"])
            self._mean      = [float(m) for m in w["scaler_mean"]]
            self._std       = [float(s) for s in w["scaler_std"]]
            self.loo_auc    = float(w.get("loo_auc", 0.0))
            self._mode      = "lr"
            self._loaded    = True
            log.info(
                "meta_learner: LR loaded from %s  features=%s  LOO-AUC=%.3f",
                path, self._features, self.loo_auc,
            )
        except Exception as exc:
            log.warning("meta_learner: LR load failed (%s)", exc)

    # ── Public API ────────────────────────────────────────────────────────────

    @property
    def is_ready(self) -> bool:
        return self._loaded

    def predict(
        self,
        results: Sequence[DetectionResult],
    ) -> tuple[float, dict]:
        score_map = {
            r.detector: float(r.p_fake)
            for r in results
            if r.error is None
        }

        raw_scores: list[float] = []
        used:     list[str]   = []
        imputed:  list[str]   = []

        for name in self._features:
            if name in score_map:
                raw_scores.append(score_map[name])
                used.append(name)
            else:
                raw_scores.append(0.5)
                imputed.append(name)

        if self._mode == "xgboost":
            return self._predict_xgb(raw_scores, used, imputed)
        else:
            return self._predict_lr(raw_scores, used, imputed)

    # ── XGBoost inference ─────────────────────────────────────────────────────

    def _predict_xgb(
        self,
        raw_scores: list[float],
        used: list[str],
        imputed: list[str],
    ) -> tuple[float, dict]:
        import numpy as np
        import xgboost as xgb

        X = np.array([raw_scores], dtype=np.float32)
        dmat = xgb.DMatrix(X, feature_names=self._features)
        p_fake = float(self._booster.predict(dmat)[0])

        details = {
            "method": "meta_learner_xgb",
            "cv_auc": self.loo_auc,
            "per_detector": {
                name: {"raw_score": round(raw_scores[i], 4), "imputed": name in imputed}
                for i, name in enumerate(self._features)
            },
            "used":    used,
            "imputed": imputed,
        }
        return p_fake, details

    # ── LR inference ──────────────────────────────────────────────────────────

    def _predict_lr(
        self,
        raw_scores: list[float],
        used: list[str],
        imputed: list[str],
    ) -> tuple[float, dict]:
        z = [
            (x - m) / max(s, 1e-9)
            for x, m, s in zip(raw_scores, self._mean, self._std)
        ]
        logit = self._intercept + sum(c * zi for c, zi in zip(self._coef, z))
        logit = max(-15.0, min(15.0, logit))
        p_fake = 1.0 / (1.0 + math.exp(-logit))

        details = {
            "method": "meta_learner_lr",
            "loo_auc": self.loo_auc,
            "logit": round(logit, 4),
            "per_detector": {
                name: {
                    "raw_score": round(raw_scores[i], 4),
                    "z": round(z[i], 4),
                    "coef": round(self._coef[i], 4),
                    "contribution": round(self._coef[i] * z[i], 4),
                    "imputed": name in imputed,
                }
                for i, name in enumerate(self._features)
            },
            "used":    used,
            "imputed": imputed,
        }
        return float(p_fake), details


# ── Module-level singleton ────────────────────────────────────────────────────

_instance: MetaLearnerFusion | None = None


def get_meta_learner() -> MetaLearnerFusion | None:
    global _instance
    if _instance is None:
        _instance = MetaLearnerFusion()
    return _instance if _instance.is_ready else None