backend/model_loader.py

"""
OmniDiag — Model Loader
=======================
Lazy-loads the model, preprocessors, and SHAP explainer for a given disease.
Supports both TreeExplainer (for XGBoost, RandomForest, etc.) and
DeepExplainer (for neural networks like TabNet, PyTorch).

The loader is instantiated per disease and caches loaded objects
so they are only loaded once (on first request).
"""

import os
import json
import logging
import traceback
import joblib
import pandas as pd
import numpy as np
import shap
import importlib
from typing import Optional, Dict, Any, List

from backend.shap_service import generate_shap_explanation

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

# Project root: resolve relative paths from the config
_PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
class ModelLoader:
    """
    Lazy loader for a single disease's model artifacts.
    
    Attributes:
        config: The disease configuration dictionary.
        _model: Cached model object (loaded on first access).
        _explainer: Cached SHAP explainer (loaded on first access).
        _preprocessors: Cached preprocessor objects (loaded on first access).
    """
    
    def __init__(self, config: dict):
        """
        Initialize the loader with a disease config.
        
        Args:
            config: Parsed YAML config dict for the disease.
                    Must contain 'model' key with 'weights_path' and 'explainer_type'.
        """
        self.config = config
        self._model: Optional[object] = None
        self._explainer: Optional[object] = None
        self._preprocessors: Optional[Dict[str, object]] = None
        self._feature_names: Optional[List[str]] = None
        self._project_root = _PROJECT_ROOT
        self._feature_engineer: Optional[object] = None
    
    # ------------------------------------------------------------------
    # Properties with lazy loading
    # ------------------------------------------------------------------
    
    @property
    def model(self) -> object:
        """Lazy-load and cache the trained model."""
        if self._model is None:
            weights_path = self._resolve_weights_path()
            log.debug(f"Loading model weights from: {weights_path}")
            log.debug(f"File exists: {os.path.exists(weights_path)}")
            if os.path.exists(weights_path):
                log.debug(f"File size: {os.path.getsize(weights_path)} bytes")
            self._model = joblib.load(weights_path)
            log.debug(f"Model loaded successfully. Type: {type(self._model).__name__}")

            # XGBoost 3.x compatibility patch: base_score may be stored as a
            # bracket-wrapped string (e.g. '[5.85041E-1]') in the model's raw
            # UBJSON serialization (save_raw()). Older SHAP versions parse this
            # with float() which fails on the bracket-wrapped format.
            #
            # Neither booster.load_config() nor booster.set_attr() modify the
            # save_raw() output, because base_score lives in the C-level
            # learner_model_param that is serialized from internal state.
            # Therefore we monkey-patch booster.save_raw() to strip the
            # brackets from the base_score string in the UBJSON output.
            import types as _types

            try:
                booster = self._model.get_booster()
                cfg = json.loads(booster.save_config())
                raw = cfg["learner"]["learner_model_param"]["base_score"]
                if isinstance(raw, str) and raw.startswith("[") and raw.endswith("]"):
                    raw_clean = raw.strip("[]")

                    _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: "
                        "%s -> %s",
                        raw,
                        raw_clean,
                    )
            except Exception:
                log.debug(
                    "base_score save_raw() patch skipped "
                    "(not an XGBoost model or already clean)"
                )

        return self._model
    
    @property
    def explainer(self) -> object:
        """Lazy-load and cache the SHAP explainer.

        The model's ``base_score`` is already patched at load time (see
        ``model`` property) so ``TreeExplainer`` should never encounter the
        XGBoost 3.x bracket-wrapped string format.
        """
        if self._explainer is None:
            explainer_type = self.config.get("model", {}).get("explainer_type", "tree")
            if explainer_type == "tree":
                self._explainer = shap.TreeExplainer(self.model)
            elif explainer_type == "deep":
                self._explainer = shap.DeepExplainer(self.model)
            else:
                raise ValueError(
                    f"Unknown explainer_type '{explainer_type}' for disease "
                    f"'{self.config.get('disease', {}).get('name', 'unknown')}'. "
                    f"Supported: 'tree', 'deep'."
                )
        return self._explainer
    
    @property
    def preprocessors(self) -> Dict[str, object]:
        """Lazy-load and cache preprocessors (label encoders, scaler)."""
        if self._preprocessors is None:
            preprocessors_path = self.config.get("model", {}).get("preprocessors_path", "")
            # Resolve relative to project root
            if preprocessors_path and not os.path.isabs(preprocessors_path):
                preprocessors_path = os.path.join(self._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)
        return self._preprocessors
    
    # ------------------------------------------------------------------
    # Feature engineering
    # ------------------------------------------------------------------
    
    def _get_feature_engineer(self):
        """
        Lazy-load the feature engineer class specified in the config.
        
        Returns:
            An instance of the feature engineer (subclass of BaseFeatureEngineer).
        """
        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}' from "
                        f"'{module_path}': {e}"
                    )
        return self._feature_engineer
    
    def _engineer_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Apply feature engineering to a DataFrame.
        
        Runs two feature engineering pipelines in dependency order:
            1. Heuristic (Statistical): Age_BP_Interaction, HR_Age_Ratio, Chol_Age_Ratio
            2. Medical (Cardiology): RPP, Exercise_Risk_Index
        
        Note: Age_Bins and Global_Risk_Score were tested (v5.1 beta) but did NOT
        improve accuracy — removed per A/B diagnostic (diagnose_v5_drop.py).
        
        Args:
            df: Raw patient DataFrame (12 base features).
        
        Returns:
            DataFrame with engineered features appended (up to 16 total features).
        """
        engineer = self._get_feature_engineer()
        if engineer:
            df = engineer.engineer_heuristic(df)     # Age_BP_Interaction, HR_Age_Ratio, Chol_Age_Ratio
            df = engineer.engineer_medical(df)       # RPP, Exercise_Risk_Index
        return df
    
    # ------------------------------------------------------------------
    # Preprocessing helpers
    # ------------------------------------------------------------------
    
    def _apply_preprocessors(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Apply label encoders to categorical columns and scaler to numeric columns.
        
        Args:
            df: Raw DataFrame with string categoricals and raw numeric values.
        
        Returns:
            DataFrame with encoded categoricals and scaled numerics.
        """
        df = df.copy()
        preprocessors = self.preprocessors
        
        # Apply label encoders to categorical columns
        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))
        
        # Apply standard scaler to numeric columns
        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
    # ------------------------------------------------------------------
    
    def predict(self, patient_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Run prediction on a single patient's data.

        Pipeline order must match the training pipeline exactly:
          1. Apply label encoders + StandardScaler (preprocessing)
          2. Engineer features FROM the already-scaled/encoded values

        This matches how final_ready_data.csv was built (encode → scale → save),
        and how the training script loaded it before calling engineer_*_features().
        Reversing the order causes ~5% accuracy loss because engineered features
        like Age_BP_Interaction are computed at completely different magnitudes.

        Args:
            patient_data: Dictionary of feature_name -> value.

        Returns:
            Dictionary with keys:
                - prediction: int (0 = Negative, 1 = Positive)
                - confidence: float (probability of positive class)
                - diagnosis: str ("Positive" or "Negative")
        """
        df = pd.DataFrame([patient_data])
        df = self._apply_preprocessors(df)   # encode → scale first
        df = self._engineer_features(df)     # then engineer from scaled values
        raw_pred = int(self.model.predict(df)[0])
        raw_proba = self.model.predict_proba(df)[0]
        # The heart disease training CSV has inverted labels:
        # HeartDisease=0 means "has disease", HeartDisease=1 means "healthy".
        # Verified via feature correlations (Oldpeak, Age, MaxHR are all sign-flipped
        # vs clinical expectations). We correct here so the API returns clinical truth.
        has_disease = (raw_pred == 0)
        confidence = float(raw_proba[0])  # P(class=0) = P(disease)
        return {
            "prediction": 1 if has_disease else 0,
            "confidence": confidence,
            "diagnosis": "Positive" if has_disease else "Negative"
        }
    
    
    def explain(self, patient_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Run SHAP explanation on a single patient's data.
        
        Applies feature engineering FIRST on raw data, then preprocessing
        (label encoding + scaling), then SHAP explanation. Returns structured
        chart data and a human-readable textual explanation — no images or
        matplotlib.
        
        Args:
            patient_data: Dictionary of feature_name -> value.
        
        Returns:
            Dictionary with keys:
                - chart_data: List of {"feature": str, "shap_value": float}
                  sorted by |shap_value| descending, for ShapBarChart.jsx.
                - text_explanation: Human-readable string identifying the
                  top 3 most impactful features with direction labels.
                - base_value: Base (expected) value from the explainer.
        """
        try:
            df = pd.DataFrame([patient_data])
            log.debug(f"Explain: raw data columns={list(df.columns)}")
            df = self._apply_preprocessors(df)   # encode → scale first (matches training pipeline)
            log.debug(f"Explain: after preprocessors columns={list(df.columns)}, shape={df.shape}")
            df = self._engineer_features(df)     # then engineer from scaled values
            log.debug(f"Explain: after engineering columns={list(df.columns)}, shape={df.shape}")
            
            # 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")

            log.debug("Creating SHAP explainer...")
            explainer = self.explainer
            log.debug(f"SHAP explainer ready: {type(explainer).__name__}")
            
            log.debug("Computing SHAP values...")
            shap_values = explainer(df)
            log.debug(f"SHAP values computed, shape={shap_values.values.shape}")
            
            feature_names = list(df.columns)
            
            result = generate_shap_explanation(shap_values, feature_names)

            # Attach prediction + confidence (same inversion fix as predict())
            raw_pred = int(self.model.predict(df)[0])
            raw_proba = self.model.predict_proba(df)[0]
            has_disease = (raw_pred == 0)
            result["prediction"] = 1 if has_disease else 0
            result["confidence"] = float(raw_proba[0])
            result["diagnosis"] = "Positive" if has_disease else "Negative"

            log.debug("SHAP explanation generated successfully")
            return result
        except Exception as e:
            log.error(f"Explain failed: {type(e).__name__}: {e}")
            log.error(traceback.format_exc())
            raise
    
    def get_feature_names(self) -> List[str]:
        """Return the feature names expected by the model."""
        if self._feature_names is None:
            # Try booster feature names first
            try:
                booster = self.model.get_booster()
                if booster.feature_names and all(n != '' for n in booster.feature_names):
                    self._feature_names = list(booster.feature_names)
            except Exception:
                pass
            # Fallback: try sklearn's feature_names_in_
            if self._feature_names is None:
                try:
                    self._feature_names = list(self.model.feature_names_in_)
                except (AttributeError, Exception):
                    self._feature_names = []
        return self._feature_names
    
    def generate_counterfactuals(self, patient_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Generate What-If counterfactual scenarios for heart disease.

        Uses random perturbation of mutable clinical features to find the
        minimal changes that flip the prediction from Positive to Negative.
        """
        import random

        baseline_result = self.predict(patient_data)
        baseline_pred = baseline_result["prediction"]
        baseline_prob = baseline_result["confidence"]

        if baseline_pred == 0:
            return {
                "status": "not_applicable",
                "counterfactuals": [],
                "baseline_probability": baseline_prob,
                "message": "Patient is already at low risk. No counterfactuals needed.",
            }

        # Heart disease mutable features and their perturbation ranges
        MUTABLE = {
            "RestingBP":   (90,  180,  False),   # (min, max, is_binary)
            "Cholesterol": (100, 400,  False),
            "MaxHR":       (60,  200,  False),
            "Oldpeak":     (0.0, 6.2,  False),
            "FastingBS":   (0,   1,    True),
            "ExerciseAngina": (None, None, True),  # Y/N toggle
        }

        rng = random.Random(42)
        candidates = []

        for _ in range(800):
            cf = dict(patient_data)
            changed: List[str] = []

            for feat, (lo, hi, is_binary) in MUTABLE.items():
                if feat not in cf:
                    continue
                if rng.random() < 0.4:
                    original = cf[feat]
                    if feat == "ExerciseAngina":
                        cf[feat] = "N" if str(original).upper() == "Y" else "Y"
                    elif is_binary:
                        cf[feat] = 1 - int(original)
                    else:
                        cf[feat] = round(rng.uniform(lo, hi), 1)
                    if cf[feat] != original:
                        changed.append(feat)

            if not changed:
                continue

            try:
                result = self.predict(cf)
                if result["prediction"] == 0:
                    candidates.append({
                        "features": cf,
                        "changed": changed,
                        "probability": result["confidence"],
                        "distance": len(changed),
                    })
            except Exception:
                continue

        # Sort by fewest changes, then by lowest probability
        candidates.sort(key=lambda c: (c["distance"], c["probability"]))

        # Pick top 3 with diversity (different primary change)
        selected = []
        seen_primary = set()
        for c in candidates:
            primary = c["changed"][0] if c["changed"] else ""
            if primary not in seen_primary:
                seen_primary.add(primary)
                selected.append(c)
            if len(selected) >= 3:
                break

        counterfactuals = []
        for c in selected:
            scenario_changes = []
            for feat in c["changed"]:
                original_val = patient_data.get(feat)
                new_val = c["features"].get(feat)
                scenario_changes.append({
                    "feature": feat,
                    "original_value": original_val,
                    "counterfactual_value": new_val,
                    "direction": "decrease" if (
                        isinstance(new_val, (int, float)) and isinstance(original_val, (int, float))
                        and new_val < original_val
                    ) else "increase",
                })
            counterfactuals.append({
                "scenario_id": len(counterfactuals) + 1,
                "probability": c["probability"],
                "changes": scenario_changes,
            })

        return {
            "status": "success",
            "counterfactuals": counterfactuals,
            "baseline_probability": baseline_prob,
        }

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _resolve_weights_path(self) -> str:
        """
        Resolve the model weights path, trying primary then fallback.
        
        Returns:
            The path to the model weights file.
        
        Raises:
            FileNotFoundError: If neither primary nor fallback path exists.
        """
        primary = self.config.get("model", {}).get("weights_path", "")
        fallback = self.config.get("model", {}).get("fallback_weights_path", "")
        
        # Resolve relative paths to absolute from project root
        primary_abs = primary if os.path.isabs(primary) else os.path.join(self._project_root, primary)
        fallback_abs = fallback if os.path.isabs(fallback) else os.path.join(self._project_root, fallback) if fallback else ""
        
        log.debug(f"Resolving weights path...")
        log.debug(f"  Project root: {self._project_root}")
        log.debug(f"  Primary (config): {primary}")
        log.debug(f"  Primary (abs): {primary_abs}")
        log.debug(f"  Fallback (config): {fallback}")
        log.debug(f"  Fallback (abs): {fallback_abs}")
        log.debug(f"  Primary exists: {os.path.exists(primary_abs)}")
        log.debug(f"  Fallback exists: {os.path.exists(fallback_abs) if fallback_abs else 'N/A'}")
        
        if os.path.exists(primary_abs):
            return primary_abs
        elif fallback_abs and os.path.exists(fallback_abs):
            return fallback_abs
        else:
            raise FileNotFoundError(
                f"Model weights not found. Tried:\n"
                f"  Primary: {primary_abs}\n"
                f"  Fallback: {fallback_abs}\n"
                f"Please ensure the model file exists at one of these paths."
            )