backend/ensemble_loader.py

"""
OmniDiag — Ensemble Model Loader (Stacking / Voting)
======================================================
Wraps multiple base models + a meta-learner for ensemble inference.

Key Design:
    - NEW class (does NOT modify existing ModelLoader)
    - Each base model is loaded as a standalone sklearn-compatible estimator
    - Feature engineering + preprocessing is applied ONCE, shared across models
    - Supports both stacking (meta-learner) and voting (equal-weight average)
    - SHAP explanations are weighted averages of per-model SHAP values

Usage:
    loader = EnsembleModelLoader(config)
    result = loader.predict(patient_data)
    explanation = loader.explain(patient_data)
"""

import os
import json
import logging
import traceback
import importlib
from typing import Dict, Any, List, Optional

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

from backend.shap_service import generate_shap_explanation
from backend.counterfactual_generator import CounterfactualGenerator

log = logging.getLogger("omnidiag.ensemble_loader")

# Project root for resolving relative paths
_PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


class EnsembleModelLoader:
    """
    Lazy-loading ensemble inference wrapper.

    Supports two modes:
      - 'stacking': base model probabilities → Logistic Regression meta-learner
      - 'voting':   average of base model probabilities (equal weight)

    Attributes:
        config: Full disease configuration dict with 'ensemble' section.
        _base_models: Dict of model_name → loaded model object.
        _meta_learner: Loaded Logistic Regression (stacking) or None (voting).
        _feature_engineer: Cached feature engineer instance.
        _preprocessors: Cached dict of preprocessor objects.
        _feature_names: Cached list of feature names expected by the models.
        _ensemble_type: 'stacking' or 'voting'.
    """

    def __init__(self, config: dict):
        """
        Initialize the ensemble loader with a disease config.

        Args:
            config: Parsed YAML config dict. Must contain a 'model.ensemble' key
                    with 'base_models' list and optionally 'meta_learner'.
        """
        self.config = config
        self.ensemble_config = config.get("model", {}).get("ensemble", {})
        self._ensemble_type = self.ensemble_config.get("type", "stacking")

        # Clinical inference threshold — read from config, default 0.5
        # Updated by training script after clinical threshold optimisation.
        # FN cost (2×) vs FP cost (1×) shifts threshold below 0.5 to catch
        # more true diabetics at the cost of increased false alarms — a
        # clinically justified trade-off (HbA1c follow-up resolves FPs).
        self._inference_threshold: float = float(
            config.get("model", {}).get("inference_threshold", 0.5)
        )

        # Cached objects
        self._base_models: Optional[Dict[str, object]] = None
        self._meta_learner: Optional[object] = None
        self._feature_engineer: Optional[object] = None
        self._preprocessors: Optional[Dict[str, object]] = None
        self._feature_names: Optional[List[str]] = None
        self._meta_feature_names: Optional[List[str]] = None

        log.info(
            f"EnsembleModelLoader initialized: type={self._ensemble_type}, "
            f"base_models={[m['name'] for m in self.ensemble_config.get('base_models', [])]}, "
            f"inference_threshold={self._inference_threshold}"
        )

    # ------------------------------------------------------------------
    # Properties with lazy loading
    # ------------------------------------------------------------------

    @property
    def base_models(self) -> Dict[str, object]:
        """Lazy-load all base models from their configured paths."""
        if self._base_models is None:
            self._base_models = {}
            base_configs = self.ensemble_config.get("base_models", [])

            for bcfg in base_configs:
                name = bcfg["name"]
                weights_path = self._resolve_path(bcfg["weights_path"])

                log.debug(f"Loading base model '{name}' from: {weights_path}")
                if not os.path.exists(weights_path):
                    raise FileNotFoundError(
                        f"Base model '{name}' weights not found at: {weights_path}"
                    )

                model = joblib.load(weights_path)
                self._base_models[name] = model
                log.info(f"Loaded base model '{name}': {type(model).__name__}")

        return self._base_models

    @property
    def meta_learner(self) -> Optional[object]:
        """Lazy-load the meta-learner (only for stacking)."""
        if self._ensemble_type != "stacking":
            return None

        if self._meta_learner is None:
            meta_config = self.ensemble_config.get("meta_learner", {})
            weights_path = self._resolve_path(meta_config.get("weights_path", ""))

            log.debug(f"Loading meta-learner from: {weights_path}")
            if not os.path.exists(weights_path):
                raise FileNotFoundError(
                    f"Meta-learner weights not found at: {weights_path}"
                )

            self._meta_learner = joblib.load(weights_path)
            log.info(
                f"Loaded meta-learner: {type(self._meta_learner).__name__}"
            )

        return self._meta_learner

    @property
    def preprocessors(self) -> Dict[str, object]:
        """Lazy-load preprocessors (standard_scaler, label_encoders, feature_names)."""
        if self._preprocessors is None:
            preprocessors_path = self.config.get("model", {}).get(
                "preprocessors_path", ""
            )
            if preprocessors_path and not os.path.isabs(preprocessors_path):
                preprocessors_path = os.path.join(
                    _PROJECT_ROOT, preprocessors_path
                )

            self._preprocessors = {}
            if preprocessors_path and os.path.isdir(preprocessors_path):
                for filename in os.listdir(preprocessors_path):
                    if filename.endswith(".pkl"):
                        filepath = os.path.join(preprocessors_path, filename)
                        key = filename.replace(".pkl", "")
                        self._preprocessors[key] = joblib.load(filepath)

                # Load feature names from JSON
                fn_path = os.path.join(preprocessors_path, "feature_names.json")
                if os.path.exists(fn_path):
                    with open(fn_path) as f:
                        self._feature_names = json.load(f)

                # Load meta-learner feature names
                mfn_path = os.path.join(
                    preprocessors_path, "meta_feature_names.json"
                )
                if os.path.exists(mfn_path):
                    with open(mfn_path) as f:
                        self._meta_feature_names = json.load(f)

        return self._preprocessors

    # ------------------------------------------------------------------
    # Feature Engineering
    # ------------------------------------------------------------------

    def _get_feature_engineer(self):
        """Lazy-load the feature engineer from config."""
        if self._feature_engineer is None:
            module_path = self.config.get("features", {}).get("module", "")
            class_name = self.config.get("features", {}).get("class", "")
            if module_path and class_name:
                try:
                    module = importlib.import_module(module_path)
                    engineer_class = getattr(module, class_name)
                    self._feature_engineer = engineer_class(self.config)
                except (ImportError, AttributeError) as e:
                    raise ImportError(
                        f"Could not load feature engineer '{class_name}' "
                        f"from '{module_path}': {e}"
                    )
        return self._feature_engineer

    def _engineer_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """Apply heuristic + medical feature engineering (NOT clinical)."""
        engineer = self._get_feature_engineer()
        if engineer:
            df = engineer.engineer_heuristic(df)
            df = engineer.engineer_medical(df)
        return df

    # ------------------------------------------------------------------
    # Preprocessing
    # ------------------------------------------------------------------

    def _apply_preprocessors(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Apply label encoders + standard scaler.
        Same logic as ModelLoader._apply_preprocessors().
        """
        df = df.copy()
        preprocessors = self.preprocessors

        # Label encoders
        label_encoders = preprocessors.get("label_encoders", {})
        if isinstance(label_encoders, dict):
            for col, encoder in label_encoders.items():
                if col in df.columns:
                    df[col] = encoder.transform(df[col].astype(str))

        # Standard scaler
        scaler = preprocessors.get("standard_scaler")
        if scaler is not None:
            numeric_cols = self.config.get("features", {}).get(
                "numerical_columns", []
            )
            numeric_cols_present = [c for c in numeric_cols if c in df.columns]
            if numeric_cols_present:
                df[numeric_cols_present] = scaler.transform(
                    df[numeric_cols_present]
                )

        return df

    # ------------------------------------------------------------------
    # Public API: predict
    # ------------------------------------------------------------------

    def predict(self, patient_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Run ensemble prediction on a single patient's data.

        For stacking: each base model predicts → meta-learner combines them.
        For voting:   each base model predicts → average probabilities.

        Returns:
            Dictionary with keys:
                - prediction: int (0 = Negative, 1 = Positive)
                - confidence: float (probability of positive class)
                - diagnosis: str ("Positive" or "Negative")
                - model_contributions: Dict of model_name → probability
                - ensemble_variance: float (std of base model probabilities)
                - model_agreement: str ("high", "moderate", "low")
        """
        # ── Feature engineering + preprocessing (shared) ──────────────
        df = pd.DataFrame([patient_data])
        df = self._engineer_features(df)
        df = self._apply_preprocessors(df)

        # ── Get probabilities from each base model ────────────────────
        probas = {}
        for name, model in self.base_models.items():
            try:
                proba = float(model.predict_proba(df)[0][1])
                probas[name] = proba
            except Exception as e:
                log.warning(
                    f"Base model '{name}' prediction failed: {e}. Using 0.5"
                )
                probas[name] = 0.5

        # ── Combine predictions ───────────────────────────────────────
        if self._ensemble_type == "stacking" and self.meta_learner is not None:
            # Stacking: meta-learner combines base model probabilities
            X_meta = np.array(
                [probas[name] for name in self.base_models.keys()]
            ).reshape(1, -1)
            final_proba = float(
                self.meta_learner.predict_proba(X_meta)[0][1]
            )
        else:
            # Voting: simple average
            final_proba = float(np.mean(list(probas.values())))

        final_pred = int(final_proba >= self._inference_threshold)

        # ── Ensemble variance (P1) ──────────────────────────────────────
        proba_values = list(probas.values())
        ensemble_variance = float(np.std(proba_values)) if len(proba_values) > 0 else 0.0
        # model_agreement: high (<0.05), moderate (0.05-0.15), low (>0.15)
        if ensemble_variance < 0.05:
            model_agreement = "high"
        elif ensemble_variance < 0.15:
            model_agreement = "moderate"
        else:
            model_agreement = "low"

        return {
            "prediction": final_pred,
            "confidence": final_proba,
            "diagnosis": "Positive" if final_pred == 1 else "Negative",
            "model_contributions": probas,
            "ensemble_type": self._ensemble_type,
            "inference_threshold": self._inference_threshold,
            "ensemble_variance": round(ensemble_variance, 4),
            "model_agreement": model_agreement,
        }

    # ------------------------------------------------------------------
    # Public API: generate_counterfactuals (P2: DiCE-inspired)
    # ------------------------------------------------------------------

    def generate_counterfactuals(self, patient_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Generate diverse counterfactual explanations for a patient.
        
        Uses the DiCE-inspired CounterfactualGenerator with random sampling
        and diversity selection to produce 3 actionable "what-if" scenarios
        that flip Positive → Negative predictions.
        
        Args:
            patient_data: Raw patient features dict (e.g., 21 BRFSS fields).
        
        Returns:
            Dictionary with:
                - counterfactuals: List of counterfactual scenario dicts
                - status: "generated" or "not_applicable" (if already low risk)
                - baseline_probability: Current probability of positive class
        """
        # Get baseline probability for metadata
        baseline_pred = self.predict(patient_data)
        baseline_proba = baseline_pred.get("confidence", 0.5)
        
        # If patient is already low risk (< threshold), no counterfactuals needed
        if baseline_proba < self._inference_threshold:
            return {
                "counterfactuals": [],
                "status": "not_applicable",
                "baseline_probability": baseline_proba,
                "message": "Patient is already below the clinical threshold. "
                          "Counterfactuals focus on reducing high-risk predictions.",
            }
        
        # Build the pipeline function that raw → engineered → preprocessed → predict
        def pipeline_fn(df: pd.DataFrame) -> pd.DataFrame:
            """Apply feature engineering + preprocessing to raw DataFrame."""
            df = self._engineer_features(df)
            df = self._apply_preprocessors(df)
            return df
        
        # Build the predict function for the generator (wraps predict)
        def predict_fn(df: pd.DataFrame) -> Dict[str, Any]:
            """Get the ensemble's predict_proba output for a DataFrame."""
            # df is already engineered+preprocessed, so we skip that in predict
            # We need a method that works on preprocessed data directly
            probas = {}
            for name, model in self.base_models.items():
                try:
                    proba = float(model.predict_proba(df)[0][1])
                    probas[name] = proba
                except Exception:
                    probas[name] = 0.5
            
            if self._ensemble_type == "stacking" and self.meta_learner is not None:
                X_meta = np.array(
                    [probas[name] for name in self.base_models.keys()]
                ).reshape(1, -1)
                final_proba = float(
                    self.meta_learner.predict_proba(X_meta)[0][1]
                )
            else:
                final_proba = float(np.mean(list(probas.values())))
            
            return {"confidence": final_proba}
        
        # Get feature names from preprocessed data
        sample_df = pipeline_fn(pd.DataFrame([patient_data]))
        feature_names = list(sample_df.columns)
        
        # Get raw feature names (all non-engineered features)
        raw_feature_names = [
            c for c in list(patient_data.keys())
            if c not in (
                "BMI_Age_Interaction", "Health_Index", "Lifestyle_Score",
                "SES_Composite", "Diabetes_Clinical_Risk"
            )
        ]
        
        # Create and run the counterfactual generator
        generator = CounterfactualGenerator(
            predict_fn=predict_fn,
            pipeline_fn=pipeline_fn,
            feature_names=feature_names,
            raw_feature_names=raw_feature_names,
            n_samples=500,
            n_counterfactuals=3,
            random_state=42,
            inference_threshold=self._inference_threshold,
        )
        
        counterfactuals = generator.generate(
            patient_data=patient_data,
            desired_class=0,  # Flip Positive → Negative
        )
        
        return {
            "counterfactuals": counterfactuals,
            "status": "generated" if counterfactuals else "no_valid_counterfactuals",
            "baseline_probability": baseline_proba,
            "message": None if counterfactuals else (
                "Could not find valid counterfactuals for this patient. "
                "The patient may need more significant lifestyle or medical changes "
                "than typical recommendations can provide."
            ),
        }
    
    # ------------------------------------------------------------------
    # Public API: explain
    # ------------------------------------------------------------------

    def explain(self, patient_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Generate a weighted-average SHAP explanation across all base models.

        For stacking: weights = meta-learner coefficients (absolute value).
        For voting:   weights = 1 / n_models (equal).

        If a base model's TreeExplainer fails (e.g. sklearn version incompatibility),
        that model's weight is redistributed to the working models.

        Returns:
            Dictionary with keys:
                - chart_data: List of {"feature": str, "shap_value": float}
                - text_explanation: Human-readable top-3 feature impacts
                - base_value: Weighted average base value
                - per_model_shap: Dict of model_name → individual SHAP result
                - shap_weights: Dict of model_name → weight used in average
                - ensemble_variance: float (std of base model probabilities)
                - model_agreement: str ("high", "moderate", "low")
        """
        try:
            df = pd.DataFrame([patient_data])
            df = self._engineer_features(df)
            df = self._apply_preprocessors(df)
            feature_names = list(df.columns)

            # ── Ensemble variance from base model predictions (P1) ─────
            probas = {}
            for name, model in self.base_models.items():
                try:
                    proba = float(model.predict_proba(df)[0][1])
                    probas[name] = proba
                except Exception:
                    probas[name] = 0.5
            proba_values = list(probas.values())
            ensemble_variance = float(np.std(proba_values)) if len(proba_values) > 0 else 0.0
            if ensemble_variance < 0.05:
                model_agreement = "high"
            elif ensemble_variance < 0.15:
                model_agreement = "moderate"
            else:
                model_agreement = "low"

            # ── Compute per-model SHAP values (P0: with robust fallback) ─
            # Cast any object dtype columns to category for SHAP TreeExplainer compatibility
            for col in df.columns:
                if df[col].dtype == 'object':
                    df[col] = df[col].astype('category')
                    log.debug(f"Cast column '{col}' from object to category for SHAP compatibility")

            # Cast to float64 for SHAP compatibility with sklearn 1.8.0
            df_float = df.astype(np.float64)
            per_model_shap = {}
            failed_models = set()

            # Build a mapping of model name → weights_path for fresh reload
            shap_model_paths = {}
            for bcfg in self.ensemble_config.get("base_models", []):
                shap_model_paths[bcfg["name"]] = self._resolve_path(
                    bcfg["weights_path"]
                )

            for name in self.base_models:
                try:
                    # Freshly reload model for SHAP (avoids cached-object issues
                    # with sklearn 1.8.0 + SHAP TreeExplainer)
                    fresh_model = joblib.load(shap_model_paths[name])

                    # XGBoost 3.x compatibility: base_score may be stored as a
                    # bracket-wrapped string (e.g. '[5.000088E-1]') in UBJSON.
                    # Safely parse using ast.literal_eval before float conversion.
                    import ast as _ast
                    try:
                        booster = fresh_model.get_booster()
                        cfg = json.loads(booster.save_config())
                        raw_bs = cfg["learner"]["learner_model_param"]["base_score"]
                        if isinstance(raw_bs, str):
                            stripped = raw_bs.strip()
                            if stripped.startswith("[") and stripped.endswith("]"):
                                parsed = _ast.literal_eval(stripped)
                                if isinstance(parsed, (list, tuple)) and len(parsed) == 1:
                                    clean_bs = float(parsed[0])
                                    import types as _types
                                    # Patch save_raw() to return clean UBJSON without brackets
                                    _original_save_raw = booster.save_raw
                                    def _patched_save_raw(self, raw_format="ubj"):
                                        raw = _original_save_raw(raw_format=raw_format)
                                        raw_bytes = bytes(raw)
                                        marker = b"base_score"
                                        idx = raw_bytes.find(marker)
                                        if idx >= 0:
                                            s_pos = idx + len(marker)
                                            if raw_bytes[s_pos] == 0x53 and raw_bytes[s_pos + 1] == 0x4C:
                                                len_bytes = raw_bytes[s_pos + 2 : s_pos + 2 + 8]
                                                old_len = int.from_bytes(len_bytes, "big")
                                                content_start = s_pos + 2 + 8
                                                old_content = raw_bytes[content_start : content_start + old_len]
                                                if old_content.startswith(b"[") and old_content.endswith(b"]"):
                                                    new_content = old_content[1:-1]
                                                    new_len = len(new_content)
                                                    patched = bytearray(raw_bytes)
                                                    patched[s_pos + 2 : s_pos + 2 + 8] = new_len.to_bytes(8, "big")
                                                    patched = patched[:content_start] + new_content + patched[content_start + old_len:]
                                                    return bytearray(patched)
                                        return raw
                                    booster.save_raw = _types.MethodType(_patched_save_raw, booster)
                                    log.info(
                                        "Patched XGBoost base_score via save_raw() monkey-patch "
                                        "for ensemble model '%s': '%s' -> %s",
                                        name, raw_bs, clean_bs,
                                    )
                    except Exception:
                        log.debug(
                            "base_score save_raw() patch skipped for ensemble model '%s' "
                            "(not XGBoost or already clean)",
                            name,
                        )

                    explainer = shap.TreeExplainer(fresh_model)
                    sv = explainer(df_float)
                    # For binary classifiers, sv.values shape is (1, n_features, 2)
                    # Take class 1 (positive) values from the last axis
                    vals = sv.values
                    if vals.ndim == 3:
                        vals = vals[:, :, -1]
                    # base_values shapes differ by model:
                    #   XGBoost/LightGBM: (1,)        — single class 1 value
                    #   Random Forest:    (1, 2)       — [class0, class1] both outputs
                    # np.ravel flattens any shape to 1D; [-1] takes the positive class
                    bv_raw = sv.base_values if hasattr(sv, "base_values") else 0.0
                    bv_scalar = float(np.ravel(bv_raw)[-1])
                    per_model_shap[name] = {
                        "values": vals.flatten().tolist(),
                        "base_value": bv_scalar,
                    }
                    nz = sum(1 for v in vals.flatten() if abs(v) > 1e-10)
                    log.debug(
                        f"SHAP for '{name}' succeeded: "
                        f"non-zero={nz}/{len(feature_names)}"
                    )
                except Exception as e:
                    log.warning(
                        f"SHAP explanation for '{name}' failed: "
                        f"{type(e).__name__}: {e}. Skipping."
                    )
                    log.warning(traceback.format_exc())
                    per_model_shap[name] = {
                        "values": [0.0] * len(feature_names),
                        "base_value": 0.0,
                    }
                    failed_models.add(name)

            # ── Determine weights ─────────────────────────────────────
            if (
                self._ensemble_type == "stacking"
                and self.meta_learner is not None
            ):
                # Use absolute meta-learner coefficients as weights
                raw_weights = np.abs(self.meta_learner.coef_[0])
                # Ensure same order as base_models
                model_names = list(self.base_models.keys())
                if len(raw_weights) == len(model_names):
                    shap_weights = {
                        name: float(w)
                        for name, w in zip(model_names, raw_weights)
                    }
                else:
                    # Fallback: equal weights
                    n = len(self.base_models)
                    shap_weights = {
                        name: 1.0 / n for name in self.base_models
                    }
            else:
                # Voting: equal weights
                n = len(self.base_models)
                shap_weights = {name: 1.0 / n for name in self.base_models}

            # ── Zero-out weights for failed models and re-normalize ───
            for failed in failed_models:
                shap_weights[failed] = 0.0

            # Normalize weights to sum to 1
            total_weight = sum(shap_weights.values())
            if total_weight > 0:
                shap_weights = {
                    k: v / total_weight for k, v in shap_weights.items()
                }
            else:
                # Catastrophic: all models failed — equal fallback
                n = len(self.base_models)
                shap_weights = {name: 1.0 / n for name in self.base_models}

            # ── Weighted average of SHAP values ───────────────────────
            n_features = len(feature_names)
            weighted_shap_values = np.zeros(n_features)
            weighted_base_value = 0.0

            for name, weight in shap_weights.items():
                sv = per_model_shap.get(name, {})
                values = np.array(sv.get("values", [0.0] * n_features))
                # Ensure length matches
                if len(values) != n_features:
                    values = np.array(
                        values[:n_features]
                        if len(values) > n_features
                        else list(values)
                        + [0.0] * (n_features - len(values))
                    )
                weighted_shap_values += weight * values
                weighted_base_value += weight * sv.get("base_value", 0.0)

            # ── Build structured explanation ──────────────────────────
            # Create a SHAP-like object for generate_shap_explanation
            class ShapValuesContainer:
                """Mimics shap.Explanation for use with generate_shap_explanation."""

                def __init__(self, values, base_values, feature_names):
                    self.values = np.array([values])
                    self.base_values = np.array([base_values])
                    self.data = None
                    self.feature_names = feature_names

            shap_container = ShapValuesContainer(
                weighted_shap_values, weighted_base_value, feature_names
            )

            result = generate_shap_explanation(shap_container, feature_names)

            # Attach prediction + confidence so /explain mirrors /predict output
            pred_result = self.predict(patient_data)
            result["prediction"] = pred_result["prediction"]
            result["confidence"] = pred_result["confidence"]
            result["diagnosis"] = pred_result["diagnosis"]

            # Add ensemble-specific metadata
            result["per_model_shap"] = per_model_shap
            result["shap_weights"] = shap_weights
            result["ensemble_variance"] = round(ensemble_variance, 4)
            result["model_agreement"] = model_agreement

            log.debug(
                f"Ensemble SHAP generated: {len(result.get('chart_data', []))} features, "
                f"weights={shap_weights}, "
                f"ensemble_variance={ensemble_variance:.4f}"
            )
            return result

        except Exception as e:
            log.error(
                f"Ensemble explain failed: {type(e).__name__}: {e}"
            )
            log.error(traceback.format_exc())
            raise

    # ------------------------------------------------------------------
    # Internal Helpers
    # ------------------------------------------------------------------

    def _resolve_path(self, path: str) -> str:
        """Resolve a potentially relative path against the project root."""
        if not path:
            return ""
        if os.path.isabs(path):
            return path
        return os.path.join(_PROJECT_ROOT, path)